[GitHub] [incubator-mxnet] josephevans commented on a change in pull request #19775: [v1.x] Onnx export support for batch_dot

[GitBox <gi...@apache.org>]

josephevans commented on a change in pull request #19775:
URL: https://github.com/apache/incubator-mxnet/pull/19775#discussion_r562934760



##########
File path: python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
##########
@@ -3275,3 +3275,78 @@ def convert_gather_nd(node, **kwargs):
     ]
 
     return nodes
+
+
+@mx_op.register("batch_dot")
+def convert_batch_dot(node, **kwargs):
+    """Map MXNet's gather_ND operator attributes to onnx's operator.
+    """
+    from onnx.helper import make_node
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    lhs = input_nodes[0]
+    rhs = input_nodes[1]
+    input_type = kwargs['in_type']
+    dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type]
+    transpose_a = str(attrs.get('transpose_a', '0'))

Review comment:
       Would it make sense to set the defaults to these to 'False', to simplify where you check transpose_[ab] below?

##########
File path: python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
##########
@@ -3275,3 +3275,78 @@ def convert_gather_nd(node, **kwargs):
     ]
 
     return nodes
+
+
+@mx_op.register("batch_dot")
+def convert_batch_dot(node, **kwargs):
+    """Map MXNet's gather_ND operator attributes to onnx's operator.

Review comment:
       comment is wrong ;)

##########
File path: python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
##########
@@ -3275,3 +3275,78 @@ def convert_gather_nd(node, **kwargs):
     ]
 
     return nodes
+
+
+@mx_op.register("batch_dot")
+def convert_batch_dot(node, **kwargs):
+    """Map MXNet's gather_ND operator attributes to onnx's operator.
+    """
+    from onnx.helper import make_node
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    lhs = input_nodes[0]
+    rhs = input_nodes[1]
+    input_type = kwargs['in_type']
+    dtype = onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type]
+    transpose_a = str(attrs.get('transpose_a', '0'))
+    transpose_b = str(attrs.get('transpose_b', '0'))
+    perm = [0, 2, 1]
+
+    nodes = [
+        create_tensor([-2], name+'_-2', kwargs['initializer']),
+        create_tensor([-1], name+'_-1', kwargs['initializer']),
+        create_tensor([0], name+'_0', kwargs['initializer']),
+        create_tensor([2], name+'_2', kwargs['initializer']),
+        create_tensor([100], name+'_100', kwargs['initializer']),
+        create_tensor([0], name+'_0f', kwargs['initializer'], dtype=dtype),
+    ]
+
+    if transpose_a in ['0', 'False']:
+        nodes += [
+            make_node('Add', [lhs, name+'_0f'], [name+'_lhs']),
+        ]
+    else:
+        nodes += [
+            make_node('Shape', [lhs], [name+'_lhs_shape']),
+            make_node('Shape', [name+'_lhs_shape'], [name+'_lhs_dim']),
+            # make_node('Sub', [name+'_lhs_dim', name+'_2'], [name+'_lhs_sub']),

Review comment:
       remove commented out line?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org