Implement DivNoNan conversion (#1499)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/backend_test_base.py

@@ -9,7 +9,7 @@
 from __future__ import unicode_literals
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test,import-outside-toplevel
-# pylint: disable=wrong-import-position
+# pylint: disable=wrong-import-position,invalid-unary-operand-type
 
 import logging
 import os
@@ -106,7 +106,8 @@ def run_backend(self, g, outputs, input_dict, large_model=False, postfix=""):
             raise ValueError("unknown backend")
         return y
 
-    def assert_results_equal(self, expected, actual, rtol, atol, check_value=True, check_shape=True, check_dtype=True):
+    def assert_results_equal(self, expected, actual, rtol, atol, mtol=None,
+                             check_value=True, check_shape=True, check_dtype=True):
         for expected_val, actual_val in zip(expected, actual):
             if check_value:
                 if expected_val.dtype == np.object:
@@ -115,6 +116,11 @@ def assert_results_equal(self, expected, actual, rtol, atol, check_value=True, c
                     expected_val_str = decode(expected_val)
                     self.assertAllEqual(expected_val_str, actual_val)
                 else:
+                    if mtol is not None:
+                        expected_val = np.minimum(expected_val, mtol)
+                        expected_val = np.maximum(expected_val, -mtol)
+                        actual_val = np.minimum(actual_val, mtol)
+                        actual_val = np.maximum(actual_val, -mtol)
                     self.assertAllClose(expected_val, actual_val, rtol=rtol, atol=atol)
             if check_dtype:
                 self.assertEqual(expected_val.dtype, actual_val.dtype)
@@ -285,10 +291,10 @@ def get_shape(info):
                 self.assertEqual(onnx_shape, tf2onnx_shape)
             self.assertEqual(info.type.tensor_type.elem_type, graph.get_dtype(info.name))
 
-    def run_test_case(self, func, feed_dict, input_names_with_port, output_names_with_port, rtol=1e-07, atol=1e-5,
-                      convert_var_to_const=True, constant_fold=True, check_value=True, check_shape=True,
-                      check_dtype=True, process_args=None, onnx_feed_dict=None, graph_validator=None, as_session=False,
-                      large_model=False, premade_placeholders=False):
+    def run_test_case(self, func, feed_dict, input_names_with_port, output_names_with_port,
+                      rtol=1e-07, atol=1e-5, mtol=None, convert_var_to_const=True, constant_fold=True,
+                      check_value=True, check_shape=True, check_dtype=True, process_args=None, onnx_feed_dict=None,
+                      graph_validator=None, as_session=False, large_model=False, premade_placeholders=False):
         test_tf = not self.config.skip_tf_tests
         test_tflite = not self.config.skip_tflite_tests
         run_tfl_consistency_test = test_tf and test_tflite and self.config.run_tfl_consistency_test
@@ -330,19 +336,19 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
                 g = optimizer.optimize_graph(g, catch_errors=False)
                 actual = self.run_backend(g, output_names_with_port, onnx_feed_dict, large_model)
 
-            self.assert_results_equal(expected, actual, rtol, atol, check_value, check_shape, check_dtype)
+            self.assert_results_equal(expected, actual, rtol, atol, mtol, check_value, check_shape, check_dtype)
             self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker)
 
             if graph_validator:
                 self.assertTrue(graph_validator(g))
 
         if test_tflite:
-            tfl_results, tfl_outputs = self.run_tflite(tflite_path, feed_dict)
-            test_tflite = tfl_results is not None
+            tfl_res, tfl_outputs = self.run_tflite(tflite_path, feed_dict)
+            test_tflite = tfl_res is not None
 
         if test_tflite:
             if run_tfl_consistency_test:
-                self.assert_results_equal(expected, tfl_results, rtol, atol, check_value, check_shape, check_dtype)
+                self.assert_results_equal(expected, tfl_res, rtol, atol, mtol, check_value, check_shape, check_dtype)
 
             tfl_process_args = process_args.copy()
             if 'inputs_as_nchw' in tfl_process_args:
@@ -358,9 +364,9 @@ def run_test_case(self, func, feed_dict, input_names_with_port, output_names_wit
                                  **tfl_process_args)
             g = optimizer.optimize_graph(g)
             onnx_feed_dict_without_port = {k.split(':')[0]: v for k, v in onnx_feed_dict.items()}
-            onnx_from_tfl_res = self.run_backend(g, tfl_outputs, onnx_feed_dict_without_port, postfix="_from_tflite")
+            onnx_tfl_res = self.run_backend(g, tfl_outputs, onnx_feed_dict_without_port, postfix="_from_tflite")
 
-            self.assert_results_equal(tfl_results, onnx_from_tfl_res, rtol, atol, check_value, check_shape, check_dtype)
+            self.assert_results_equal(tfl_res, onnx_tfl_res, rtol, atol, mtol, check_value, check_shape, check_dtype)
             self.assert_shapes_correct(g, self.config.allow_missing_shapes, not self.config.skip_onnx_checker)
 
             if graph_validator:

tests/test_backend.py

@@ -983,6 +983,17 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_tf_min_version("1.14")
+    @check_opset_min_version(11, "float equality")
+    def test_div_no_nan(self):
+        x_val = np.array([1.0, 2.0, -3.0, -4.0, 5.0, 0.0, float("nan"), float("-inf"), float("inf")], dtype=np.float32)
+        y_val = np.array([1.0, 0.5, 0.0, -4.0, 0.0, 0.0, 0.0, 2.0, 0.0], dtype=np.float32)
+        def func(x, y):
+            x_ = tf.math.divide_no_nan(x, y)
+            return tf.identity(x_, name=_TFOUTPUT)
+        # TFLite expresses infinity as a value > 1e38
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val}, mtol=1e38)
+
     @check_onnxruntime_incompatibility("Exp")
     def test_exp(self):
         x_val = np.array([1.0, 2.0, -3.0, -4.0], dtype=np.float32).reshape((2, 2))

