TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorNoBuffer) {
  auto allocator = std::make_unique<AsyncValueAllocator>();
  tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});

  EXPECT_THAT(
      GetPjRtCBufferFromTensor(&tensor),
      StatusIs(error::INTERNAL, HasSubstr(absl::StrCat(
                                    "Input tensor does not have PjRtBuffer"))));
}

#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/pjrt_c_api_client.h"
#include "tensorflow/core/framework/tensor.h"

namespace tensorflow {

StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor);

absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
                                    xla::PjRtCApiClient* c_api_client,
                                    Tensor* tensor);

#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {

absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor) {
  tensorflow::AsyncValueTensor* av_tensor =
      tensorflow::AsyncValueTensor::FromTensor(tensor);
  if (av_tensor == nullptr || av_tensor->GetBuffer() == nullptr) {

  auto s = tensorflow::TF_TensorToTensor(c_tensor, &tensor);
  if (!s.ok()) {
    status->status = s;
    return nullptr;
  }
  absl::StatusOr<PJRT_Buffer*> c_buffer =
      tensorflow::GetPjRtCBufferFromTensor(&tensor);
  if (!c_buffer.ok()) {
    status->status = c_buffer.status();
    return nullptr;
  }
  status->status = absl::OkStatus();
  return *c_buffer;
}