ONNX export: Gather

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

@@ -1457,6 +1457,23 @@ def convert_dynamic_reshape(node, **kwargs):
     """
     return create_basic_op_node('Reshape', node, kwargs)
 
+@mx_op.register("take")
+def convert_take(node, **kwargs):
+    """Map MXNet's Take operator attributes to onnx's Gather operator.
+    """
+    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]
+
 # Changing shape and type.
 @mx_op.register("Reshape")
 def convert_reshape(node, **kwargs):

tests/python-pytest/onnx/test_cases.py

@@ -78,10 +78,10 @@
              'test_max_',
              'test_softplus',
              'test_reduce_',
-             'test_split_equal'
+             'test_split_equal',
+             'test_gather'
              ],
-    'import': ['test_gather',
-               'test_softsign',
+    'import': ['test_softsign',
                'test_mean',
                'test_averagepool_1d',
                'test_averagepool_2d_pads_count_include_pad',