ONNX export: logsoftmax

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

@@ -1655,3 +1655,23 @@ def convert_size(node, **kwargs):
     and return the created node.
     """
     return create_basic_op_node('Size', node, kwargs)
+
+
+@mx_op.register("log_softmax")
+def convert_logsoftmax(node, **kwargs):
+    """Map MXNet's log_softmax operator attributes to onnx's LogSoftMax operator
+    and return the created node.
+    """
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    # Converting to float32
+    axis = int(attrs.get("axis", 1))
+
+    node = onnx.helper.make_node(
+        'LogSoftmax',
+        input_nodes,
+        [name],
+        axis=abs(axis),
+        name=name
+    )
+    return [node]