tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
  TF_ASSERT_OK_AND_ASSIGN(
      auto pjrt_client,
      xla::GetTfrtCpuClient(/*asynchronous=*/true, /*cpu_device_count=*/1));
  std::vector<int32_t> data(1, 0);
  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
  TF_ASSERT_OK_AND_ASSIGN(
      auto buffer,
      pjrt_client->BufferFromHostBuffer(

  if (av_tensor == nullptr) {
    TF_ASSIGN_OR_RETURN(
        *tensor, MakeTensorFromPjRtBuffer(tensor->dtype(), tensor->shape(),
                                          std::move(buffer)));
  } else {
    av_tensor->SetBuffer(std::move(buffer));
  }
  return absl::OkStatus();
}

absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
    const DeviceType& device_type) {