Clarify NonZero behavior for scalar input in spec (onnx#4113)

jcwchen · gramalingam · web-flow · commit 31b6337dda5e · 2022-05-18T03:29:05.000Z
* bump NonZero with updated shape inference

Signed-off-by: Chun-Wei Chen &lt;jacky82226@gmail.com&gt;

* revert and add tests for NonZero and Unsqueeze

Signed-off-by: Chun-Wei Chen &lt;jacky82226@gmail.com&gt;

* clarify in spec

Signed-off-by: Chun-Wei Chen &lt;jacky82226@gmail.com&gt;

* update spec

Signed-off-by: Chun-Wei Chen &lt;jacky82226@gmail.com&gt;

* use extern to handle shared string across files

Signed-off-by: Chun-Wei Chen &lt;jacky82226@gmail.com&gt;

Co-authored-by: G. Ramalingam &lt;grama@microsoft.com&gt;
diff --git a/docs/Changelog.md b/docs/Changelog.md
@@ -8648,7 +8648,8 @@ This version of the operator has been available since version 9 of the default O
   Returns the indices of the elements that are non-zero
       (in row-major order - by dimension).
       NonZero behaves similar to numpy.nonzero:
-      https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html
+      https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html,
+      but for scalar input, NonZero produces output shape (0, N) instead of (1, N), which is different from Numpy's behavior.
 
 #### Version
 
@@ -16905,7 +16906,8 @@ This version of the operator has been available since version 13 of the default
   Returns the indices of the elements that are non-zero
       (in row-major order - by dimension).
       NonZero behaves similar to numpy.nonzero:
-      https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html
+      https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html,
+      but for scalar input, NonZero produces output shape (0, N) instead of (1, N), which is different from Numpy's behavior.
 
 #### Version
 
diff --git a/docs/Operators.md b/docs/Operators.md
@@ -13418,7 +13418,8 @@ expect(node, inputs=[boxes, scores, max_output_boxes_per_class, iou_threshold, s
   Returns the indices of the elements that are non-zero
       (in row-major order - by dimension).
       NonZero behaves similar to numpy.nonzero:
-      https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html
+      https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html,
+      but for scalar input, NonZero produces output shape (0, N) instead of (1, N), which is different from Numpy's behavior.
 
 #### Version
 
diff --git a/onnx/defs/tensor/defs.cc b/onnx/defs/tensor/defs.cc
@@ -3108,18 +3108,11 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
 
-static const char* NonZero_ver13_doc = R"DOC(
-    Returns the indices of the elements that are non-zero
-    (in row-major order - by dimension).
-    NonZero behaves similar to numpy.nonzero:
-    https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html
-)DOC";
-
 ONNX_OPERATOR_SET_SCHEMA(
     NonZero,
     13,
     OpSchema()
-        .SetDoc(NonZero_ver13_doc)
+        .SetDoc(NonZero_ver9_doc)
         .Input(
             0,
             "X",
diff --git a/onnx/defs/tensor/old.cc b/onnx/defs/tensor/old.cc
@@ -2203,11 +2203,12 @@ ONNX_OPERATOR_SET_SCHEMA(
           }
         }));
 
-static const char* NonZero_ver9_doc = R"DOC(
+const char* NonZero_ver9_doc = R"DOC(
     Returns the indices of the elements that are non-zero
     (in row-major order - by dimension).
     NonZero behaves similar to numpy.nonzero:
-    https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html
+    https://docs.scipy.org/doc/numpy/reference/generated/numpy.nonzero.html,
+    but for scalar input, NonZero produces output shape (0, N) instead of (1, N), which is different from Numpy's behavior.
 )DOC";
 
 ONNX_OPERATOR_SET_SCHEMA(
diff --git a/onnx/defs/tensor/utils.h b/onnx/defs/tensor/utils.h
@@ -31,4 +31,6 @@ void resizeShapeInferenceHelper_opset7_to_10(
     const TensorShapeProto& input_shape,
     const std::vector<float>& scales_data,
     TensorShapeProto* output_shape);
+
+extern const char* NonZero_ver9_doc;
 } // namespace ONNX_NAMESPACE
diff --git a/onnx/test/shape_inference_test.py b/onnx/test/shape_inference_test.py
@@ -69,6 +69,7 @@ def _compare_value_infos(self, vi_type: TypeProto, inferred_vi_type: TypeProto)
             assert vi_type.tensor_type.elem_type == inferred_vi_type.tensor_type.elem_type
             assert vi_type.tensor_type.HasField('shape') == inferred_vi_type.tensor_type.HasField('shape')
             if vi_type.tensor_type.HasField('shape'):
+                assert len(vi_type.tensor_type.shape.dim) == len(inferred_vi_type.tensor_type.shape.dim)
                 for dim_i in range(len(vi_type.tensor_type.shape.dim)):
                     dim = vi_type.tensor_type.shape.dim[dim_i]
                     inferred_dim = inferred_vi_type.tensor_type.shape.dim[dim_i]
@@ -632,6 +633,15 @@ def test_unsqueeze_negative_axes(self) -> None:
             initializer=[make_tensor('axes', TensorProto.INT64, (2,), (0, -1))])
         self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.FLOAT, (1, 3, 4, 5, 1))])
 
+    def test_unsqueeze_scalar(self) -> None:
+        graph = self._make_graph(
+            [('x', TensorProto.FLOAT, ()),
+            ('axes', TensorProto.INT64, ())],
+            [make_node('Unsqueeze', ['x', 'axes'], 'y')],
+            [],
+            initializer=[make_tensor('axes', TensorProto.INT64, (), (-1,))])
+        self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.FLOAT, (1,))])
+
     def test_slice_without_input_shape(self) -> None:
         graph = self._make_graph(
             [('x', TensorProto.FLOAT, (3, 2)), ('starts', TensorProto.INT64, (1,)), ('ends', TensorProto.INT64, (1,))],
@@ -3816,6 +3826,13 @@ def test_nonzero_existing_dim_param(self) -> None:
             [make_tensor_value_info('y', TensorProto.INT64, (None, 'NZ'))])
         self._assert_inferred(graph, [make_tensor_value_info('y', TensorProto.INT64, (1, 'NZ'))])  # type: ignore
 
+    def test_nonzero_scalar(self) -> None:
+        graph = self._make_graph(
+            [('x', TensorProto.FLOAT, ())],
+            [make_node('NonZero', ['x'], ['out'])],
+            [])
+        self._assert_inferred(graph, [make_tensor_value_info('out', TensorProto.INT64, (0, None))])  # type: ignore
+
     def test_optional_construct_empty_tensor(self) -> None:
         tensor_type_proto = helper.make_tensor_type_proto(elem_type=TensorProto.FLOAT, shape=[1, 2, 3])
         optional_type_proto = helper.make_optional_type_proto(tensor_type_proto)
