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) {
  TF_ASSIGN_OR_RETURN(tsl::StatusOr<xla::PjRtClient*> pjrt_client,

          nullptr, pjrt_client->addressable_devices()[0]));
  tensorflow::AsyncValueTensor* av_tensor =
      tensorflow::AsyncValueTensor::FromTensor(&tensor);
  av_tensor->SetBuffer(std::move(buffer));

  EXPECT_THAT(
      GetPjRtCBufferFromTensor(&tensor),
      StatusIs(
          error::INTERNAL,
          HasSubstr(absl::StrCat(