Implement RandomShuffle op (#1658)

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

Author: TomWildenhain-Microsoft

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")

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/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):
     """