Merge branch 'master' into tflfft

Author: guschmue

ci_build/azure_pipelines/templates/setup.yml

@@ -41,6 +41,11 @@ steps:
       then
         pip install tensorflow-text>=2.5
       fi
+      if [[ $CI_TF_VERSION == 2.6* ]] ;
+      then
+        # FIXME: make it >= 2.6 after offical tensorflow-text was released
+        pip install tensorflow-text==2.6.0rc0
+      fi
     fi
 
     python setup.py install

ci_build/azure_pipelines/unit_test.yml

@@ -3,6 +3,17 @@
 stages:
   - stage:
     jobs:
+    - template: 'templates/job_generator.yml'
+      parameters:
+        # tf 2.5
+        python_versions: ['3.8']
+        tf_versions: ['2.6.0rc2']
+        onnx_opsets: ['14']
+        job:
+          steps:
+          - template: 'unit_test.yml'
+        report_coverage: 'True'
+
     - template: 'templates/job_generator.yml'
       parameters:
         # TFJS tf 2.5

tests/keras2onnx_applications/nightly_build/test_resnext.py

@@ -7,7 +7,12 @@
 import numpy as np
 from mock_keras2onnx.proto import keras
 from mock_keras2onnx.proto.tfcompat import is_tf2
-from keras.regularizers import l2
+if is_tf2:
+    def l2(weight_decay):
+        # old keras layer expects a tuple but tf keras wants a single value
+        return keras.regularizers.l2(weight_decay[0])
+else:
+    from keras.regularizers import l2
 from os.path import dirname, abspath
 sys.path.insert(0, os.path.join(dirname(abspath(__file__)), '../../keras2onnx_tests/'))
 img_path = os.path.join(os.path.dirname(__file__), '../data', 'street.jpg')

tests/test_backend.py

@@ -2093,6 +2093,21 @@ def func(shape):
         self.assertTrue(-0.1 < np.mean(results) < 0.1)
         self.assertTrue(0.9 < np.std(results) < 1.1)
 
+    @check_opset_min_version(10, "TopK")
+    def test_random_shuffle(self):
+        x_val = make_xval([5, 4, 3])
+        def func(x):
+            x_ = tf.random.shuffle(x)
+            return tf.identity(x_, name=_TFOUTPUT)
+        # since results are random, compare the shapes only
+        g = self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, check_value=False, check_shape=True)
+        feed_dict = {_INPUT: x_val}
+        if "input" in g.input_names:
+            # TFLite inputs don't have port numbers
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
+        results = self.run_backend(g, g.outputs, feed_dict)
+        np.testing.assert_allclose(x_val, np.sort(results[0], axis=0))
+
     def test_randomuniform_int(self):
         def func():
             shape = tf.constant([100, 3], name="shape")
