Repository: rwth-i6/returnn

Issue:
RandomLayer behaves different when optimized out of loop
This test fails:

```python
from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
time_dim = SpatialDim('time')
input_dim = FeatureDim('input', 3)
config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, input_dim)}}})
random_state_dim = FeatureDim('random-state', 3)
net_dict = {
  'output': {
    'class': 'rec',
    'from': "data",
    # "optimize_move_layers_out": False,
    'unit': {
      'rnd': {
        'class': 'subnetwork',
        'from': [],
        'subnetwork': {
          'random': {
            'class': 'random',
            'shape': [
              batch_dim,
              input_dim
            ],
            'distribution': 'normal',
            'mean': 0.0,
            'stddev': 1.0,
            'explicit_state': 'base:base:rnd_state_var0',
            'auto_update_state': True
          },
          'output': {
            'class': 'copy',
            'from': 'random',
            'out_shape': {batch_dim, input_dim}
          }
        },
        'out_shape': {batch_dim, input_dim}
      },
      'output': {
        'class': 'eval',
        'from': ['data:source', 'rnd'],
        'eval': 'source(0) * 0.0 + source(1)',
        'out_shape': {batch_dim, input_dim}
      },
    },
    'axis': time_dim,
    'out_shape': {batch_dim, time_dim, input_dim},
    'name_scope': ''
  },
  'random_state_init': {
    'class': 'random_state_init',
    'out_dim': random_state_dim,
    'out_shape': {random_state_dim}
  },
  'rnd_state_var0': {
    'class': 'variable',
    'shape': [
      random_state_dim
    ],
    'param_name': 'param',
    'dtype': 'int64',
    'init_by_layer': 'random_state_init',
    'name_scope': 'rnd/state_var0'
  }
}
with make_scope() as session:
  net = TFNetwork(config=config)
  net.construct_from_dict(net_dict)
  net.initialize_params(session)
  out = net.get_default_output_layer().output.copy_as_time_major()
  out_np = session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
  print(out_np)
  out0_np = out_np[:1]  # [1,B,D]
  print((out0_np == out_np))
  assert not (out0_np == out_np).all()  # not all the same
```

When you put in the `optimize_move_layers_out=False`, it works.
This violates the [RETURNN principle](https://github.com/rwth-i6/returnn/wiki/RETURNN-principles) that the optimization should not have an effect on the outcome.

This is also reported on returnn-common side: https://github.com/rwth-i6/returnn_common/issues/184



Discussion:

Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -2540,6 +2540,21 @@ def transform_config_dict(cls, d, network, get_layer):
     for attrib in ("mean", "stddev", "bound", "minval", "maxval", "explicit_state"):
       if attrib in d and isinstance(d[attrib], str):
         d[attrib] = get_layer(d[attrib])
+    # We need special care in case this is inside a RecLayer.
+    # https://github.com/rwth-i6/returnn/issues/1044
+    # When this layer stays inside the loop, all is fine.
+    # When it is moved out, we must add the loop dim
+    # such that we get different random values in each loop iteration.
+    # When the loop dim is not defined yet because of a dynamic loop ("end" layer),
+    # this is not possible, so we must make sure the layer stays inside the loop.
+    inside_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=True)
+    over_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=False)
+    if over_rec_time_dim:
+      collocate_with = list(d.get("collocate_with", None) or [])
+      collocate_with.append("end")  # in case the "end" layer is used, it will be collocated
+      d["collocate_with"] = collocate_with
+    if over_rec_time_dim and not inside_rec_time_dim:  # moved out of loop
+      d["shape"] = [over_rec_time_dim] + list(d["shape"])
 
   @classmethod
   def get_out_data_from_opts(cls, name, shape, dtype="float32", **kwargs):


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -4862,6 +4862,81 @@ def test_rand_indices():
     assert_equal(output.shape, (n_batch, n_time, sz[-1].dimension, feature_dim.dimension))
 
 
+def test_RandomLayer_in_loop():
+  # https://github.com/rwth-i6/returnn/issues/1044
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+  time_dim = SpatialDim('time')
+  input_dim = FeatureDim('input', 3)
+  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, input_dim)}}})
+  random_state_dim = FeatureDim('random-state', 3)
+  net_dict = {
+    'output': {
+      'class': 'rec',
+      'from': "data",
+      'unit': {
+        'rnd': {
+          'class': 'subnetwork',
+          'from': [],
+          'subnetwork': {
+            'random': {
+              'class': 'random',
+              'shape': [
+                batch_dim,
+                input_dim
+              ],
+              'distribution': 'normal',
+              'mean': 0.0,
+              'stddev': 1.0,
+              'explicit_state': 'base:base:rnd_state_var0',
+              'auto_update_state': True
+            },
+            'output': {
+              'class': 'copy',
+              'from': 'random',
+              'out_shape': {batch_dim, input_dim}
+            }
+          },
+          'out_shape': {batch_dim, input_dim}
+        },
+        'output': {
+          'class': 'eval',
+          'from': ['data:source', 'rnd'],
+          'eval': 'source(0) * 0.0 + source(1)',
+          'out_shape': {batch_dim, input_dim}
+        },
+      },
+      'axis': time_dim,
+      'out_shape': {batch_dim, time_dim, input_dim},
+      'name_scope': ''
+    },
+    'random_state_init': {
+      'class': 'random_state_init',
+      'out_dim': random_state_dim,
+      'out_shape': {random_state_dim}
+    },
+    'rnd_state_var0': {
+      'class': 'variable',
+      'shape': [
+        random_state_dim
+      ],
+      'param_name': 'param',
+      'dtype': 'int64',
+      'init_by_layer': 'random_state_init',
+      'name_scope': 'rnd/state_var0'
+    }
+  }
+  with make_scope() as session:
+    net = TFNetwork(config=config)
+    net.construct_from_dict(net_dict)
+    net.initialize_params(session)
+    out = net.get_default_output_layer().output.copy_as_time_major()
+    out_np = session.run(out.placeholder, feed_dict=make_feed_dict(net.extern_data))
+    print(out_np)
+    out0_np = out_np[:1]  # [1,B,D]
+    print((out0_np == out_np))
+    assert not (out0_np == out_np).all()  # not all the same
+
+
 def test_untrainable_params():
   with make_scope() as session:
     config = Config()