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

Status SetAttrFloat(AbstractOperation*, const char* attr_name, float value,
                    ForwardOperation*);
Status SetAttrBool(AbstractOperation*, const char* attr_name, bool value,
                   ForwardOperation*);
Status SetAttrType(AbstractOperation*, const char* attr_name, DataType value,
                   ForwardOperation*);
Status SetAttrShape(AbstractOperation*, const char* attr_name,

//   // forward `op`.
//   return new AddGradientFunction;
// }
//
// Status RegisterGradients(GradientRegistry* registry) {
//   return registry->Register("Add", AddRegisterer);
// }
class GradientFunction {
 public:
  virtual Status Compute(AbstractContext* ctx,
                         absl::Span<AbstractTensorHandle* const> grad_outputs,

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

 public:
  // Add a function parameter and return the corresponding tensor.
  virtual Status AddParameter(DataType dtype, const PartialTensorShape& shape,
                              TracingTensorHandle**) = 0;

  // Finalize this context and make a function out of it. The context is in a
  // invalid state after this call and must be destroyed.
  virtual Status Finalize(OutputList* outputs, AbstractFunction**) = 0;

  // For LLVM style RTTI.

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

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

TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_TensorHandleCopySharingTensor(
    TFE_TensorHandle* h, TF_Status* status) {
  if (h == nullptr) {

                              status);
    CHECK_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);

    matmul = TFE_NewOp(ctx, "MatMulFunction", status);
    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
    TFE_OpAddInput(matmul, h0_task0, status);
    ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);

  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
  EXPECT_EQ(assign, 1);
  int id = TF_OpIsStateful("Identity", status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
  EXPECT_EQ(id, 0);
}

class ShapeInferenceTest : public ::testing::Test {
 protected:
  ShapeInferenceTest()

                                                TF_Status* status);

TF_CAPI_EXPORT extern void TF_AllocateTempForVariableInfo(
    TF_OpKernelContext* ctx, TF_VariableInfo* var_info, TF_Status* status);

TF_CAPI_EXPORT extern TF_Tensor* TF_GetTensorFromVariableInfo(
    TF_VariableInfo* var_info, TF_Status* status);

TF_CAPI_EXPORT extern void TF_DeleteVariableInfo(TF_VariableInfo* var_info);