[v1.x] Onnx support for reshape_like

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

@@ -3050,3 +3050,104 @@ def convert_contrib_AdaptiveAvgPooling2D(node, **kwargs):
         make_node("GlobalAveragePool", [input_nodes[0]], [name], name=name)
     ]
     return nodes
+
+
+@mx_op.register("reshape_like")
+def convert_reshape_like(node, **kwargs):
+    """Map MXNet's reshape_like operator attributes to onnx's operator.
+    """
+    from onnx.helper import make_node
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    lhs = input_nodes[0]
+    rhs = input_nodes[1]
+
+    lhs_begin = str(attrs.get('lhs_begin', '0'))
+    rhs_begin = str(attrs.get('rhs_begin', '0'))
+    lhs_end = str(attrs.get('lhs_end', 'None'))
+    rhs_end = str(attrs.get('rhs_end', 'None'))
+
+    if lhs_begin == 'None' or rhs_begin == 'None':
+        raise NotImplementedError("lhs_begin and rhs_begin should not be None.")
+
+    lhs_begin = int(lhs_begin)
+    rhs_begin = int(rhs_begin)
+
+    # basic case
+    if lhs_begin == 0 and lhs_end == 'None' and rhs_begin == 0 and rhs_end == 'None':
+        nodes = [
+            make_node('Shape', [rhs], [name+'_shape_rhs']),
+            make_node('Reshape', [lhs, name+'_shape_rhs'], [name], name=name)
+        ]
+        return nodes
+
+    nodes = [
+        create_tensor([0], name+'_0', kwargs["initializer"]),
+        make_node('Shape', [lhs], [name+'_lhs_shape']),
+        make_node('Shape', [name+'_lhs_shape'], [name+'_lhs_dim']),
+        make_node('Shape', [rhs], [name+'_rhs_shape']),
+        make_node('Shape', [name+'_rhs_shape'], [name+'_rhs_dim']),
+    ]
+
+    if lhs_begin >= 0:
+        nodes += [
+            create_tensor([lhs_begin], name+'_lhs_begin', kwargs["initializer"]),
+        ]
+    else:
+        nodes += [
+            create_tensor([lhs_begin], name+'_lhs_begin_neg', kwargs["initializer"]),
+            make_node('Add', [name+'_lhs_dim', name+'_lhs_begin_neg'], [name+'_lhs_begin']),
+        ]
+
+    if rhs_begin >= 0:
+        nodes += [
+            create_tensor([rhs_begin], name+'_rhs_begin', kwargs["initializer"]),
+        ]
+    else:
+        nodes += [
+            create_tensor([rhs_begin], name+'_rhs_begin_neg', kwargs["initializer"]),
+            make_node('Add', [name+'_rhs_dim', name+'_rhs_begin_neg'], [name+'_rhs_begin']),
+        ]
+
+    if lhs_end == 'None':
+        nodes += [
+            make_node('Add', [name+'_lhs_dim', name+'_0'], [name+'_lhs_end']),
+        ]
+    else:
+        lhs_end = int(lhs_end)
+        if lhs_end >= 0:
+            nodes += [
+                create_tensor([lhs_end], name+'_lhs_end', kwargs["initializer"]),
+            ]
+        else:
+            nodes += [
+                create_tensor([lhs_end], name+'_lhs_end_neg', kwargs["initializer"]),
+                make_node('Add', [name+'_lhs_dim', name+'_lhs_end_neg'], [name+'_lhs_end']),
+            ]
+
+    if rhs_end == 'None':
+        nodes += [
+            make_node('Add', [name+'_rhs_dim', name+'_0'], [name+'_rhs_end']),
+        ]
+    else:
+        rhs_end = int(rhs_end)
+        if rhs_end >= 0:
+            nodes += [
+                create_tensor([rhs_end], name+'_rhs_end', kwargs["initializer"]),
+            ]
+        else:
+            nodes += [
+                create_tensor([rhs_end], name+'_rhs_end_neg', kwargs["initializer"]),
+                make_node('Add', [name+'_rhs_dim', name+'_rhs_end_neg'], [name+'_rhs_end']),
+            ]
+
+    nodes += [
+        make_node('Slice', [name+'_lhs_shape', name+'_0', name+'_lhs_begin'], [name+'_slice0_out']),
+        make_node('Slice', [name+'_rhs_shape', name+'_rhs_begin', name+'_rhs_end'], [name+'_slice1_out']),
+        make_node('Concat', [name+'_slice0_out', name+'_slice1_out'], [name+'_concat0_out'], axis=0),
+        make_node('Slice', [name+'_lhs_shape', name+'_lhs_end', name+'_lhs_dim'], [name+'_slice2_out']),
+        make_node('Concat', [name+'_concat0_out', name+'_slice2_out'], [name+'_concat1_out'], axis=0),
+        make_node('Reshape', [lhs, name+'_concat1_out'], [name], name=name)
+    ]
+
+    return nodes

tests/python-pytest/onnx/test_operators.py

@@ -434,3 +434,21 @@ def test_onnx_export_contrib_AdaptiveAvgPooling2D(tmp_path, dtype):
     op_export_test('contrib_AdaptiveAvgPooling2D', M3, [x], tmp_path)
     M4 = def_model('contrib.AdaptiveAvgPooling2D', output_size=[1,1])
     op_export_test('contrib_AdaptiveAvgPooling2D', M4, [x], tmp_path)
+
+
+@pytest.mark.parametrize('dtype', ['int32', 'int64', 'float16', 'float32', 'float64'])
+def test_onnx_export_reshape_like(tmp_path, dtype):
+    if 'int' in dtype:
+        x = mx.nd.random.randint(0, 10, (2, 2, 3, 2), dtype=dtype)
+        y = mx.nd.random.randint(0, 10, (1, 4, 3, 2), dtype=dtype)
+    else:
+        x = mx.nd.random.normal(0, 10, (2, 2, 3, 2), dtype=dtype)
+        y = mx.nd.random.normal(0, 10, (1, 4, 3, 2), dtype=dtype)
+    M1 = def_model('reshape_like')
+    op_export_test('reshape_like1', M1, [x, y], tmp_path)
+    M2 = def_model('reshape_like', lhs_begin=0, lhs_end=2, rhs_begin=1, rhs_end=2)
+    op_export_test('reshape_like2', M2, [x, y], tmp_path)
+    M3 = def_model('reshape_like', lhs_begin=-4, lhs_end=-2, rhs_begin=-3, rhs_end=-2)
+    op_export_test('reshape_like3', M3, [x, y], tmp_path)
+    M4 = def_model('reshape_like', lhs_begin=0, lhs_end=None, rhs_begin=1, rhs_end=None)
+    op_export_test('reshape_like4', M4, [x, y], tmp_path)