//  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
//  xla::Shape shape(xla::F32, {2, 2}, {}, {})
//  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =
//      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(stream.parent());
//  done_event->Init();
//
//  XlaHostSendDeviceContext device_context(&stream, &gpu_dst,
//    shape, done_event);
//  device_context.CopyCPUTensorToDeviceSync(

  // CHECK-SAME: recv_key = "host_compute_channel_recv"
  // CHECK-SAME: send_key = "host_compute_channel_send"
  // CHECK-SAME: shape_inference_graph = @host_func
  // CHECK-SAME: shapes = [#tf_type.shape<*>, #tf_type.shape<3x?>]
  // CHECK-SAME: tpu_core = 0 : i64
  // CHECK: func @host_func
  // CHECK: %[[RECV_OUTPUT:[0-9]*]]:2 = "tf._XlaRecvAtHost"
  // CHECK-SAME: key = "host_compute_channel_send"

  return tensorflow::wrap(new tensorflow::PartialTensorShape());
}

int TF_ShapeDims(const TF_Shape* shape) {
  return tensorflow::unwrap(shape)->dims();
}

int64_t TF_ShapeDimSize(const TF_Shape* shape, int d) {
  return tensorflow::unwrap(shape)->dim_size(d);
}

void TF_DeleteShape(TF_Shape* shape) { delete tensorflow::unwrap(shape); }

namespace tensorflow {

std::string AbstractTensorHandle::DebugString() const {
  PartialTensorShape shape;
  Status s = Shape(&shape);
  std::string shape_string;
  if (!s.ok()) {
    shape_string = "<error computing shape>";
  } else {
    shape_string = shape.DebugString();
  }
  return absl::StrCat("TensorHandle(shape=", shape_string,
                      ", dtype=", DataType_Name(DataType()),

// Converts a shape from MLIR to a TensorFlow tensor shape proto.
void ConvertToTensorShapeProto(llvm::ArrayRef<int64_t> shape,
                               TensorShapeProto* output_shape);

// Converts an MLIR type to a TensorFlow tensor shape.
PartialTensorShape ConvertTypeToTensorShape(const mlir::Type& type);

// Converts an MLIR shaped type to a TensorFlow shape attribute.

    {
      "builtin_code": "SPLIT"
    }
  ],
  "subgraphs": [
    {
      "tensors": [
        {
          "shape": [
            1
          ],
          "name": "split_dim",
          "quantization": {
          }
        },
        {
          "shape": [
            256,
            32,
            32,
            3
          ],
          "name": "input",

	const expectedLocation = "runtime/pprof.nonRecursiveGenericAllocFunction[go.shape.struct {},go.shape.struct { runtime/pprof.buf [128]uint8 }];runtime/pprof.nonRecursiveGenericAllocFunction[go.shape.struct...

  values: {
    tuple_value: {
      values: {
        tensor_spec_value: {
          name : "x"
          shape: {
            dim: {
              size: 1
            }
          }
          dtype: DT_FLOAT
        }
      }
      values: {
        tensor_spec_value: {
          name : "y"
          shape: {
            dim: {
              size: 1
            }
          }
          dtype: DT_FLOAT

// CHECK:           [[LHSSHAPE:%.*]] = shape.shape_of [[LHS]] : tensor<1x4x2xf32>
// CHECK:           [[RHSSHAPE:%.*]] = shape.shape_of [[RHS]] : tensor<3x2x4xf32>
// CHECK:           [[CM2:%.*]] = arith.constant -2 : index
// CHECK:           [[LHSHEAD:%.*]], [[LHSTAIL:%.*]] = "shape.split_at"([[LHSSHAPE]], [[CM2]])
// CHECK:           [[RHSHEAD:%.*]], [[RHSTAIL:%.*]] = "shape.split_at"([[RHSSHAPE]], [[CM2]])

  // This is our 1D tensor of varying dtype.
  std::vector<typename TypeParam::type> value = {42, 100, 0, 1, 4, 29};
  // Shape is Rank 1 vector.
  std::vector<int64_t> shape;
  shape.push_back(value.size());

  Tensor original_tensor = Tensor::FromBuffer(
      /*dtype=*/dtype, /*shape=*/shape,
      /*data=*/value.data(),
      /*len=*/value.size() * sizeof(typename TypeParam::type),