[v1.x] ONNX support rewrite norm (#20195)

* fix norm

* fix sanity

Co-authored-by: Wei Chu <weichu@amazon.com>

Author: waytrue17

python/mxnet/onnx/mx2onnx/_op_translations/_op_translations_opset12.py

@@ -2393,6 +2393,7 @@ def convert_norm(node, **kwargs):
     """Map MXNet's norm operator attributes to onnx's ReduceL1 and ReduceL2 operators
     and return the created node.
     """
+    from onnx.helper import make_node
     name, input_nodes, attrs = get_inputs(node, kwargs)
 
     mx_axis = attrs.get("axis", None)
@@ -2407,24 +2408,32 @@ def convert_norm(node, **kwargs):
     onnx_op_name = "ReduceL1" if ord == 1 else "ReduceL2"
 
     if axes:
-        reduce_node = onnx.helper.make_node(
-            onnx_op_name,
-            input_nodes,
-            [name],
-            axes=axes,
-            keepdims=keepdims,
-            name=name
-        )
-        return [reduce_node]
+        if keepdims:
+            reduce_node = make_node(onnx_op_name, input_nodes, [name], axes=axes, keepdims=keepdims)
+            return [reduce_node]
+        else:
+            create_tensor([1], name+'_1', kwargs['initializer'])
+            nodes = [
+                make_node(onnx_op_name, input_nodes, [name+'_norm'], axes=axes, keepdims=keepdims),
+                make_node('Shape', [name+'_norm'], [name+'_norm_shape']),
+                make_node('Concat', [name+'_1', name+'_norm_shape'], [name+'_concat'], axis=0),
+                make_node('Reshape', [name+'_norm', name+'_concat'], [name+'_reshape']),
+                make_node('Squeeze', [name+'_reshape'], [name], axes=[0]),
+            ]
+            return nodes
     else:
-        reduce_node = onnx.helper.make_node(
-            onnx_op_name,
-            input_nodes,
-            [name],
-            keepdims=keepdims,
-            name=name
-        )
-        return [reduce_node]
+
+        if keepdims:
+            reduce_node = make_node(onnx_op_name, input_nodes, [name], keepdims=keepdims)
+            return [reduce_node]
+        else:
+            create_tensor([1], name+'_1', kwargs['initializer'])
+            nodes = [
+                make_node(onnx_op_name, input_nodes, [name+'_norm'], keepdims=keepdims),
+                make_node('Reshape', [name+'_norm', name+'_1'], [name])
+            ]
+            return nodes
+
 
 @mx_op.register("_sample_multinomial")
 def convert_multinomial(node, **kwargs):

python/mxnet/onnx/mx2onnx/_op_translations/_op_translations_opset13.py

@@ -1544,3 +1544,57 @@ def convert_squeeze(node, **kwargs):
             name=name,
         )
     return [node]
+
+
+@mx_op.register("norm", OPSET_VERSION)
+def convert_norm(node, **kwargs):
+    """Map MXNet's norm operator attributes to onnx's ReduceL1 and ReduceL2 operators
+    and return the created node.
+    """
+    from onnx.helper import make_node
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    mx_axis = attrs.get("axis", None)
+    axes = convert_string_to_list(str(mx_axis)) if mx_axis else None
+
+    keepdims = get_boolean_attribute_value(attrs, "keepdims")
+    ord = int(attrs.get("ord", 2))
+
+    onnx_op_name = "ReduceL1" if ord == 1 else "ReduceL2"
+
+    if axes:
+        if keepdims:
+            reduce_node = make_node(onnx_op_name, input_nodes, [name], axes=axes, keepdims=keepdims)
+            return [reduce_node]
+        else:
+            create_tensor([1], name+'_1', kwargs['initializer'])
+            create_tensor([0], name+'_0', kwargs['initializer'])
+            create_tensor([len(axes)], name+'_axes_dim', kwargs['initializer'])
+            nodes = [
+                make_node(onnx_op_name, input_nodes, [name+'_reduce'], axes=axes, keepdims=keepdims),
+                make_node('Shape', [name+'_reduce'], [name+'_reduce_shape']),
+                make_node('Shape', [name+'_reduce_shape'], [name+'_reduce_dim']),
+                make_node('Shape', [input_nodes[0]], [name+'_in_shape']),
+                make_node('Shape', [name+'_in_shape'], [name+'_in_dim']),
+                make_node('Equal', [name+'_axes_dim', name+'_in_dim'], [name+'_equal']),
+                make_node('Where', [name+'_equal', name+'_1', name+'_reduce_dim'], [name+'_where0']),
+                make_node('Tile', [name+'_0', name+'_where0'], [name+'_tile']),
+                make_node('Unsqueeze', [name+'_0', name+'_0'], [name+'_unsqueeze']),
+                make_node('Where', [name+'_equal', name+'_1', name+'_0'], [name+'_where1']),
+                make_node('ScatterND', [name+'_tile', name+'_unsqueeze', name+'_where1'], [name+'_SND']),
+                make_node('Reshape', [name+'_reduce', name+'_SND'], [name]),
+            ]
+            return nodes
+    else:
+
+        if keepdims:
+            reduce_node = make_node(onnx_op_name, input_nodes, [name], keepdims=keepdims)
+            return [reduce_node]
+        else:
+            create_tensor([1], name+'_1', kwargs['initializer'])
+            nodes = [
+                make_node(onnx_op_name, input_nodes, [name+'_norm'], keepdims=keepdims),
+                make_node('Reshape', [name+'_norm', name+'_1'], [name])
+            ]
+            return nodes
+            

tests/python-pytest/onnx/test_operators.py

@@ -1488,7 +1488,6 @@ def test_onnx_export_maximum_minimum(tmp_path, dtype, shape, op_name):
     op_export_test(op_name, M, [lhs, rhs], tmp_path)
 
 
-
 # onnx reduce ops do not support float64
 @pytest.mark.parametrize('dtype', ['float16', 'float32','int32', 'int64'])
 @pytest.mark.parametrize('shape', [(2, 3), (4, 5, 6)])
@@ -1557,6 +1556,24 @@ def test_onnx_export_squeeze(tmp_path, dtype, shape_axis):
     op_export_test('squeeze', M, [x], tmp_path)
 
 
+@pytest.mark.parametrize("dtype", ["float16", "float32"])
+@pytest.mark.parametrize("order", [1, 2])
+@pytest.mark.parametrize("keepdims", [0, 1])
+@pytest.mark.parametrize("axis", [None, 0, 1, 2, -1, (0, 2), (0, 1, 2)])
+@pytest.mark.parametrize("shape", [(4, 5, 6), (3, 4, 5, 6)])
+def test_onnx_export_norm(tmp_path, dtype, order, axis, shape, keepdims):
+    kwargs = {}
+    if order is not None:
+        kwargs['ord'] = order
+    if axis is not None:
+        kwargs['axis'] = axis
+    if keepdims is not None:
+        kwargs['keepdims'] = keepdims
+    M = def_model('norm', **kwargs)
+    x = mx.random.normal(0, 10, shape).astype(dtype)
+    op_export_test('norm', M, [x], tmp_path)
+
+
 @pytest.mark.parametrize('dtype', ['float16', 'float32', 'float64', 'int32', 'int64'])
 @pytest.mark.parametrize("shape", [(10,), (2,3), (4,5,6)])
 def test_onnx_export_logical_not(tmp_path, dtype, shape):