Add cast to same type before equal operator

Author: bedapisl

tests/test_equal.py

@@ -0,0 +1,39 @@
+# Licensed under the MIT license.
+
+"""Unit tests for equal"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+import tensorflow as tf
+
+from backend_test_base import Tf2OnnxBackendTestBase
+
+
+# pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
+# pylint: disable=abstract-method,arguments-differ
+
+class EqualTests(Tf2OnnxBackendTestBase):
+
+    def test_equal_with_different_parameters(self):
+        input_val = np.array([5], dtype=np.int32)
+
+        def func(input_val):
+            tensor = tf.zeros(input_val)
+            input_size = tf.size(tensor)
+            constant = tf.constant(3, dtype=tf.int32)
+            return tf.math.equal(input_size, constant, name="output")
+
+        feed_dict = {"input:0": input_val}
+        input_names_with_port = ["input:0"]
+        output_names_with_port = ["output:0"]
+
+        current_opset = self.config.opset
+        self.config.opset = 12
+        self.run_test_case(func, feed_dict, input_names_with_port, output_names_with_port)
+        self.config.opset = current_opset
+
+if __name__ == '__main__':
+    unittest_main()

tf2onnx/onnx_opset/logical.py

@@ -34,6 +34,16 @@ def _add_cast_to_inputs(graph, node, supported_dtypes, target_dtype):
             graph.set_dtype(inp_cast.output[0], target_dtype)
 
 
+def _add_cast_to_same_type_to_inputs(graph, node):
+    common_dtype = graph.get_dtype(node.input[0])
+
+    for inp in node.input[1:]:
+        if graph.get_dtype(inp) != common_dtype:
+            inp_cast = graph.insert_new_node_on_input(node, "Cast", inp, to=common_dtype)
+            graph.copy_shape(inp, inp_cast.output[0])
+            graph.set_dtype(inp_cast.output[0], common_dtype)
+
+
 @tf_op("LogicalNot", onnx_op="Not")
 class DirectOp:
     @classmethod
@@ -81,7 +91,8 @@ def version_7(cls, ctx, node, **kwargs):
 
     @classmethod
     def version_11(cls, ctx, node, **kwargs):
-        # starting with opset-11, equal supports all types
+        # starting with opset-11, equal supports all types (but both operands must be of the same type)
+        _add_cast_to_same_type_to_inputs(ctx, node)
         need_not = node.type == "NotEqual"
         if need_not:
             node.type = "Equal"

tf2onnx/onnx_opset/tensor.py

@@ -61,7 +61,7 @@ def _wrap_concat_with_cast(ctx, node):
 class Size:
     @classmethod
     def version_1(cls, ctx, node, **kwargs):
-        pass
+        ctx.set_dtype(node.output[0], onnx_pb.TensorProto.INT64)
 
 
 @tf_op("Flatten")