if (!status.ok()) {
    done(status);
    return;
  }

  xla::PjRtBuffer* device_buffer;
  AsyncValueTensor* device_tensor_av =
      tensorflow::AsyncValueTensor::FromTensor(device_tensor);
  if (use_pjrt_tensor_buffer_) {
    if (device_tensor_av) {
      done(absl::InvalidArgumentError(
          "If use_pjrt_tensor_buffer is set, the device tensor should not "

                         /*len=*/sizeof(value),
                         /*deleter=*/[](void*, size_t) {}, &status);
  ASSERT_TRUE(status.ok()) << status.message();

  TensorHandle handle =
      TensorHandle::FromTensor(original_tensor, *runtime, &status);
  ASSERT_TRUE(status.ok()) << status.message();

  Tensor tensor = handle.Resolve(&status);
  ASSERT_TRUE(status.ok()) << status.message();

  EXPECT_EQ(tensor.dims(), 0);

#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) {
    return absl::InternalError("Input tensor does not have PjRtBuffer.");
  }
  auto* c_api_buffer =

          xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
          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(

  Tensor Resolve(Status* status);

  // Constructs a TensorHandle from a Tensor. If an error occurred,
  // status->ok() will be false, and the returned TensorHandle must not be used.
  static TensorHandle FromTensor(const Tensor& tensor, const Runtime& runtime,
                                 Status* status);

  // TensorHandle is movable, and not copyable
  TensorHandle(TensorHandle&&) = default;

/*static*/ XlaTensor* XlaTensor::FromTensor(const Tensor* tensor) {
  if (tensor->NumElements() == 0) {
    return nullptr;
  }
  XlaTensor* xla_tensor =
      FromOpaquePointer(const_cast<char*>(tensor->tensor_data().data()));
  return xla_tensor;
}

/*static*/ se::DeviceMemoryBase XlaTensor::DeviceMemoryFromTensor(
    const Tensor& tensor) {
  const XlaTensor* xla_tensor = FromTensor(&tensor);
  if (xla_tensor) {

    const int dst_device_ordinal =
        dst_xla_context->stream()->parent()->device_ordinal();

    XlaTensor* const xla_input = XlaTensor::FromTensor(input);
    TF_RET_CHECK(xla_input != nullptr && xla_input->has_shaped_buffer());
    XlaTensor* const xla_output = XlaTensor::FromTensor(output);
    TF_RET_CHECK(xla_output != nullptr && !xla_output->has_shaped_buffer());
    TF_RET_CHECK(input->shape() == output->shape());

// pointer data stored in the TensorBuffer is a tagged pointer.
class XlaTensor {
 public:
  // Downcast from a Tensor to an XlaTensor. Return nullptr if the downcast
  // fails.
  static XlaTensor* FromTensor(const Tensor* tensor);

  // Create a DeviceMemoryBase from a Tensor. The Tensor can be an XlaTensor, in
  // which case the returned value is shaped_buffer()->root_buffer(), or a

    if (use_multiple_streams_) {
      CHECK(ctx->op_device_context() && ctx->op_device_context()->stream())
          << "Must have a stream available when using XLA tensors!";
      XlaTensor* xla_tensor = XlaTensor::FromTensor(t);
      CHECK(xla_tensor);
      xla_tensor->WaitForDefinitionEventOnStream(
          ctx->op_device_context()->stream());
    }

    arguments.emplace_back(device_shape, host_shape);

          << " " << cpu_tensor->NumElements() << " "
          << cpu_tensor->shape().DebugString() << " "
          << device_tensor->shape().DebugString();

  XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor);
  CHECK(xla_tensor);

  XlaLayoutPreference layout_preference =
      shape_determination_fns_.layout_preference_fn(
          device_tensor->shape(), device_tensor->dtype(), std::nullopt);