#include "tensorflow/c/tf_shape_internal.h"

extern "C" {

TF_DataType TF_TensorSpecDataType(const TF_TensorSpec* spec) {
  return static_cast<TF_DataType>(tensorflow::unwrap(spec)->dtype());
}

const TF_Shape* TF_TensorSpecShape(const TF_TensorSpec* spec) {
  return tensorflow::wrap(&tensorflow::unwrap(spec)->shape());
}

// Returns the shape associated with the TensorSpec. The returned Shape is not
// owned by the caller. Caller must not call TF_DeleteShape on the returned
// shape.
TF_CAPI_EXPORT extern const TF_Shape* TF_TensorSpecShape(
    const TF_TensorSpec* spec);

#ifdef __cplusplus
}  // end extern "C"
#endif  // __cplusplus

  const TF_SignatureDefParam* param_a = TF_SignatureDefParamListGet(args, 0);
  const TF_TensorSpec* tensor_spec_a = TF_SignatureDefParamTensorSpec(param_a);
  const TF_Shape* shape_a = TF_TensorSpecShape(tensor_spec_a);

  // 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));