StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor);

absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
                                    xla::PjRtCApiClient* c_api_client,
                                    Tensor* tensor);

absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
    const DeviceType& device_type);

absl::Status ResetPjRtClient(const DeviceType& device_type);

}  // namespace tensorflow

    srcs = [":device_sources"],
    hdrs = [":device_headers"],
    visibility = ["//visibility:private"],
    deps = [
        ":parallel_device_lib",
        "//tensorflow/c:c_api",
        "//tensorflow/c:tf_status_helper",
        "//tensorflow/c/eager:c_api",
        "//tensorflow/c/eager:c_api_experimental",
        "//tensorflow/c/eager:tfe_tensorhandle_internal",
        "@com_google_absl//absl/strings",

        `dtensor.shutdown_tpu_system`.

## Deprecations

*   The C++ `tensorflow::Code` and `tensorflow::Status` will become aliases of
    respectively `absl::StatusCode` and `absl::Status` in some future release.
    *   Use `tensorflow::OkStatus()` instead of `tensorflow::Status::OK()`.
    *   Stop constructing `Status` objects from `tensorflow::error::Code`.
    *   One MUST NOT access `tensorflow::errors::Code` fields. Accessing

  // When `reset_context` is true, initialize new cluster context state based
  // on cluster configurations provided in `server_def`; otherwise, update
  // existing context state with the provided `server_def`. Contexts created
  // on remote tasks will be considered stale and garbage collected after
  // `keep_alive_secs` of inactivity.
  virtual Status SetOrUpdateServerDef(const ServerDef& server_def,

                                                   TF_Status* status);

// Returns the status for the tensor handle. In TFRT, a tensor handle can carry
// error info if error happens. If so, the status will be set with the error
// info. If not, status will be set as OK.
TF_CAPI_EXPORT extern void TFE_TensorHandleGetStatus(TFE_TensorHandle* h,
                                                     TF_Status* status);

  if (h == nullptr) {
    status->status = tensorflow::errors::InvalidArgument("Invalid handle");
    return -1;
  }
  return tensorflow::unwrap(h)->DeviceId(&status->status);
}

TF_CAPI_EXPORT extern void TFE_TensorHandleGetStatus(TFE_TensorHandle* h,
                                                     TF_Status* status) {
  status->status = tensorflow::unwrap(h)->TensorHandleStatus();
}

  }
  return pjrt_c_api_client;
}

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();
}

  TF_DeleteStatus(status);

  return var_handle;
}

TFE_Op* AddOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) {
  TF_Status* status = TF_NewStatus();

  TFE_Op* op = TFE_NewOp(ctx, "AddV2", status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpAddInput(op, a, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpAddInput(op, b, status);

                                              TF_Status* status) {
  Session* session;
  status->status = NewSession(opt->options, &session);
  if (status->status.ok()) {
    return new TF_DeprecatedSession({session});
  } else {
    DCHECK_EQ(nullptr, session);
    return nullptr;
  }
}

void TF_CloseDeprecatedSession(TF_DeprecatedSession* s, TF_Status* status) {
  status->status = s->session->Close();
}

  Status CallBackwardFunction(
      const string& op_type, GradientFunction* gradient_function,
      const std::vector<int64_t>& unneeded_gradients,
      gtl::ArraySlice<AbstractTensorHandle*> output_gradients,
      absl::Span<AbstractTensorHandle*> result) const override;

  // Builds a tensor filled with ones with the same shape and dtype as `t`.
  Status BuildOnesLike(const TapeTensor& t,