Merge pull request #431 from lucienwang1009/equal_more_dtype

support more dtype for Equal

Author: nbcsm

tests/common.py

@@ -197,6 +197,7 @@ def check_onnxruntime_incompatibility(op):
         "AveragePool": 7,  # AveragePool-1
         "Div": 7,  # Div-1, Div-6
         "Elu": 6,  # Elu-1
+        "Equal": 7,  # Equal-1
         "Exp": 6,  # Exp-1
         "Greater": 7,  # Greater-1
         "Less": 7,  # Less-1

tests/test_backend.py

@@ -672,6 +672,26 @@ def test_less_unsupport_type(self):
         _ = tf.identity(mi, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
 
+    @check_onnxruntime_incompatibility("Equal")
+    def test_equal(self):
+        x_val1 = np.array([4, 2, 4, 1], dtype=np.int32).reshape((2, 2))
+        x_val2 = np.array([2, 4, 4, 1], dtype=np.int32).reshape((2, 2))
+        x1 = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT)
+        x2 = tf.placeholder(tf.int32, [2, 2], name=_TFINPUT1)
+        mi = tf.equal(x1, x2)
+        _ = tf.identity(mi, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
+
+        tf.reset_default_graph()
+        x_val1 = np.array([4, 2, 4, 1], dtype=np.float32).reshape((2, 2))
+        x_val2 = np.array([2, 4, 4, 1], dtype=np.float32).reshape((2, 2))
+        x1 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT)
+        x2 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT1)
+        mi = tf.equal(x1, x2)
+        _ = tf.identity(mi, name=_TFOUTPUT)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
+
+
     def test_sequeeze_no_axis_specified(self):
         x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2, 1))
         x = tf.placeholder(tf.float32, [2, 2, 1], name=_TFINPUT)

tf2onnx/onnx_opset/logical.py

@@ -11,7 +11,7 @@
 
 import logging
 
-from onnx import onnx_pb
+from onnx import TensorProto
 from tf2onnx import utils
 from tf2onnx.handler import tf_op
 from tf2onnx.onnx_opset import common
@@ -21,22 +21,17 @@
 
 # pylint: disable=unused-argument,missing-docstring
 
-def logical_compare_op(ctx, node, **kwargs):
-    # T2 output = Greater(T1 x, T1 y), T2=tensor(bool)
-    # T2 output = Less(T1 x, T1 y), T2=tensor(bool)
-    # Great/Less in opset7 only supports limited types, insert Cast if needed
-    if ctx.opset < 9:
-        supported_dtypes = [
-            onnx_pb.TensorProto.FLOAT,
-            onnx_pb.TensorProto.FLOAT16,
-            onnx_pb.TensorProto.DOUBLE
-        ]
-        target_dtype = onnx_pb.TensorProto.FLOAT
+def _add_cast_to_inputs(graph, node, supported_dtypes, target_dtype):
+    is_support = True
+    for inp in node.input:
+        if graph.get_dtype(inp) not in supported_dtypes:
+            is_support = False
+            break
+    if not is_support:
         for inp in node.input:
-            if ctx.get_dtype(inp) not in supported_dtypes:
-                inp_cast = ctx.insert_new_node_on_input(node, "Cast", inp, to=target_dtype)
-                ctx.copy_shape(inp, inp_cast.output[0])
-                ctx.set_dtype(inp_cast.output[0], target_dtype)
+            inp_cast = graph.insert_new_node_on_input(node, "Cast", inp, to=target_dtype)
+            graph.copy_shape(inp, inp_cast.output[0])
+            graph.set_dtype(inp_cast.output[0], target_dtype)
 
 
 @tf_op(["LogicalNot", "NotEqual"], onnx_op="Not")
@@ -46,30 +41,56 @@ def version_4(cls, ctx, node, **kwargs):
         pass
 
 
-@tf_op(["Equal", "Greater", "Less"])
 @tf_op("LogicalAnd", onnx_op="And")
 @tf_op("LogicalOr", onnx_op="Or")
 class BroadcastOp(common.BroadcastOp):
     pass
 
 
+@tf_op("Equal")
+class Equal:
+    @classmethod
+    def version_4(cls, ctx, node, **kwargs):
+        common.BroadcastOp.version_4(ctx, node, **kwargs)
+
+    @classmethod
+    def version_7(cls, ctx, node, **kwargs):
+        # T2 output = Equal(T1, x, T1 y), T1 \in {bool, int32, int64}
+        supported_dtypes = [
+            TensorProto.BOOL,
+            TensorProto.INT32,
+            TensorProto.INT64
+        ]
+        target_dtype = TensorProto.INT32
+        _add_cast_to_inputs(ctx, node, supported_dtypes, target_dtype)
+
+
 @tf_op(["Greater", "Less"])
-class Greater:
+class GreaterLess:
     @classmethod
     def version_4(cls, ctx, node, **kwargs):
         common.BroadcastOp.version_4(ctx, node, **kwargs)
 
     @classmethod
     def version_7(cls, ctx, node, **kwargs):
-        logical_compare_op(ctx, node, **kwargs)
+        # T2 output = Greater(T1 x, T1 y), T2=tensor(bool)
+        # T2 output = Less(T1 x, T1 y), T2=tensor(bool)
+        # Great/Less in opset7 only supports limited types, insert Cast if needed
+        supported_dtypes = [
+            TensorProto.FLOAT,
+            TensorProto.FLOAT16,
+            TensorProto.DOUBLE
+        ]
+        target_dtype = TensorProto.FLOAT
+        _add_cast_to_inputs(ctx, node, supported_dtypes, target_dtype)
 
 
 @tf_op("GreaterEqual", onnx_op="Less")
 @tf_op("LessEqual", onnx_op="Greater")
 class GreaterLessEqual:
     @classmethod
     def version_7(cls, ctx, node, **kwargs):
-        logical_compare_op(ctx, node, **kwargs)
+        GreaterLess.version_7(ctx, node, **kwargs)
         output_name = node.output[0]
         new_node = ctx.insert_new_node_on_output("Not", output_name, name=utils.make_name(node.name))
         ctx.copy_shape(output_name, new_node.output[0])