}
  return data;
}

void* allocate_tensor(const char* operation, size_t len) {
  return allocate_tensor(operation, len, cpu_allocator());
}

void deallocate_buffer(void* data, size_t len, void* arg) {
  Allocator* allocator = nullptr;
  if (arg == nullptr) {
    allocator = cpu_allocator();
  } else {
    allocator = reinterpret_cast<Allocator*>(arg);
  }
  if (LogMemory::IsEnabled() && data != nullptr) {

  DeviceSetup device_setup;
  device_setup.AddDevicesAndSetUp({DEVICE_GPU});
  Device* device = device_setup.GetDevice(DEVICE_GPU);
  std::vector<int64_t> dimensions = {2, 3};
  Tensor dest_cpu_tensor(cpu_allocator(), tensorflow::DT_INT32,
                         tensorflow::TensorShape(dimensions));
  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, GetPjRtClient(DEVICE_GPU));
  std::vector<int32_t> data{1, 2, 3, 4, 5, 6};

namespace tensorflow {
namespace {

class DummyDevice : public DeviceBase {
 public:
  explicit DummyDevice(Env* env) : DeviceBase(env) {}
  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
    return cpu_allocator();
  }
};

void TestBitcastOp(Tensor* input_tensor, DataType out_type,
                   TensorShape expected_shape, error::Code expected_code) {
  Status status;
  NodeDef def;
  def.set_op("Bitcast");

  host_alloc_attrs.set_on_host(true);
  Allocator* cpu_allocator = ctx->device()->GetAllocator(host_alloc_attrs);

  // Async compilation returns nullptr executable without an error.
  if (!executable && !pjrt_executable) {
    DCHECK(!must_compile_);
    Tensor compilation_key(cpu_allocator, DT_STRING, TensorShape({}));
    Tensor compilation_successful(cpu_allocator, DT_BOOL, TensorShape({}));

  void FillAllocationDescription(
      tensorflow::AllocationDescription* proto) const override {
    int64_t rb = size();
    proto->set_requested_bytes(rb);
    proto->set_allocator_name(tensorflow::cpu_allocator()->Name());
  }

  bool OwnsMemory() const override { return owns_memory_; }

 private:
  const size_t len_;
  void (*const deallocator_)(void* data, size_t len, void* arg);
  void* const deallocator_arg_;

namespace tensorflow {
namespace {

class DummyDevice : public DeviceBase {
 public:
  explicit DummyDevice(Env* env) : DeviceBase(env) {}
  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
    return cpu_allocator();
  }
};

// Helper for comparing output and expected output
void ExpectSummaryMatches(const Summary& actual, const string& expected_str) {
  Summary expected;

  mutex_lock lock(mu_);
  return GetAllocatorLocked(attr);
}

Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) {
  if (attr.on_host()) {
    return cpu_allocator();
  }

  if (xla_allocator_ == nullptr) {
    if (UsePjRtForSingleDeviceCompilation(device_name_)) {
      VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator";

  TF_DeleteStatus(s);
}

void Deallocator(void* data, size_t, void* arg) {
  tensorflow::cpu_allocator()->DeallocateRaw(data);
  *reinterpret_cast<bool*>(arg) = true;
}

TEST(CAPI, Tensor) {
  const int num_bytes = 6 * sizeof(float);
  float* values =
      reinterpret_cast<float*>(tensorflow::cpu_allocator()->AllocateRaw(
          TF_TensorDefaultAlignment(), num_bytes));
  int64_t dims[] = {2, 3};

  // DT_VARIANT tensors must be allocated on CPU since they wrap C++
  // objects which can not be efficiently represented in GPU memory.
  int numa_node = cc_ctx->device()->NumaNode();
  Tensor out(::tensorflow::cpu_allocator(numa_node), ::tensorflow::DT_VARIANT,
             ::tensorflow::TensorShape({}));
  Variant* out_v = &(out.scalar<Variant>()());
  Status cc_status = ZerosLikeVariant(cc_ctx, v, out_v, zeros_like_func);

}
#undef EXPECT_TF_SIZE

class DummyDevice : public DeviceBase {
 public:
  explicit DummyDevice(Env* env) : DeviceBase(env) {}
  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
    return cpu_allocator();
  }
};

TEST(TestKernel, TestInputAndOutputCount) {
  const char* node_name = "InputOutputCounterKernel";
  const char* op_name = "BarOp";
  const char* device_name = "FakeDeviceName2";