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 =
      dynamic_cast<xla::PjRtCApiBuffer*>(av_tensor->GetBuffer().get());
  if (c_api_buffer == nullptr) {
    return absl::InternalError(

          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(

    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 = {};
  DLTensor* dltensor_in = &dlm_in.dl_tensor;
  dltensor_in->data = data.data();
  dltensor_in->device = {kDLCPU, 0};
  dltensor_in->ndim = static_cast<int32_t>(shape.size());
  dltensor_in->dtype = {kDLFloat, 32, 1};
  dltensor_in->shape = shape.data();