ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  TFE_OpSetAttrType(op.get(), "dtype", TF_FLOAT);
  TFE_OpSetAttrShape(op.get(), "shape", {}, 0, status.get());
  TFE_OpSetAttrString(op.get(), "container", "", 0);
  TFE_OpSetAttrString(op.get(), "shared_name", "", 0);
  TFE_OpSetDevice(op.get(), name, status.get());

      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,
                       AbstractTensorHandle** result) const override;

  // Looks up the ID of a Gradient.

namespace tensorflow {

std::string AbstractTensorHandle::DebugString() const {
  PartialTensorShape shape;
  Status s = Shape(&shape);
  std::string shape_string;
  if (!s.ok()) {
    shape_string = "<error computing shape>";
  } else {
    shape_string = shape.DebugString();
  }
  return absl::StrCat("TensorHandle(shape=", shape_string,
                      ", dtype=", DataType_Name(DataType()),

  auto c_api_client = down_cast<xla::PjRtCApiClient*>(pjrt_client->get());
  std::vector<int32_t> data(1, 0);
  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});

  auto buffer = c_api_client->pjrt_c_client()->client->BufferFromHostBuffer(
      data.data(), shape.element_type(), shape.dimensions(),
      /*byte_strides=*/std::nullopt,
      xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr,

      if (combined_shape.dims() < 0 ||
          combined_shape.dims() != component_shape.dims()) {
        PartialTensorShape first_shape;
        TF_RETURN_IF_ERROR(unwrap(tensors_[0].get())->Shape(&first_shape));
        return errors::Unimplemented(absl::StrCat(
            "Computing the shape of a ParallelTensor when the components do "
            "not all have the same rank is not supported. One tensor had "
            "shape ",

      TF_ASSERT_OK(s);

      EXPECT_EQ(DT_FLOAT, in_props.dtype());
      EXPECT_FALSE(in_props.shape().unknown_rank());
      EXPECT_EQ(2, in_props.shape().dim_size());
      EXPECT_EQ(10, in_props.shape().dim(0).size());
      EXPECT_EQ(1, in_props.shape().dim(1).size());

      for (int i = 0; i < in_props_buf.size(); i++)
        TF_DeleteBuffer(in_props_buf[i]);
    }
  }

// number of dimensions of a parallel tensor.
int ParallelTensorNumDims(void* data, TF_Status* status) {
  const std::vector<int64_t>* shape;
  Status s = reinterpret_cast<ParallelTensor*>(data)->Shape(&shape);
  if (!s.ok()) {
    tsl::Set_TF_Status_from_Status(status, s);
    return -1;
  }
  return shape->size();
}

// Used as an argument to TFE_NewCustomDeviceTensorHandle, for computing a
// dimension of a parallel tensor.

  AbstractTensorHandlePtr sum_dims;
  {
    vector<int32_t> vals(num_dims_out);
    int64_t vals_shape[] = {num_dims_out};
    Range(&vals, 0, num_dims_out);
    AbstractTensorHandle* sum_dims_raw = nullptr;
    TF_RETURN_IF_ERROR(TestTensorHandleWithDims<int32_t, TF_INT32>(
        ctx, vals.data(), vals_shape, 1, &sum_dims_raw));
    sum_dims.reset(sum_dims_raw);
  }

  // Reduce sum the output on all dimensions.

      TFE_NewOp(context, "VarHandleOp", status), TFE_DeleteOp);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_OpSetAttrType(op.get(), "dtype", type);
  TFE_OpSetAttrShape(op.get(), "shape", dims, num_dims, status);
  TFE_OpSetAttrString(op.get(), "container", "", 0);
  // Use the special GUID for no buffer sharing
  //
  // TODO(allenl): Should we provide a better API for this? AFAIK this is the

    status->status = cc_status;
    return;
  }

  cc_status = cc_ctx->allocate_temp(var_info->var_info.var()->tensor()->dtype(),
                                    var_info->var_info.var()->tensor()->shape(),
                                    var_info->var_info.var()->tensor());
  status->status = cc_status;
}

TF_Tensor* TF_GetTensorFromVariableInfo(TF_VariableInfo* var_info,