switch (ctx.device_type) {
    case DLDeviceType::kDLCPU:
      return "CPU:0";
    case DLDeviceType::kDLCUDA:
      return absl::StrCat("GPU:", ctx.device_id);
    case DLDeviceType::kDLROCM:
      return absl::StrCat("GPU:", ctx.device_id);
    default:
      return absl::nullopt;
  }
}

// Converts DLPack data type to TF_DATATYPE.

#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/public/session_options.h"

using tensorflow::GraphDef;
using tensorflow::NodeDef;

static void BoolDeallocator(void* data, size_t, void* arg) {
  delete[] static_cast<bool*>(data);
}

    const std::string& fname, TransactionToken* token,
    std::unique_ptr<RandomAccessFile>* result) {
  if (ops_->new_random_access_file == nullptr)
    return errors::Unimplemented(tensorflow::strings::StrCat(
        "Filesystem for ", fname, " does not support NewRandomAccessFile()"));

  UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
  auto file = std::make_unique<TF_RandomAccessFile>();

      root_dir_ = tensorflow::strings::StrCat(
          "/", tmp_dir_,
          tensorflow::strings::StrCat("tf_fs_", rng_val_, "_", test_name), "/");
    } else {
      root_dir_ = tensorflow::io::JoinPath(
          tmp_dir_,
          tensorflow::strings::StrCat("tf_fs_", rng_val_, "_", test_name));
    }
    if (!GetParam().empty()) {
      root_dir_ = tensorflow::strings::StrCat(GetParam(), "://", cloud_path_,

  for (int i = 0; i < num_values; ++i) {
    const auto num_dims_i = num_dims[i];

    if (num_dims_i > TensorShape::MaxDimensions()) {
      return errors::InvalidArgument(
          strings::StrCat("Value specified for `", attr_name, "` has ",
                          num_dims_i, " dimensions which is over the limit of ",
                          TensorShape::MaxDimensions(), "."));
    }
    if (num_dims_i < 0) {

  auto client_job = cluster->add_job();
  client_job->set_name("localhost");
  int client_port = tensorflow::testing::PickUnusedPortOrDie();
  client_job->mutable_tasks()->insert(
      {0, strings::StrCat("localhost:", client_port)});
  server_def.set_job_name("localhost");
  auto serialized = server_def.SerializeAsString();
  TFE_ContextSetServerDef(ctx, 0, serialized.data(), serialized.size(), status);

  std::strncpy(buffer, "VALID", BUFFER_SIZE);
  TFE_OpSetAttrString(op, "padding", buffer, std::strlen(buffer));
  // Overwriting value in "buffer", should be fine since TFE_Op
  // shouldn't be holding on to it.
  std::strncpy(buffer, "NHWC", BUFFER_SIZE);
  TFE_OpSetAttrString(op, "data_format", buffer, std::strlen(buffer));

  TFE_OpSetAttrType(op, "T", TF_FLOAT);

  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"))));
}

TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorIncoorectType) {
  auto allocator = std::make_unique<AsyncValueAllocator>();

#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/resource_loader.h"
#include "tensorflow/core/platform/str_util.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/util/equal_graph_def.h"

    // Special-cased operation for packing per-device tensors into one parallel
    // tensor.
    if (inputs.size() != parallel_device.num_underlying_devices()) {
      std::string message(absl::StrCat(
          "The parallel device ", parallel_device_name, " expected ",
          parallel_device.num_underlying_devices(),
          " inputs to TPUReplicatedInput, but got ", inputs.size()));