}
  return RunPjRtExecutable(/*num_missing_prefix_ctx_inputs=*/0, inputs,
                           variable_snapshots, variables, compilation_result,
                           pjrt_client, executable, ctx);
}

// TODO(b/289421064): Add unit test for this.
Status RunPjRtExecutable(
    int num_missing_prefix_ctx_inputs, const std::vector<const Tensor*>& inputs,

            done);
        OP_REQUIRES_OK_ASYNC(ctx, LockVariables(absl::MakeSpan(variable_infos)),
                             done);
        OP_REQUIRES_OK_ASYNC(
            ctx,
            RunPjRtExecutable(inputs, variable_infos, *compilation_result,
                              pjrt_client, pjrt_executable, ctx),
            done);
      }
      VLOG(2) << "Done executing with PJRT.";
      done();
    };