owner->reference.Unref();
  delete owner;
}

// Converts TF_DATAType to DLPack data type.
DLDataType GetDlDataType(TF_DataType data_type, TF_Status* status) {
  DLDataType dtype;
  dtype.lanes = 1;
  dtype.bits = TF_DataTypeSize(data_type) * 8;
  switch (data_type) {
    case TF_DataType::TF_BOOL:
      dtype.code = DLDataTypeCode::kDLBool;
      break;
    case TF_DataType::TF_HALF:
    case TF_DataType::TF_FLOAT:

    input_ids[i] = ToId(inputs[i]);
    input_dtypes[i] = inputs[i]->DataType();
  }
  std::vector<TapeTensor> tape_tensors;
  tape_tensors.reserve(outputs.size());
  for (auto t : outputs) {
    tape_tensors.push_back(TapeTensor(t));
  }
  GradientTape::RecordOperation(
      op_name, tape_tensors, input_ids, input_dtypes,
      [gradient_function]() -> GradientFunction* { return gradient_function; },

  TF_RETURN_IF_ERROR(ValidateTPOptimizerConfigs(optimizer_configs));

  CGraphOptimizerRegister(
      [=]() { return new CGraphOptimizer(optimizer, params.device_type); },
      optimizer_configs, params.device_type);

  return absl::OkStatus();
}

}  // namespace grappler
}  // namespace tensorflow

void TF_GetNodesToPreserveListSize(const TF_GrapplerItem* item, int* num_values,

  }
  return found_dtype && found_shape;
}

bool IsScalarConst(const tensorflow::NodeDef& node_def, int v) {
  if (node_def.op() != "Const" || node_def.name() != "scalar") {
    return false;
  }
  bool found_dtype = false;
  bool found_value = false;
  for (const auto& attr : node_def.attr()) {
    if (attr.first == "dtype") {
      if (attr.second.type() == tensorflow::DT_INT32) {
        found_dtype = true;

  executed = false;
  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());
  EXPECT_FALSE(executed);
  ASSERT_FALSE(TF_GetCode(status.get()) == TF_OK);

  // Custom devices can refuse to do type-based dispatch (as hcustom1 is
  // configured to do)
  matmul.reset(MatMulOp(context.get(), hcustom1.get(), hcpu.get()));
  num_retvals = 1;
  executed = false;

}

TFE_TensorHandle* TFE_NewTensorHandleFromScalar(TF_DataType data_type,
                                                void* data, size_t len,
                                                TF_Status* status) {
  auto dtype = static_cast<tensorflow::DataType>(data_type);
  DCHECK(tensorflow::DataTypeCanUseMemcpy(dtype));

  tensorflow::Tensor tensor(dtype, tensorflow::TensorShape({}));

      tensorflow::DeviceType(device_type), std::move(*pjrt_client));
  status->status = s;
}

void TF_ResetPjRtCClient(const char* device_type, TF_Status* status) {
  status->status =
      tensorflow::ResetPjRtClient(tensorflow::DeviceType(device_type));
}

PJRT_Client* TF_GetPjRtCClient(const char* device_type, TF_Status* status) {
  absl::StatusOr<xla::PjRtCApiClient*> pjrt_c_api_client =

      "      attr {"
      "        key: 'dtype'"
      "        value {"
      "          type: DT_FLOAT"
      "        }"
      "      }"
      "    }"
      "    node_def {"
      "      name: 'read1'"
      "      op: 'ReadVariableOp'"
      "      input: 'var'"
      "      device: '/job:localhost/replica:0/task:1/device:CPU:0'"
      "      attr {"
      "        key: 'dtype'"
      "        value {"

    TF_SetStatus(status, TF_OK, "");
    PopulateDefaultParam(params);
    params->device_type = "Device0";
  };
  auto plugin_init_1 = [](TP_OptimizerRegistrationParams* const params,
                          TF_Status* const status) -> void {
    TF_SetStatus(status, TF_OK, "");
    PopulateDefaultParam(params);
    params->device_type = "Device1";
  };

  TF_ASSERT_OK(InitGraphPlugin(plugin_init_0));

Variable* Variable::Create(TFE_Context* context, TF_DataType type,
                           const int64_t* dims, const int num_dims,
                           const char* device, TF_Status* status) {
  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> op(
      TFE_NewOp(context, "VarHandleOp", status), TFE_DeleteOp);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_OpSetAttrType(op.get(), "dtype", type);