@@ -3280,6 +3295,16 @@ def func(x):
             return tf.identity(picks, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(9, "IsNaN")
+    def test_where_isnan(self):
+        x_val = np.array([1, 2, -3, float('nan'), -5, -6, float('nan'), 8, 9, 0], dtype=np.float32)
+        true_result = np.array([111, 222, 333, 444, 555, 666, 777, 888, 999, 1000],
+                               dtype=np.float32)
+        def func(x):
+            picks = tf.where(is_nan(x), true_result, x)
+            return tf.identity(picks, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
     @check_opset_min_version(9, "Where for strings needs opset 9")
     @skip_tfjs("Technically tf where doesn't support strings and tfjs doesn't like it")
     def test_where_string(self):

tf2onnx/onnx_opset/controlflow.py

@@ -184,6 +184,15 @@ def version_9(cls, ctx, node, **kwargs):
         # We can't use the mul/add trick if a NaN is involved. handles_nan is added earlier in the converter.
         handles_nan = node.get_attr_value("handles_nan", False)
         if ctx.get_dtype(node.output[0]) in [TensorProto.FLOAT, TensorProto.DOUBLE]:
+            cond_node = node.inputs[0]
+            if cond_node.type == "IsNaN":
+                handles_nan = True
+            if cond_node.type == "NotEqual" and cond_node.input[0] == cond_node.input[1]:
+                handles_nan = True
+            if cond_node.type == "Not" and cond_node.inputs[0].type == "Equal":
+                eq_node = cond_node.inputs[0]
+                if eq_node.input[0] == eq_node.input[1]:
+                    handles_nan = True
             for inp in node.inputs[1:]:
                 if inp.is_const() and np.any(np.isnan(inp.get_tensor_value(as_list=False))):
                     handles_nan = True

tf2onnx/onnx_opset/generator.py

@@ -109,6 +109,28 @@ class PassThroughOp:
     def version_1(cls, ctx, node, **kwargs):
         pass
 
+@tf_op(["RandomShuffle"])
+class RandomShuffleOp:
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        inp_shape = ctx.make_node("Shape", [node.input[0]]).output[0]
+        dim_0 = GraphBuilder(ctx).make_slice({'data': inp_shape, 'starts': [0], 'ends': [1], 'axes': [0]})
+        zeros = ctx.make_node("ConstantOfShape", [dim_0], shapes=[[-1]]).output[0]
+
+        seed = node.get_attr_value("seed", 0)
+        seed2 = node.get_attr_value("seed2", 0)
+        onnx_seed = utils.combine_seeds(seed, seed2)
+        rand_attr = {'dtype': onnx_pb.TensorProto.FLOAT}
+        if onnx_seed is not None:
+            rand_attr['seed'] = onnx_seed
+
+        random_floats = ctx.make_node("RandomUniformLike", [zeros], op_name_scope=node.name, shapes=[[-1]],
+                                      attr=rand_attr).output[0]
+        # Use indices of the TopK to get a random ordering
+        _, random_ordering = ctx.make_node("TopK", [random_floats, dim_0], output_count=2, attr={'axis': -1}).output
+        shuffled_res = ctx.make_node("Gather", [node.input[0], random_ordering]).output[0]
+        ctx.replace_all_inputs(node.output[0], shuffled_res)
+
 @tf_op("Fill")
 class Fill:
     @classmethod

tf2onnx/onnx_opset/nn.py

@@ -990,11 +990,10 @@ def version_9(cls, ctx, node, **kwargs):
 
         seed = node.get_attr_value("seed", 0)
         seed2 = node.get_attr_value("seed2", 0)
+        onnx_seed = utils.combine_seeds(seed, seed2)
         rand_attr = {}
-        if seed != 0 or seed2 != 0:
-            # Produce a unique value depending on both seeds. (diagonal grid traversal)
-            combined_seed = (seed + seed2 + 1) * (seed + seed2 + 2) // 2 - seed
-            rand_attr['seed'] = float(combined_seed)
+        if onnx_seed is not None:
+            rand_attr['seed'] = onnx_seed
 
         min_aspect_ratio, max_aspect_ratio = node.get_attr_value("aspect_ratio_range", [0.75, 1.33])
         ratio_range = max_aspect_ratio - min_aspect_ratio

tf2onnx/tf_loader.py

@@ -161,7 +161,18 @@ def make_tensor_proto_wrapped(values, dtype=None, shape=None, verify_shape=False
     tensor_util.make_tensor_proto = make_tensor_proto_wrapped
 
     try:
-        converter_data = _FunctionConverterData(func=func, lower_control_flow=False, aggressive_inlining=True)
+        function_converter = _FunctionConverterData
+        if LooseVersion(tf.__version__) >= "2.6.0":
+            from tensorflow.python.eager import context
+            from tensorflow.python.framework.convert_to_constants import _FunctionConverterDataInEager, \
+                _FunctionConverterDataInGraph
+            if context.executing_eagerly():
+                function_converter = _FunctionConverterDataInEager
+            else:
+                function_converter = _FunctionConverterDataInGraph
+        else:
+            function_converter = _FunctionConverterData
+        converter_data = function_converter(func=func, lower_control_flow=False, aggressive_inlining=True)
         frozen_graph_def, _ = _replace_variables_by_constants(converter_data=converter_data)
     finally:
         tensor_util.make_tensor_proto = make_tensor_proto_original

tf2onnx/utils.py

@@ -259,6 +259,13 @@ def make_sure(bool_val, error_msg, *args):
     if not bool_val:
         raise ValueError("make_sure failure: " + error_msg % args)
 
+def combine_seeds(seed, seed2):
+    """Produces an onnx float seed from two tf int seeds. Returns None if both seeds are 0."""
+    if seed != 0 or seed2 != 0:
+        # Produce a unique value depending on both seeds. (diagonal grid traversal)
+        combined_seed = (seed + seed2 + 1) * (seed + seed2 + 2) // 2 - seed
+        return float(combined_seed)
+    return None
 
 def topological_sort(dependencies):
     """