ONNX export: Shape operator

apache · Oct 25, 2018 · 0394bc3 · 0394bc3
1 parent 68e0b43
commit 0394bc3
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -2300,3 +2300,24 @@ 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.
+    """
+    onnx = import_onnx_modules()
+    name = node["name"]
+    proc_nodes = kwargs["proc_nodes"]
+    inputs = node["inputs"]
+
+    input_node_id = kwargs["index_lookup"][inputs[0][0]]
+    input_node = proc_nodes[input_node_id].name
+
+    node = onnx.helper.make_node(
+        "Shape",
+        [input_node],
+        [name],
+        name=name,
+    )
+    return [node]
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 = [