float data[], int64_t dims[],
                                                  int num_dims) {
  TF_Status* status = TF_NewStatus();
  TF_Tensor* t =
      TFE_AllocateHostTensor(ctx, TF_FLOAT, &dims[0], num_dims, status);
  memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t));
  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);

                         DataType value, ForwardOperation*);
absl::Status SetAttrShape(AbstractOperation*, const char* attr_name,
                          const int64_t* dims, const int num_dims,
                          ForwardOperation*);
absl::Status SetAttrFunction(AbstractOperation*, const char* attr_name,
                             const AbstractOperation* value, ForwardOperation*);

    return nullptr;
  }
  tensorflow::PartialTensorShape partial_shape;
  if (shape.num_dims != -1) {
    DCHECK(shape.dim_sizes != nullptr);
    Status status = tensorflow::PartialTensorShape::MakePartialShape(
        reinterpret_cast<int64_t*>(shape.dim_sizes), shape.num_dims,
        &partial_shape);
    if (!status.ok()) {
      tsl::Set_TF_Status_from_Status(s, status);
      return nullptr;

    return nullptr;
  }

  size_t total_bytes = dl_tensor->dtype.bits / 8;
  for (int i = 0; i < num_dims; i++) {
    total_bytes *= dims[i];
  }

  if (dl_tensor->strides != nullptr &&
      !IsValidStrideCompactRowMajorData(dl_tensor->shape, dl_tensor->strides,
                                        num_dims)) {
    status->status = tensorflow::errors::InvalidArgument(
        "Invalid strides array from DLPack");

TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values) {
  int64_t num_values = 1;
  for (int i = 0; i < num_dims; ++i) {
    num_values *= dims[i];
  }
  TF_Tensor* t =
      TF_AllocateTensor(TF_INT8, dims, num_dims, sizeof(char) * num_values);
  memcpy(TF_TensorData(t), values, sizeof(char) * num_values);
  return t;
}

TF_Tensor* Int32Tensor(const int64_t* dims, int num_dims,

TEST(CAPI, TestTensorNonScalarBytesAllocateDelete) {
  const int batch_size = 4;
  const int num_dims = 2;
  int64_t* dims = new int64_t[num_dims];
  int64_t num_elements = 1;
  dims[0] = batch_size;
  dims[1] = 1;
  for (int64_t i = 0; i < num_dims; ++i) {
    num_elements *= dims[i];
  }
  TF_Tensor* t = TF_AllocateTensor(TF_STRING, dims, num_dims,
                                   sizeof(TF_TString) * num_elements);

  virtual absl::Status SetAttrType(const char* attr_name, DataType value) = 0;
  virtual absl::Status SetAttrShape(const char* attr_name, const int64_t* dims,
                                    const int num_dims) = 0;
  virtual absl::Status SetAttrShape(const char* attr_name,
                                    const PartialTensorShape shape);
  virtual absl::Status SetAttrFunction(const char* attr_name,

// Read the shape of a variable and write to `dims`
TF_CAPI_EXPORT extern void TF_CheckpointReaderGetVariableShape(
    TF_CheckpointReader* reader, const char* name, int64_t* dims, int num_dims,
    TF_Status* status);
// Get the number of dimension of a variable
TF_CAPI_EXPORT extern int TF_CheckpointReaderGetVariableNumDims(
    TF_CheckpointReader* reader, const char* name);
// Load the weight of a variable

tensorflow::shape_inference::ShapeHandle ShapeHandleFromDims(
    tensorflow::shape_inference::InferenceContext* ic, int num_dims,
    const int64_t* dims) {
  if (num_dims != -1) {
    std::vector<tensorflow::shape_inference::DimensionHandle> dim_vec;
    dim_vec.reserve(num_dims);
    for (int i = 0; i < num_dims; ++i) {
      dim_vec.push_back(ic->MakeDim(dims[i]));
    }
    return ic->MakeShape(dim_vec);
  } else {

  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  TFE_OpSetAttrType(handle_op.get(), "dtype", TF_FLOAT);
  TFE_OpSetAttrShape(handle_op.get(), "shape", /*dims=*/nullptr, /*num_dims=*/0,
                     status.get());
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  auto outputs =
      parallel_device.Execute(context.get(), std::vector<ParallelTensor*>(),