tf2onnx/onnx_opset/math.py

@@ -321,6 +321,18 @@ def version_7(cls, ctx, node, **kwargs):
         pass
 
 
+@tf_op("DivNoNan")
+class DivNoNan:
+    @classmethod
+    def version_9(cls, ctx, node, **kwargs):
+        node.type = "Div"
+        np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[1]))
+        zero_const = ctx.make_const(utils.make_name("const_zero"), np.array(0, np_dtype)).output[0]
+        is_zero = ctx.make_node("Equal", [node.input[1], zero_const]).output[0]
+        where_node = ctx.make_node("Where", [is_zero, zero_const, node.output[0]])
+        ctx.insert_node_on_output(where_node, node.output[0])
+
+
 @tf_op("LRN")
 class LRN:
     @classmethod

tf2onnx/tflite_rewriters/__init__.py

@@ -4,10 +4,10 @@
 
 from tf2onnx.tflite_rewriters.tfl_scan_output_rewriter import rewrite_tfl_scan_outputs
 from tf2onnx.tflite_rewriters.tfl_qdq_rewriter import rewrite_tfl_qdq
-from tf2onnx.tflite_rewriters.tfl_select_zero_mul_rewriter import rewrite_tfl_select_zero_mul
+from tf2onnx.tflite_rewriters.tfl_select_zero_rewriter import rewrite_tfl_select_zero
 
 __all__ = [
     "rewrite_tfl_scan_outputs",
     "rewrite_tfl_qdq",
-    "rewrite_tfl_select_zero_mul",
+    "rewrite_tfl_select_zero",
 ]

tf2onnx/tflite_rewriters/tfl_select_zero_mul_rewriter.py renamed to tf2onnx/tflite_rewriters/tfl_select_zero_rewriter.py

@@ -2,22 +2,22 @@
 
 
 """
-tf2onnx.tflite_rewriters.tfl_select_zero_mul_rewriter - TFLite has a pattern to remove NaN when multiplying by 0
+tf2onnx.tflite_rewriters.tfl_select_zero_rewriter - TFLite has a pattern to remove NaN when multiplying/dividing by 0
 """
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 
 
 # pylint: disable=missing-docstring,unused-argument
 
-def rewrite_tfl_select_zero_mul(g, ops):
+def rewrite_tfl_select_zero(g, ops):
     pattern0 = \
         OpTypePattern('TFL_SELECT_V2', name='select', inputs=[
             OpTypePattern('TFL_EQUAL', name='equal', inputs=[
                 OpTypePattern('Const|ConstV2', name='const_eq'),
                 OpTypePattern('*', name='term_eq'),
             ], allow_reorder=True),
             OpTypePattern('Const|ConstV2', name='const_select'),
-            OpTypePattern('TFL_MUL', name='mul', inputs=[
+            OpTypePattern('TFL_MUL|TFL_DIV', name='mul', inputs=[
                 OpTypePattern('*', name='term_mul1'),
                 OpTypePattern('*', name='term_mul2'),
             ]),

tf2onnx/tfonnx.py

@@ -544,7 +544,7 @@ def process_parsed_graph(g, custom_op_handlers, inputs_as_nchw, continue_on_erro
         if dequantize:
             tfl_rewriters.append(rewrite_tfl_qdq)
         tfl_rewriters.append(rewrite_tfl_scan_outputs)
-        tfl_rewriters.append(rewrite_tfl_select_zero_mul)
+        tfl_rewriters.append(rewrite_tfl_select_zero)
         run_rewriters(g, tfl_rewriters, continue_on_error)
         tfl_ops_mapping = handler.tfl_op.create_tfl_to_tf_mapping()
         _, _, exceptions = tensorflow_onnx_mapping(g, tfl_ops_mapping, is_tflite=True, dequantize=False)