Implement RandomStandardNormal

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

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -1962,6 +1962,45 @@ def func():
         # since results are random, compare the shapes only
         self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
 
+    def test_random_std_normal(self):
+        def func():
+            shape = tf.constant([20, 10, 5], name="shape")
+            x_ = tf.random.normal(shape)
+            return tf.identity(x_, name=_TFOUTPUT)
+        # since results are random, compare the shapes only
+        g = self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
+        results = self.run_backend(g, g.outputs, {})[0]
+        self.assertTrue(-0.1 < np.mean(results) < 0.1)
+        self.assertTrue(0.9 < np.std(results) < 1.1)
+
+    def test_randomnormal(self):
+        def func():
+            shape = tf.constant([20, 10, 5], name="shape")
+            x_ = tf.random.normal(shape, mean=10, stddev=2)
+            return tf.identity(x_, name=_TFOUTPUT)
+        # since results are random, compare the shapes only
+        g = self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
+        results = self.run_backend(g, g.outputs, {})[0]
+        if not 9.9 < np.mean(results) < 10.1:
+            np.testing.assert_allclose(np.mean(results), 0)
+        self.assertTrue(1.9 < np.std(results) < 2.1)
+
+    @check_opset_min_version(9, "RandomNormalLike")
+    def test_randomnormal_unknown_shape(self):
+        shape_val = np.array([20, 10, 5], np.int32)
+        def func(shape):
+            x_ = tf.random.normal(shape)
+            return tf.identity(x_, name=_TFOUTPUT)
+        # since results are random, compare the shapes only
+        feed_dict = {_INPUT: shape_val}
+        g = self._run_test_case(func, [_OUTPUT], feed_dict, check_value=False, check_shape=True)
+        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)[0]
+        self.assertTrue(-0.1 < np.mean(results) < 0.1)
+        self.assertTrue(0.9 < np.std(results) < 1.1)
+
     def test_randomuniform_int(self):
         def func():
             shape = tf.constant([100, 3], name="shape")

tf2onnx/onnx_opset/generator.py

@@ -29,7 +29,7 @@ def version_1(cls, ctx, node, **kwargs):
         pass
 
 
-@tf_op(["RandomNormal", "RandomUniform", "RandomUniformInt"])
+@tf_op(["RandomNormal", "RandomUniform", "RandomUniformInt", "RandomStandardNormal"])
 class RandomOp:
     @classmethod
     def randuniform_int(cls, ctx, rand_node, rand_out, min_inp, max_inp):
@@ -66,7 +66,7 @@ def version_1(cls, ctx, node, **kwargs):
         # const we make it an attribute.
         seed = node.get_attr("seed")
         node.set_attr("seed", float(seed.f))
-        utils.make_sure(node.inputs[0].is_const(), "%s node with non-const shape requires opset >= 9")
+        utils.make_sure(node.inputs[0].is_const(), "%s node with non-const shape requires opset >= 9", node.type)
         shape = node.inputs[0].get_tensor_value()
         ctx.remove_input(node, node.input[0], 0)
         if len(shape) == 0:
@@ -84,6 +84,8 @@ def version_1(cls, ctx, node, **kwargs):
             cls.randuniform_int(ctx, node, rand_out, node.input[0], node.input[1])
             node.type = "RandomUniform"
             ctx.replace_inputs(node, [])
+        elif node.type == "RandomStandardNormal":
+            node.type = "RandomNormal"
 
     @classmethod
     def version_9(cls, ctx, node, **kwargs):
@@ -99,6 +101,8 @@ def version_9(cls, ctx, node, **kwargs):
             if node.type == "RandomUniformInt":
                 cls.randuniform_int(ctx, node, node.output[0], inputs[1], inputs[2])
                 node.type = "RandomUniformLike"
+            elif node.type == "RandomStandardNormal":
+                node.type = "RandomNormalLike"
             else:
                 node.type = node.type + 'Like'
 

tf2onnx/rewriter/random_normal_rewriter.py

@@ -39,23 +39,29 @@ def rewrite_random_normal(g, ops):
             else:
                 # pattern 2
                 mean = 0.0
+            input2 = match.get_op('input2')
+            if input2.type == 'Mul':
+                scale = input2.inputs[1].get_tensor_value()
+            else:
+                scale = 1.0
             dtype = g.get_dtype(output.output[0])
             op_name = utils.make_name("RandomNormal")
             out_name = utils.port_name(op_name)
 
             rn_op = match.get_op('input1')
             seed = float(rn_op.get_attr('seed2').i)
 
+            attr = {"mean": mean, "scale": scale, "dtype": dtype, "seed": seed}
             if rn_op.inputs[0].type == "Shape":
                 shape_node = rn_op.inputs[0]
                 new_node = g.make_node("RandomNormalLike", [shape_node.input[0]], outputs=[out_name], name=op_name,
-                                       attr={"mean": mean, "scale": 1.0, "dtype": dtype, "seed": seed})
+                                       attr=attr)
             else:
                 shape = g.get_shape(output.output[0])
                 if shape is None or -1 in shape:
                     continue
-                new_node = g.make_node("RandomNormal", [], outputs=[out_name], name=op_name,
-                                       attr={"shape": shape, "mean": mean, "scale": 1.0, "dtype": dtype, "seed": seed})
+                attr['shape'] = shape
+                new_node = g.make_node("RandomNormal", [], outputs=[out_name], name=op_name, attr=attr)
 
             g.replace_all_inputs(output.output[0], new_node.output[0], ops=ops)
             g.safe_remove_nodes(match.get_nodes())