Convert RandomUniformInt (#1347)

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

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1800,16 +1800,61 @@ def func():
         # since results are random, compare the shapes only
         self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
 
-    @unittest.skip("TF RandomUniformInt is not supported")
     def test_randomuniform_int(self):
         def func():
-            shape = tf.constant([2, 3], name="shape")
-            x_ = random_uniform(shape, name="rand", dtype=tf.int32, maxval=10)
+            shape = tf.constant([100, 3], name="shape")
+            x_ = random_uniform(shape, name="rand", dtype=tf.int32, minval=2, maxval=10)
             x_ = tf.identity(x_, name="output1")
             x_ = tf.identity(x_, name="output2")
             return tf.identity(x_, name=_TFOUTPUT)
         # since results are random, compare the shapes only
-        self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
+        g = self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
+        results = self.run_backend(g, [_OUTPUT], {})
+        numbers = set(results[0].flatten())
+        self.assertEqual(sorted(numbers), list(range(2, 10)))
+
+    def test_randomuniform_int_nonconst_max(self):
+        m_val = np.array(8, dtype=np.int32)
+        def func(m):
+            shape = tf.constant([100, 3], name="shape")
+            x_ = random_uniform(shape, name="rand", dtype=tf.int32, minval=0, maxval=m)
+            x_ = tf.identity(x_, name="output1")
+            x_ = tf.identity(x_, name="output2")
+            return tf.identity(x_, name=_TFOUTPUT)
+        g = self._run_test_case(func, [_OUTPUT], {_INPUT: m_val}, check_value=False, check_shape=True)
+        results = self.run_backend(g, [_OUTPUT], {_INPUT: m_val})
+        numbers = set(results[0].flatten())
+        self.assertEqual(sorted(numbers), list(range(8)))
+
+    def test_randomuniform_int_nonconst_min_max(self):
+        n_val = np.array(2, dtype=np.int32)
+        m_val = np.array(10, dtype=np.int32)
+        def func(n, m):
+            shape = tf.constant([100, 3], name="shape")
+            x_ = random_uniform(shape, name="rand", dtype=tf.int32, minval=n, maxval=m)
+            x_ = tf.identity(x_, name="output1")
+            x_ = tf.identity(x_, name="output2")
+            return tf.identity(x_, name=_TFOUTPUT)
+        g = self._run_test_case(func, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val}, check_value=False, check_shape=True)
+        results = self.run_backend(g, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val})
+        numbers = set(results[0].flatten())
+        self.assertEqual(sorted(numbers), list(range(2, 10)))
+
+    @check_opset_min_version(9, "RandomUniformLike")
+    def test_randomuniform_int_nonconst_min_max_shape(self):
+        n_val = np.array(2, dtype=np.int32)
+        m_val = np.array(10, dtype=np.int32)
+        s_val = np.array([100, 3], dtype=np.int64)
+        def func(n, m, s):
+            x_ = random_uniform(s, name="rand", dtype=tf.int32, minval=n, maxval=m)
+            x_ = tf.identity(x_, name="output1")
+            x_ = tf.identity(x_, name="output2")
+            return tf.identity(x_, name=_TFOUTPUT)
+        g = self._run_test_case(func, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val, _INPUT2: s_val},
+                                check_value=False, check_shape=True)
+        results = self.run_backend(g, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val, _INPUT2: s_val})
+        numbers = set(results[0].flatten())
+        self.assertEqual(sorted(numbers), list(range(2, 10)))
 
     @skip_caffe2_backend()
     @check_opset_after_tf_version("2.2", 9, "RandomUniform")

tf2onnx/onnx_opset/generator.py

@@ -28,8 +28,33 @@ def version_1(cls, ctx, node, **kwargs):
         pass
 
 
-@tf_op(["RandomNormal", "RandomUniform"])
+@tf_op(["RandomNormal", "RandomUniform", "RandomUniformInt"])
 class RandomOp:
