Merge branch 'master' into tom/MatrixBandPart

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")
@@ -2981,7 +3026,7 @@ def func(input_x):
 
     @check_opset_min_version(11, "CumSum")
     def test_matrix_band_part_3(self):
-        for low, high in [(-1, 3), (2, 3), (4, 3), (0, -1), (0, 0)]:
+        for low, high in [(-1, 3), (2, 3), (4, 3), (0, -1), (0, 0), (-1, -1)]:
             input_val = np.random.randint(0, 666, (10, 15)).astype(np.int32)
             def func(input_x):
                 res = tf.linalg.band_part(input_x, low, high)

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/onnx_opset/nn.py

@@ -1188,7 +1188,7 @@ def version_11(cls, ctx, node, **kwargs):
 @tf_op("MatrixBandPart")
 class MatrixBandPart:
     @classmethod
-    def version_7(cls, opset, ctx, node, **kwargs):
+    def version_7(cls, ctx, node, **kwargs):
         # T output = MatrixBandPart(T input, int num_lower, int num_upper)
         # data-flow: first generate mask matrix and then use element-wise mul op
         input_rank = len(ctx.get_shape(node.input[0]))
@@ -1274,7 +1274,6 @@ def version_11(cls, ctx, node, **kwargs):
         rank = ctx.get_rank(data)
         int_max_val = utils.get_max_value(np.int64)
         dtype = ctx.get_dtype(data)
-        np_dtype = utils.map_onnx_to_numpy_type(dtype)
         if rank == 2:
             shape = ctx.make_node("Shape", [data]).output[0]
         else:
@@ -1288,8 +1287,6 @@ def version_11(cls, ctx, node, **kwargs):
                 zero_tensor = helper.make_tensor("value", dtype, dims=[1], vals=[0])
                 const_of_shape = ctx.make_node("ConstantOfShape", [shape], attr={'value': zero_tensor}).output[0]
                 identity_node = ctx.make_node("EyeLike", [const_of_shape]).output[0]
-            one_const = ctx.make_const(utils.make_name("one"), np.array(1, np_dtype)).output[0]
-            mask = ctx.make_node("Sub", [one_const, identity_node]).output[0]
             shapes = node.output_shapes
             dtypes = node.output_dtypes
             ctx.remove_node(node.name)
@@ -1316,12 +1313,16 @@ def version_11(cls, ctx, node, **kwargs):
         if num_upper_const is None or num_upper_const >= 0:
             if ctx.get_dtype(num_upper) != TensorProto.INT64:
                 num_upper = ctx.make_node("Cast", [num_upper], attr={'to': TensorProto.INT64}).output[0]
-            conditions.append(ctx.make_node("LessOrEqual", [idx_diff, num_upper]).output[0])
+            greater = ctx.make_node("Greater", [idx_diff, num_upper]).output[0]
+            less_or_equal = ctx.make_node("Not", [greater]).output[0]
+            conditions.append(less_or_equal)
         if num_lower_const is None or num_lower_const >= 0:
             if ctx.get_dtype(num_lower) != TensorProto.INT64:
                 num_lower = ctx.make_node("Cast", [num_lower], attr={'to': TensorProto.INT64}).output[0]
             num_lower_neg = ctx.make_node("Neg", [num_lower]).output[0]
-            conditions.append(ctx.make_node("LessOrEqual", [num_lower_neg, idx_diff]).output[0])
+            greater = ctx.make_node("Greater", [num_lower_neg, idx_diff]).output[0]
+            less_or_equal = ctx.make_node("Not", [greater]).output[0]
+            conditions.append(less_or_equal)
         if len(conditions) == 0:
             node.type = "Identity"
             ctx.replace_inputs(node, [data])

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