#include "tensorflow/c/tf_shape_internal.h"
#include "tensorflow/core/framework/tensor_shape.h"

extern "C" {

TF_Shape* TF_NewShape() {
  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);
}

// calling TF_DeleteShape to deallocate and destroy the returned shape.
TF_CAPI_EXPORT extern TF_Shape* TF_NewShape();

// Returns the rank of `shape`. If `shape` has unknown rank, returns -1.
TF_CAPI_EXPORT extern int TF_ShapeDims(const TF_Shape* shape);

// Returns the `d`th dimension of `shape`. If `shape` has unknown rank,
// invoking this function is undefined behavior. Returns -1 if dimension is
// unknown.

  // Input "a" is a scalar, float32 tensor
  EXPECT_EQ("a", std::string(TF_SignatureDefParamName(param_a)));
  EXPECT_EQ(TF_FLOAT, TF_TensorSpecDataType(tensor_spec_a));
  EXPECT_EQ(0, TF_ShapeDims(shape_a));

  const TF_SignatureDefParam* param_b = TF_SignatureDefParamListGet(args, 1);
  const TF_TensorSpec* tensor_spec_b = TF_SignatureDefParamTensorSpec(param_b);