ONNX export: Comparison operators

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

@@ -1565,3 +1565,24 @@ def convert_hardsigmoid(node, **kwargs):
         name=name
     )
     return [node]
+
+@mx_op.register("broadcast_lesser")
+def convert_lesser(node, **kwargs):
+    """Map MXNet's broadcast_lesser operator attributes to onnx's Less operator
+    and return the created node.
+    """
+    return create_basic_op_node('Less', node, kwargs)
+
+@mx_op.register("broadcast_greater")
+def convert_greater(node, **kwargs):
+    """Map MXNet's broadcast_greater operator attributes to onnx's Greater operator
+    and return the created node.
+    """
+    return create_basic_op_node('Greater', node, kwargs)
+
+@mx_op.register("broadcast_equal")
+def convert_equal(node, **kwargs):
+    """Map MXNet's broadcast_equal operator attributes to onnx's Equal operator
+    and return the created node.
+    """
+    return create_basic_op_node('Equal', node, kwargs)

tests/python-pytest/onnx/export/mxnet_export_test.py

@@ -238,6 +238,36 @@ def test_square():
 
     npt.assert_almost_equal(result, numpy_op)
 
+
+@with_seed()
+def test_comparison_ops():
+    """Test greater, lesser, equal"""
+    def test_ops(op_name, inputs, input_tensors, numpy_op):
+        outputs = [helper.make_tensor_value_info("output", TensorProto.FLOAT, shape=np.shape(inputs[0]))]
+        nodes = [helper.make_node(op_name, ["input"+str(i+1) for i in range(len(inputs))], ["output"])]
+        graph = helper.make_graph(nodes,
+                                  op_name + "_test",
+                                  input_tensors,
+                                  outputs)
+        model = helper.make_model(graph)
+        bkd_rep = backend.prepare(model)
+        output = bkd_rep.run(inputs)
+        npt.assert_almost_equal(output[0], numpy_op)
+
+    input_data = [np.random.rand(1, 3, 4, 5).astype("float32"),
+                  np.random.rand(1, 5).astype("float32")]
+    input_tensor = []
+    for idx, ip in enumerate(input_data):
+        input_tensor.append(helper.make_tensor_value_info("input" + str(idx + 1),
+                                                          TensorProto.FLOAT, shape=np.shape(ip)))
+    test_ops("Greater", input_data, input_tensor,
+             np.greater(input_data[0], input_data[1]).astype(np.float32))
+    test_ops("Less", input_data, input_tensor,
+             np.less(input_data[0], input_data[1]).astype(np.float32))
+    test_ops("Equal", input_data, input_tensor,
+             np.equal(input_data[0], input_data[1]).astype(np.float32))
+
+
 if __name__ == '__main__':
     test_models("bvlc_googlenet", (1, 3, 224, 224), (1, 1000))
     test_models("bvlc_reference_caffenet", (1, 3, 224, 224), (1, 1000))