ONNX export: Shape operator

apache · Nov 19, 2018 · acf64f2 · acf64f2
1 parent 940b6ec
commit acf64f2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -1562,3 +1562,10 @@ def convert_sum(node, **kwargs):
             name=name
         )
     return [node]
+
+@mx_op.register("shape_array")
+def convert_shape(node, **kwargs):
+    """Map MXNet's shape_array operator attributes to onnx's Shape operator
+    and return the created node.
+    """
+    return create_basic_op_node('Shape', node, kwargs)
diff --git a/tests/python-pytest/onnx/export/onnx_backend_test.py b/tests/python-pytest/onnx/export/onnx_backend_test.py
@@ -95,7 +95,8 @@
     'test_clip'
     'test_cast',
     'test_depthtospace',
-    'test_instancenorm'
+    'test_instancenorm',
+    'test_shape'
     ]
 
 BASIC_MODEL_TESTS = [