std::vector<int64_t> shape,
                          std::vector<int64_t> strides) {
  size_t num_elements = 1;
  for (int i = 0; i < static_cast<int32_t>(shape.size()); ++i) {
    num_elements *= shape[i];
  }
  std::vector<float> data(num_elements);
  for (size_t j = 0; j < num_elements; ++j) {
    data[j] = j;
  }
  DLManagedTensor dlm_in = {};

namespace gradients {

using namespace std;

// ================== Helper functions =================

// Fills data with values [start,end) with given step size.
void Range(vector<int32_t>* data, int32_t start, int32_t end,
           int32_t step = 1) {
  for (int32_t i = start; i < end; i += step) {
    (*data)[i] = i;
  }
}

// Fills out_dims with the dimensions of the given tensor.
void GetDims(const TF_Tensor* t, int64_t* out_dims) {

  while (n > 0) {
    // Some platforms, notably macs, throw `EINVAL` if `pread` is asked to read
    // more than fits in a 32-bit integer.
    size_t requested_read_length;
    if (n > INT32_MAX)
      requested_read_length = INT32_MAX;
    else
      requested_read_length = n;

    // `pread` returns a `ssize_t` on POSIX, but due to interface being
    // cross-platform, return type of `Read` is `int64_t`.

  }
  auto pjrt_client = xla::GetCApiClient(DEVICE_CPU);
  CHECK_OK(pjrt_client.status());
  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(),

    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    X.reset(X_raw);
  }
  // Label
  int32_t Y_vals[] = {1, 0, 1};
  int64_t Y_dims[] = {3};
  AbstractTensorHandlePtr Y;
  {
    AbstractTensorHandle* Y_raw;
    status_ = TestTensorHandleWithDims<int32_t, TF_INT32>(
        immediate_execution_ctx_.get(), Y_vals, Y_dims, 1, &Y_raw);
    ASSERT_EQ(errors::OK, status_.code()) << status_.message();

    TFE_ContextSetExecutorForThread(ctx, executor);
    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

    TF_Tensor* t =
        TF_AllocateTensor(TF_INT32, nullptr, 0, 1 * sizeof(tensorflow::int32));
    *reinterpret_cast<tensorflow::int32*>(TF_TensorData(t)) = 42;
    TFE_TensorHandle* h = TFE_NewTensorHandle(t, status);
    ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
    TF_DeleteTensor(t);

  TF_Tensor* out = csession.output_tensor(0);
  ASSERT_TRUE(out != nullptr);
  EXPECT_EQ(TF_INT32, TF_TensorType(out));
  EXPECT_EQ(0, TF_NumDims(out));  // scalar
  ASSERT_EQ(sizeof(int32), TF_TensorByteSize(out));
  int32* output_contents = static_cast<int32*>(TF_TensorData(out));
  EXPECT_EQ(3 + 2, *output_contents);

  // Add another operation to the graph.

                       const int32_t* values) {
  int64_t num_values = 1;
  for (int i = 0; i < num_dims; ++i) {
    num_values *= dims[i];
  }
  TF_Tensor* t =
      TF_AllocateTensor(TF_INT32, dims, num_dims, sizeof(int32_t) * num_values);
  memcpy(TF_TensorData(t), values, sizeof(int32_t) * num_values);
  return t;
}

TF_Tensor* Int32Tensor(const std::vector<int32_t>& values) {
  int64_t dims = values.size();

}

std::unique_ptr<ParallelTensor> ParallelDevice::DeviceIDs(
    TFE_Context* context, TF_Status* status) const {
  std::vector<int32_t> ids;
  ids.reserve(num_underlying_devices());
  for (int i = 0; i < num_underlying_devices(); ++i) {
    ids.push_back(i);
  }
  return ScalarsFromSequence<int32_t>(ids, context, status);
}

absl::optional<std::vector<std::unique_ptr<ParallelTensor>>>

    case TF_DataType::TF_FLOAT:
    case TF_DataType::TF_DOUBLE:
      dtype.code = DLDataTypeCode::kDLFloat;
      break;
    case TF_DataType::TF_INT8:
    case TF_DataType::TF_INT16:
    case TF_DataType::TF_INT32:
    case TF_DataType::TF_INT64:
      dtype.code = DLDataTypeCode::kDLInt;
      break;
    case TF_DataType::TF_UINT8:
    case TF_DataType::TF_UINT16:
    case TF_DataType::TF_UINT32: