ONNX export: Gather

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

@@ -2080,3 +2080,21 @@ def convert_topk(node, **kwargs):
     )
 
     return [topk_node]
+
+
+@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]

tests/python-pytest/onnx/test_cases.py

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