Size check

Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby · Sep 13, 2023 · 5eea907 · 5eea907
1 parent 5d702f8
commit 5eea907
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/onnx/defs/math/defs.cc b/onnx/defs/math/defs.cc
@@ -3020,7 +3020,8 @@ ONNX_OPERATOR_SET_SCHEMA(
             // Skip if axis is not known
             return;
           }
-          if (axis_tensor->dims_size() != 1) {
+          // TODO(justinchuby): Create invariance checking functions to ensure shapes and sizes
+          if (axis_tensor->dims_size() != 0) {
             fail_shape_inference("axis input must be a scalar.");
           }
           const int64_t axis = ONNX_NAMESPACE::defs::math::utils::GetScalarValueFromTensor<int64_t>(axis_tensor);