Add cast to same type before equal operator

Signed-off-by: bedapisl <[email protected]>

Author: bedapisl

tests/test_backend.py

@@ -3602,6 +3602,23 @@ def test_conv2d_1_kernel_as_input(self):
                           [1., 1., 4.]], dtype=np.float32).reshape(_KERNEL3x3)
         self._conv_kernel_as_input_test(x_val, w_val)
 
+    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")