+    @classmethod
+    def randuniform_int(cls, ctx, node, min_inp, max_inp):
+        dtype = ctx.get_dtype(node.output[0])
+        min_node = ctx.get_node_by_output(min_inp)
+        max_node = ctx.get_node_by_output(max_inp)
+        if min_node.is_const() and max_node.is_const():
+            node.set_attr('low', float(min_node.get_tensor_value()))
+            node.set_attr('high', float(max_node.get_tensor_value()))
+            out = node.output[0]
+        elif min_node.is_const() and min_node.get_tensor_value() == 0:
+            max_float = ctx.make_node("Cast", [max_inp], attr={'to': onnx_pb.TensorProto.FLOAT}).output[0]
+            mul_node = ctx.insert_new_node_on_output("Mul", node.output[0], inputs=[node.output[0], max_float])
+            out = mul_node.output[0]
+        else:
+            min_float = ctx.make_node("Cast", [min_inp], attr={'to': onnx_pb.TensorProto.FLOAT}).output[0]
+            max_float = ctx.make_node("Cast", [max_inp], attr={'to': onnx_pb.TensorProto.FLOAT}).output[0]
+            diff = ctx.make_node("Sub", [max_float, min_float]).output[0]
+            diff_float = ctx.make_node("Cast", [diff], attr={'to': onnx_pb.TensorProto.FLOAT}).output[0]
+            mul_node = ctx.insert_new_node_on_output("Mul", node.output[0], inputs=[node.output[0], diff_float])
+            mul = mul_node.output[0]
+            add_node = ctx.insert_new_node_on_output("Add", mul, inputs=[mul, min_float])
+            out = add_node.output[0]
+        floor_node = ctx.insert_new_node_on_output("Floor", out)
+        ctx.insert_new_node_on_output("Cast", floor_node.output[0], to=dtype)
+
     @classmethod
     def version_1(cls, ctx, node, **kwargs):
         # in tf-2.0 grappler optimizes the graph pretty well and our matching logic
@@ -43,6 +68,10 @@ def version_1(cls, ctx, node, **kwargs):
             ctx.remove_input(node, node.input[0], 0)
             node.set_attr("shape", shape)
             ctx.set_shape(node.output[0], shape)
+        if node.type == "RandomUniformInt":
+            cls.randuniform_int(ctx, node, node.input[0], node.input[1])
+            node.type = "RandomUniform"
+            ctx.replace_inputs(node, [])
 
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
@@ -51,10 +80,15 @@ def version_9(cls, ctx, node, **kwargs):
         else:
             seed = node.get_attr("seed")
             node.set_attr("seed", float(seed.f))
-            cast_node = ctx.make_node("Cast", node.input, attr={'to': onnx_pb.TensorProto.INT64})
+            cast_node = ctx.make_node("Cast", [node.input[0]], attr={'to': onnx_pb.TensorProto.INT64})
             const_node = ctx.make_node("ConstantOfShape", cast_node.output)
+            inputs = node.input.copy()
             ctx.replace_inputs(node, const_node.output.copy())
-            node.type = node.type + 'Like'
+            if node.type == "RandomUniformInt":
+                cls.randuniform_int(ctx, node, inputs[1], inputs[2])
+                node.type = "RandomUniformLike"
+            else:
+                node.type = node.type + 'Like'
 
 
 @tf_op(["RandomNormalLike", "RandomUniformLike"])

tf2onnx/tf_utils.py

@@ -221,6 +221,7 @@ def is_huge_shape(x):
                     outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
                     progress = True
             can_fold = node.type not in ['Enter', 'Placeholder', 'PlaceholderWithDefault']
+            can_fold = can_fold and not node.type.startswith('Random')
             can_fold = can_fold and len(input_names) > 0 and all(inp in outputs_to_values for inp in input_names)
             # We can only fold nodes with a single output
             can_fold = can_fold and len(output_names) == 1 and output_names[0] not in outputs_to_values