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 = {};

TF_DEVICELIST_METHOD(const char*, TF_DeviceListType, device_type().c_str(),
                     nullptr);
TF_DEVICELIST_METHOD(int64_t, TF_DeviceListMemoryBytes, memory_limit(), -1);
TF_DEVICELIST_METHOD(uint64_t, TF_DeviceListIncarnation, incarnation(), 0);

#undef TF_DEVICELIST_METHOD

}  // end extern "C"

// --------------------------------------------------------------------------
// New Graph and Session API

      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:
    case TF_DataType::TF_UINT64:
      dtype.code = DLDataTypeCode::kDLUInt;
      break;
    case TF_DataType::TF_BFLOAT16:

    // 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`.
    int64_t r = int64_t{pread(posix_file->fd, dst, requested_read_length,

    EXPECT_EQ(calls.find(offset), calls.end()) << "at offset " << offset;
    calls.insert(offset);
    memset(buffer, 'x', n);
    TF_SetStatus(status, TF_OK, "");
    return n;
  };
  const uint32 block_count = 256;
  tf_gcs_filesystem::RamFileBlockCache cache(
      block_size, block_count * block_size, 0, fetcher);
  std::vector<char> out;
  out.resize(block_count, 0);

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

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

}

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

  TF_Tensor* a =
      TF_AllocateTensor(TF_UINT64, dims, 2, 6 * TF_DataTypeSize(TF_UINT64));
  TF_Tensor* b =
      TF_AllocateTensor(TF_UINT64, nullptr, 0, TF_DataTypeSize(TF_UINT64));
  EXPECT_NE(a, nullptr);
  EXPECT_NE(b, nullptr);

  EXPECT_EQ(6, TF_TensorElementCount(a));
  EXPECT_EQ(1, TF_TensorElementCount(b));
  EXPECT_EQ(6 * TF_DataTypeSize(TF_UINT64), TF_TensorByteSize(a));

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