ONNX export: Gather

apache · Dec 11, 2018 · 3afe2fe · 3afe2fe
1 parent a4d88ef
commit 3afe2fe
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git a/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py b/python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
@@ -1677,3 +1677,23 @@ def convert_multinomial(node, **kwargs):
         name=name,
     )
     return [node]
+
+
+@mx_op.register("take")
+def convert_gather(node, **kwargs):
+    """Map MXNet's size_array operator attributes to onnx's Size operator
+    and return the created node.
+    """
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    axis = int(attrs.get('axis', '0'))
+
+    node = onnx.helper.make_node(
+        "Gather",
+        input_nodes,
+        [name],
+        axis=axis,
+        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
@@ -98,7 +98,8 @@
     'test_hardsigmoid',
     'test_instancenorm',
     'test_shape',
-    'test_size'
+    'test_size',
+    'test_gather'
     ]
 
 BASIC_MODEL_TESTS = [

diff --git a/tests/python-pytest/onnx/import/test_cases.py b/tests/python-pytest/onnx/import/test_cases.py
@@ -86,7 +86,8 @@
     'test_operator_params',
     'test_operator_permute2',
     'test_depthtospace',
-    'test_size'
+    'test_size',
+    'test_gather'
     ]
 
 BASIC_MODEL_TESTS = [