Merge pull request #663 from onnx/gs/fix-equal

opset-11 supports equal for all types

Author: wayuanho

tests/test_backend.py

@@ -808,6 +808,16 @@ def test_less_unsupport_type(self):
         _ = tf.identity(mi, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val1, _INPUT1: x_val2})
 
+    @check_opset_min_version(11, "Equal")
+    def test_equal_float(self):
+        x_val1 = np.array([0., 1., 2., 3., 4., -1., -2], dtype=np.float32)
+        x_val2 = np.array([0., 1., 2.1, 3.5, 4.6, -1.1, -2.9], dtype=np.float32)
+        x1 = tf.placeholder(tf.float32, x_val1.shape, name=_TFINPUT)
+        x2 = tf.placeholder(tf.float32, x_val2.shape, 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_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))

tf2onnx/onnx_opset/logical.py

@@ -79,6 +79,17 @@ def version_7(cls, ctx, node, **kwargs):
             ctx.copy_shape(output_name, not_node.output[0])
             ctx.copy_dtype(output_name, not_node.output[0])
 
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        # starting with opset-11, equal supports all types
+        need_not = node.type == "NotEqual"
+        if need_not:
+            node.type = "Equal"
+            output_name = node.output[0]
+            not_node = ctx.insert_new_node_on_output("Not", output_name, name=utils.make_name(node.name))
+            ctx.copy_shape(output_name, not_node.output[0])
+            ctx.copy_dtype(output_name, not_node.output[0])
+
 
 @tf_op(["Greater", "Less"])
 class GreaterLess: