Add data cast to fix LessOrEqual issue with tf 2.4 (onnx#839)

Add data cast to fix LessOrEqual issue with tf 2.4, since
tf LessEqual doesn't support uint16, uint32, uint64. Also
add test code to verify unsupported data types.

Signed-off-by: Chin Huang <[email protected]>

Co-authored-by: Winnie Tsang <[email protected]>

Author: chinhuang007

onnx_tf/handlers/backend/less_or_equal.py

@@ -4,12 +4,48 @@
 from onnx_tf.handlers.handler import onnx_op
 from onnx_tf.handlers.handler import tf_func
 from .control_flow_mixin import ComparisonMixin
+from onnx_tf.common import sys_config
+from onnx_tf.common import exception
+import onnx_tf.common.data_type as data_type
 
 
 @onnx_op("LessOrEqual")
 @tf_func(tf.less_equal)
 class LessOrEqual(ComparisonMixin, BackendHandler):
+  cast_map = {tf.uint16: tf.int32, tf.uint32: tf.int64}
+  supported_types = [
+      tf.uint8, tf.int8, tf.int16, tf.int32, tf.int64, tf.float16,
+      tf.float32, tf.float64, tf.bfloat16
+  ]
+
+  @classmethod
+  def args_check(cls, node, **kwargs):
+    # update cast map based on the auto_cast config option
+    cls.cast_map[tf.uint64] = tf.int64 if sys_config.auto_cast else None
+
+    x = kwargs["tensor_dict"][node.inputs[0]]
+    y = kwargs["tensor_dict"][node.inputs[1]]
+
+    # throw an error if the data type is not natively supported by
+    # Tensorflow, cannot be safely cast, and auto_cast option is False
+    if x.dtype in cls.cast_map and cls.cast_map[x.dtype] is None:
+      exception.DTYPE_NOT_CAST_EXCEPT(
+          "LessOrEqual input " + node.inputs[0] + " with data type '" +
+          data_type.tf_to_np_str(x.dtype) + "'",
+          data_type.tf_to_np_str_list(cls.supported_types))
+    if y.dtype in cls.cast_map and cls.cast_map[y.dtype] is None:
+      exception.DTYPE_NOT_CAST_EXCEPT(
+          "LessOrEqual input " + node.inputs[1] + " with data type '" +
+          data_type.tf_to_np_str(y.dtype) + "'",
+          data_type.tf_to_np_str_list(cls.supported_types))
 
   @classmethod
   def version_12(cls, node, **kwargs):
-    return [cls.make_tensor_from_onnx_node(node, **kwargs)]
+    def dtype_cast(x):
+      return tf.cast(x, cls.cast_map[x.dtype]) if x.dtype in cls.cast_map else x
+
+    # handle data types that are not natively supported by Tensorflow
+    x = dtype_cast(kwargs["tensor_dict"][node.inputs[0]])
+    y = dtype_cast(kwargs["tensor_dict"][node.inputs[1]])
+
+    return [cls.make_tensor_from_onnx_node(node, inputs=[x, y])]

test/backend/test_node.py

@@ -1302,11 +1302,11 @@ def test_less_or_equal(self):
     node_def = helper.make_node('LessOrEqual', ['X', 'Y'], ['Z'])
     shape = [2, 3, 4, 5]
     x = self._get_rnd_int(
-        np.iinfo(np.uint64).min,
-        np.iinfo(np.uint64).max, shape, np.uint64)
+        np.iinfo(np.int64).min,
+        np.iinfo(np.int64).max, shape, np.int64)
     y = self._get_rnd_int(
-        np.iinfo(np.uint64).min,
-        np.iinfo(np.uint64).max, shape, np.uint64)
+        np.iinfo(np.int64).min,
+        np.iinfo(np.int64).max, shape, np.int64)
     output = run_node(node_def, [x, y])
     np.testing.assert_equal(output['Z'], np.less_equal(x, y))
     # test with broadcast
@@ -1319,6 +1319,22 @@ def test_less_or_equal(self):
         np.finfo(np.float16).max, shape2).astype(np.float16)
     output = run_node(node_def, [x, y])
     np.testing.assert_equal(output['Z'], np.less_equal(x, y))
+    # test data types that are not natively supported by Tensorflow
+    x = self._get_rnd_int(
+        np.iinfo(np.uint32).min,
+        np.iinfo(np.uint32).max, shape, np.uint32)
+    y = self._get_rnd_int(
+        np.iinfo(np.uint32).min,
+        np.iinfo(np.uint32).max, shape, np.uint32)
+    output = run_node(node_def, [x, y])
+    np.testing.assert_equal(output['Z'], np.less_equal(x, y))
+    x = self._get_rnd_int(
+        np.iinfo(np.uint64).min,
+        np.iinfo(np.uint64).max, shape, np.uint64)
+    y = self._get_rnd_int(
+        np.iinfo(np.uint64).min,
+        np.iinfo(np.uint64).max, shape, np.uint64)
+    self.assertRaises(RuntimeError, run_node, node_def, [x, y])
 
   def test_lp_normalization(self):
     for ordr in range(1, 3):
@@ -1627,12 +1643,12 @@ def test_matmul(self):
     a = self._get_rnd_float32(shape=[5, 6])
     b = self._get_rnd_float32(shape=[6, 5])
     output = run_node(node_def, [a, b])
-    np.testing.assert_almost_equal(output["Y"], np.matmul(a, b))
+    np.testing.assert_allclose(output["Y"], np.matmul(a, b), rtol=1e-6, atol=1e-6)
     # test data types that are not natively supported by Tensorflow
     a = self._get_rnd_int(0, 1000, [10, 10], np.uint32)
     b = self._get_rnd_int(0, 1000, [10, 10], np.uint32)
     output = run_node(node_def, [a, b])
-    np.testing.assert_almost_equal(output["Y"], np.matmul(a, b))
+    np.testing.assert_allclose(output["Y"], np.matmul(a, b), rtol=1e-6, atol=1e-6)
     # sys_config.auto_cast=False and a or b dtype=uint64 should throw exception
     self.assertRaises(
         RuntimeError, run_node, node_def,