*handle = buf;
  }
}

void TF_DeletePRunHandle(const char* handle) {
  delete[] handle;
  // TODO(suharshs): Free up any resources held by the partial run state.
}

void TF_SessionPRun(TF_Session* session, const char* handle,
                    const TF_Output* inputs, TF_Tensor* const* input_values,
                    int ninputs, const TF_Output* outputs,

    {
      tensorflow::mutex_lock l(execution_mutex_);
      while (execution_state_ == ExecutionState::kIdle ||
             execution_state_ == ExecutionState::kHasResult) {
        start_execute_.wait(l);
      }
      if (execution_state_ == ExecutionState::kShuttingDown) {
        return;
      } else if (execution_state_ == ExecutionState::kReadyToExecute) {
        // op_outputs_ may have been std::moved

          downloaded_block = true;
          block->state = FetchState::FINISHED;
        } else {
          block->state = FetchState::ERROR;
        }
        block->cond_var.SignalAll();
        return;
      case FetchState::FETCHING:
        block->cond_var.WaitWithTimeout(&block->mu, absl::Minutes(1));
        if (block->state == FetchState::FINISHED) {
          return TF_SetStatus(status, TF_OK, "");

  TFE_DeleteContext(ctx);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TF_DeleteStatus(status);
}
BENCHMARK(BM_InitOp);

void BM_Execute(::testing::benchmark::State& state) {
  const int async = state.range(0);
  state.SetLabel(async ? "ExecuteAsync" : "Execute");
  TF_Status* status = TF_NewStatus();
  TFE_ContextOptions* opts = TFE_NewContextOptions();

  int32_t init_timeout_in_ms = 300000;
  TFE_Context* ctx_0 =
      CreateContext(serialized_server_def_0,
                    /*isolate_session_state=*/false, init_timeout_in_ms);
  TFE_Context* ctx_1 =
      CreateContext(serialized_server_def_1,
                    /*isolate_session_state=*/false, init_timeout_in_ms);

  // Remote device on `worker2`.
  const char remote_device[] = "/job:localhost/replica:0/task:2/device:CPU:0";

std::string ModularFileSystemTest::cloud_path_;
std::string ModularFileSystemTest::tmp_dir_;

// As some of the implementations might be missing, the tests should still pass
// if the returned `Status` signals the unimplemented state.
bool UnimplementedOrReturnsCode(Status actual_status, Code expected_code) {
  Code actual_code = actual_status.code();
  return (actual_code == Code::UNIMPLEMENTED) || (actual_code == expected_code);
}

      output_gradients, result, /*build_default_zeros_grads*/ false));
  return absl::OkStatus();
}

// Helper functions which delegate to `AbstractOperation`, update
// the state of the ForwardOperation and call the tape as appropriate.
// These APIs are mainly to facilitate testing and are subject to change.
namespace internal {
Status Reset(AbstractOperation* op_, const char* op,

absl::Status ResetPjRtClient(const DeviceType& device_type) {
  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
  PjRtState* pjrt_state;
  TF_RETURN_IF_ERROR(rmgr->Lookup(rmgr->default_container(),
                                  kPjRtStateResourceName, &pjrt_state));
  TF_RETURN_IF_ERROR(pjrt_state->MovePjRtClientToUnused(device_type));
  return absl::OkStatus();
}

TFE_Context* CreateContext(const std::string& serialized_server_def,
                           bool isolate_session_state,
                           int64_t init_timeout_in_ms) {
  TF_Status* status = TF_NewStatus();
  TFE_ContextOptions* opts = TFE_NewContextOptions();
  opts->session_options.options.config.set_isolate_session_state(
      isolate_session_state);
  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(false));

      TF_SetPayload(&s, tensorflow::CoordinationErrorPayloadKey().data(),
                    error.SerializeAsString().c_str());
    }
    *state_iter = std::move(s);
    ++state_iter;
  }
  status->status = absl::OkStatus();
}

void TFE_WaitAtBarrier(TFE_Context* ctx, const char* barrier_id,
                       int64_t barrier_timeout_in_ms, TF_Status* status) {