// Return a tensor handle containing 2D matrix containing given data and
// dimensions
TFE_TensorHandle* TestMatrixTensorHandleWithInput(TFE_Context* ctx,
                                                  float data[], int64_t dims[],
                                                  int num_dims);

// Get a Matrix TensorHandle with given float values and dimensions
TFE_TensorHandle* TestTensorHandleWithDimsFloat(TFE_Context* ctx, float data[],

class ImmediateExecutionTensorHandle : public AbstractTensorHandle {
 public:
  // Returns number of dimensions.
  virtual absl::Status NumDims(int* num_dims) const = 0;
  // Returns number of elements across all dimensions.
  virtual absl::Status NumElements(int64_t* num_elements) const = 0;
  // Returns size of specified dimension
  //
  // -1 indicates an unknown axis length; this is unreachable for most standard

// Returns the number of elements in dimension `dim_index`.
// Tensor representation on device can be transposed from its representation
// on host. The data contained in dimension `dim_index` on device
// can correspond to the data contained in another dimension in on-host
// representation. The dimensions are indexed using the standard TensorFlow
// major-to-minor order (slowest varying dimension first),
// not the XLA's minor-to-major order.

void Range(vector<int32_t>* data, int32_t start, int32_t end,
           int32_t step = 1) {
  for (int32_t i = start; i < end; i += step) {
    (*data)[i] = i;
  }
}

// Fills out_dims with the dimensions of the given tensor.
void GetDims(const TF_Tensor* t, int64_t* out_dims) {
  int num_dims = TF_NumDims(t);
  for (int i = 0; i < num_dims; i++) {
    out_dims[i] = TF_Dim(t, i);
  }
}

    TF_Status* status);

// Information about the shape of a Tensor and its type.
struct TF_ShapeAndType {
  // Number of dimensions. -1 indicates unknown rank.
  int num_dims;
  // Array of dimensions. -1 indicates unknown dim.
  int64_t* dims;
  // The data type. May be 0 to denote unknown type.
  TF_DataType dtype;
};

typedef struct TF_ShapeAndType TF_ShapeAndType;

// entries (e.g., the return value of TF_GraphGetTensorNumDims).
//
// If the number of dimensions in the shape is unknown or the shape is
// a scalar, `dims` will remain untouched. Otherwise, each element of
// `dims` will be set corresponding to the size of the dimension. An
// unknown dimension is represented by `-1`.
//
// Returns an error into `status` if:
//   * `output` is not in `graph`.

  if (num_dims > TensorShape::MaxDimensions()) {
    return errors::InvalidArgument("Value specified for `", attr_name, "` has ",
                                   num_dims,
                                   " dimensions which is over the limit of ",
                                   TensorShape::MaxDimensions(), ".");
  }
  TensorShapeProto proto;
  if (num_dims < 0) {
    proto.set_unknown_rank(true);
  } else {

                    /*input_shapes*/ {make_shape({3, 2}), unknown_shape()},
                    /*input_tensors*/ {},
                    /*expected_shape*/ make_shape({3, kUnknownDim}));

  // Infer shape when some dimensions are unknown.
  CheckOutputShapes(
      matmul_op,
      /*input_shapes*/ {make_shape({kUnknownDim, 2}), make_shape({2, 4})},
      /*input_tensors*/ {},
      /*expected_shape*/ make_shape({kUnknownDim, 4}));

  bool valid = true;
  int64_t expected_stride = 1;
  for (int i = ndim - 1; i >= 0; --i) {
    // Empty tensors are always compact regardless of strides.
    if (shape_arr[i] == 0) return true;
    // Note that dimensions with size=1 can have any stride.
    if (shape_arr[i] != 1 && stride_arr[i] != expected_stride) {
      valid = false;
    }
    expected_stride *= shape_arr[i];
  }
  return valid;
}
}  // namespace

void ParallelTensorDeallocator(void* data) {
  delete reinterpret_cast<ParallelTensor*>(data);
}

// Used as an argument to TFE_NewCustomDeviceTensorHandle, for computing the
// number of dimensions of a parallel tensor.
int ParallelTensorNumDims(void* data, TF_Status* status) {
  const std::vector<int64_t>* shape;
  absl::Status s = reinterpret_cast<ParallelTensor*>(data)->Shape(&shape);
  if (!s.ok()) {