Repository: rwth-i6/returnn

Issue:
Repeat layer as loss crashes
When writing the test for rwth-i6/returnn_common#200 I noticed that marking the `nn.repeat` output does not work. This does not affect me right now and maybe there is something wrong in the test, but its just something I noticed, so I wanted to report it.

```python
def test_repeat_as_loss():
  # https://github.com/rwth-i6/returnn/issues/1115
  from returnn.config import Config
  from returnn.tf.engine import Engine
  from returnn.datasets import init_dataset
  time_dim = nn.SpatialDim("time")
  in_dim = nn.FeatureDim("in", 3)
  out_dim = nn.FeatureDim("out", 5)
  x = nn.Data("data", dim_tags=[nn.batch_dim, time_dim, in_dim], available_for_inference=True)


  def _config_get_network(epoch: int, **_kwargs) -> dict:
    # noinspection PyStatementEffect
    epoch  # unused
    nn.reset_default_root_name_ctx()
    out, dim = nn.repeat(nn.get_extern_data(x), repetitions=1, axis=time_dim, out_dim=time_dim)
    out.mark_as_default_output()
    out.mark_as_loss()
    net_dict = nn.get_returnn_config().get_net_dict_raw_dict(nn.Module())
    return net_dict

  config = Config({
    "task": "train", "num_epochs": 1, "start_epoch": 1,
    "get_network": _config_get_network,
    "extern_data": {x.name: {"dim_tags": [nn.batch_dim, time_dim, in_dim], "available_for_inference": True}},
  })
  train_dataset = init_dataset(
    {"class": "DummyDataset", "input_dim": in_dim.dimension, "output_dim": 5, "num_seqs": 3})
  engine = Engine(config)
  engine.init_train_from_config(config, train_data=train_dataset)
  engine.train()
```

```
test_repeat_as_loss()
setup_tf_thread_pools() not yet called (via func get_tf_list_local_devices), calling it now.
Use num_threads=1 (but min 2) via OMP_NUM_THREADS.
Setup TF inter and intra global thread pools, num_threads 2, session opts {'log_device_placement': False, 'device_count': {'GPU': 0}, 'intra_op_parallelism_threads': 2, 'inter_op_parallelism_threads': 2}.
2022-09-01 10:41:59.372269: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN)to use the following CPU instructions in perfo
rmance-critical operations:  SSE4.1 SSE4.2 AVX
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-09-01 10:41:59.406917: I tensorflow/core/platform/profile_utils/cpu_utils.cc:104] CPU Frequency: 2600090000 Hz
2022-09-01 10:41:59.407581: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x581dfd0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2022-09-01 10:41:59.407692: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
2022-09-01 10:41:59.417080: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcuda.so.1
2022-09-01 10:42:00.057617: E tensorflow/stream_executor/cuda/cuda_driver.cc:314] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
2022-09-01 10:42:00.057707: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:169] retrieving CUDA diagnostic information for host: cluster-cn-211
2022-09-01 10:42:00.057725: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:176] hostname: cluster-cn-211
2022-09-01 10:42:00.057932: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:200] libcuda reported version is: 418.67.0
2022-09-01 10:42:00.058024: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:204] kernel reported version is: 418.67.0
2022-09-01 10:42:00.058040: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:310] kernel version seems to match DSO: 418.67.0
Collecting TensorFlow device list...
Device not set explicitly, and no GPU found.
Learning-rate-control: no file specified, not saving history (no proper restart possible)
2022-09-01 10:42:00.505612: W tensorflow/python/util/util.cc:348] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
Setup TF session with options {'log_device_placement': False, 'device_count': {'GPU': 0}} ...
layer /'data:data': [B,T|'time'[B],F|F'in'(3)] float32
layer /'repeat': [B,T|'time'[B],F|F'in'(3)] float32
layer /'output': [B,T|'time'[B],F|F'in'(3)] float32
Loss flattened layer /'data:data_flat' output: Data{'data:data_flat_output', [B&Packed{'time'},F|F'in'(3)]}
EXCEPTION
Traceback (most recent call last):
  File "/work/tools/asr/python/3.8.0/lib/python3.8/runpy.py", line 192, in _run_module_as_main
    line: return _run_code(code, main_globals, None,
                           "__main__", mod_spec)
    locals:
      _run_code = <global> <function _run_code at 0x7fc9d0a59670>
      code = <local> <code object <module> at 0x7fc9d0a3bbe0, file "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/test_nn_base.py", line 1>
      main_globals = <local> {'__name__': '__main__', '__doc__': '\nTest the base nn functionality\n', '__package__': 'returnn_common.tests', '__loader__': <_frozen_importlib_external.SourceFileLoader object at 
0x7fc9d0a502e0>, '__spec__': ModuleSpec(name='tests.test_nn_base', loader=<_frozen_importlib_external.SourceFileLoa..., len = 10
      mod_spec = <local> ModuleSpec(name='tests.test_nn_base', loader=<_frozen_importlib_external.SourceFileLoader object at 0x7fc9d0a502e0>, origin='/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests
/test_nn_base.py')
  File "/work/tools/asr/python/3.8.0/lib/python3.8/runpy.py", line 85, in _run_code
    line: exec(code, run_globals)
    locals:
      exec = <builtin> <built-in function exec>
      code = <local> <code object <module> at 0x7fc9d0a3bbe0, file "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/test_nn_base.py", line 1>
      run_globals = <local> {'__name__': '__main__', '__doc__': '\nTest the base nn functionality\n', '__package__': 'returnn_common.tests', '__loader__': <_frozen_importlib_external.SourceFileLoader object at 0
x7fc9d0a502e0>, '__spec__': ModuleSpec(name='tests.test_nn_base', loader=<_frozen_importlib_external.SourceFileLoa..., len = 10
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/test_nn_base.py", line 7, in <module>
    line: from . import _setup_test_env  # noqa
    locals:
      from = <not found>
      from.import = <not found>
      _setup_test_env = <not found>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/_setup_test_env.py", line 304, in <module>
    line: setup()
    locals:
      setup = <local> <function setup at 0x7fc9d0a23790>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/_setup_test_env.py", line 131, in setup
    line: _main(main_mod_)
    locals:
      _main = <global> <function _main at 0x7fc9d0a23700>
      main_mod_ = <local> <module 'returnn_common.tests.test_nn_base' from '/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/test_nn_base.py'>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/_setup_test_env.py", line 283, in _main
    line: test()
    locals:
      test = <local> <function test_repeat_as_loss at 0x7fc98fb5bca0>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn_common/tests/test_nn_base.py", line 584, in test_repeat_as_loss
    line: engine.init_train_from_config(config, train_data=train_dataset)
    locals:
      engine = <local> <returnn.tf.engine.Engine object at 0x7fc98f9d9220>
      engine.init_train_from_config = <local> <bound method Engine.init_train_from_config of <returnn.tf.engine.Engine object at 0x7fc98f9d9220>>
      config = <local> <returnn.config.Config object at 0x7fc98fd4e880>
      train_data = <not found>
      train_dataset = <local> <DummyDataset 'dataset_id140503675435040' epoch=None>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/engine.py", line 1056, in Engine.init_train_from_config
    line: self.init_network_from_config(config)
    locals:
      self = <local> <returnn.tf.engine.Engine object at 0x7fc98f9d9220>
      self.init_network_from_config = <local> <bound method Engine.init_network_from_config of <returnn.tf.engine.Engine object at 0x7fc98f9d9220>>
      config = <local> <returnn.config.Config object at 0x7fc98fd4e880>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/engine.py", line 1121, in Engine.init_network_from_config
    line: self._init_network(net_desc=net_dict, epoch=self.epoch)
    locals:
      self = <local> <returnn.tf.engine.Engine object at 0x7fc98f9d9220>
      self._init_network = <local> <bound method Engine._init_network of <returnn.tf.engine.Engine object at 0x7fc98f9d9220>>
      net_desc = <not found>
      net_dict = <local> {'repeat': {'class': 'repeat', 'from': 'data:data', 'repetitions': 1, 'axis': Dim{'time'[B]}, 'out_dim': Dim{'time'[B]}, 'loss': 'as_is', 'out_shape': {Dim{'time'[B]}, Dim{B}, Dim{F'in'(
3)}}}, 'output': {'class': 'copy', 'from': 'repeat', 'out_shape': {Dim{'time'[B]}, Dim{B}, Dim{F'in'(3)}}}}
      epoch = <local> None
      self.epoch = <local> 1
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/engine.py", line 1303, in Engine._init_network
    line: self.network, self.updater = self.create_network(
            config=self.config,
            extern_data=extern_data,
            rnd_seed=net_random_seed,
            train_flag=train_flag, eval_flag=self.use_eval_flag, search_flag=self.use_search_flag,
            initial_learning_rate=getattr(self, "initial_learning_rate", None),
            net_dict=net_desc)
    locals:
      self = <local> <returnn.tf.engine.Engine object at 0x7fc98f9d9220>
      self.network = <local> None
      self.updater = <local> None
      self.create_network = <local> <bound method Engine.create_network of <class 'returnn.tf.engine.Engine'>>
      config = <not found>
      self.config = <local> <returnn.config.Config object at 0x7fc98fd4e880>
      extern_data = <local> <ExternData data={'data': Data{'data', [B,T|'time'[B],F|F'in'(3)]}}>
      rnd_seed = <not found>
      net_random_seed = <local> 1
      train_flag = <local> <tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>
      eval_flag = <not found>
      self.use_eval_flag = <local> True
      search_flag = <not found>
      self.use_search_flag = <local> False
      initial_learning_rate = <not found>
      getattr = <builtin> <built-in function getattr>
      net_dict = <not found>
      net_desc = <local> {'repeat': {'class': 'repeat', 'from': 'data:data', 'repetitions': 1, 'axis': Dim{'time'[B]}, 'out_dim': Dim{'time'[B]}, 'loss': 'as_is', 'out_shape': {Dim{'time'[B]}, Dim{B}, Dim{F'in'(
3)}}}, 'output': {'class': 'copy', 'from': 'repeat', 'out_shape': {Dim{'time'[B]}, Dim{B}, Dim{F'in'(3)}}}}
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/engine.py", line 1353, in Engine.create_network
    line: updater = Updater(
            config=config, network=network,
            initial_learning_rate=initial_learning_rate)
    locals:
      updater = <local> None
      Updater = <global> <class 'returnn.tf.updater.Updater'>
      config = <local> <returnn.config.Config object at 0x7fc98fd4e880>
      network = <local> <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>
      initial_learning_rate = <local> 1.0
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/updater.py", line 177, in Updater.__init__
    line: self.loss = network.get_objective()
    locals:
      self = <local> <returnn.tf.updater.Updater object at 0x7fc98c0e6fd0>
      self.loss = <local> !AttributeError: 'Updater' object has no attribute 'loss'
      network = <local> <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>
      network.get_objective = <local> <bound method TFNetwork.get_objective of <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/network.py", line 1292, in TFNetwork.get_objective
    line: self.maybe_construct_objective()
    locals:
      self = <local> <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>
      self.maybe_construct_objective = <local> <bound method TFNetwork.maybe_construct_objective of <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/network.py", line 1285, in TFNetwork.maybe_construct_objective
    line: self._construct_objective()
    locals:
      self = <local> <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>
      self._construct_objective = <local> <bound method TFNetwork._construct_objective of <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/network.py", line 1267, in TFNetwork._construct_objective
    line: self._flatten_layer_with_losses()
    locals:
      self = <local> <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>
      self._flatten_layer_with_losses = <local> <bound method TFNetwork._flatten_layer_with_losses of <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/network.py", line 1584, in TFNetwork._flatten_layer_with_losses
    line: mapped_layers[layer] = _make_layer(type(layer), layer.kwargs)
    locals:
      mapped_layers = <local> {<SourceLayer 'data:data' out_type=Data{[B,T|'time'[B],F|F'in'(3)]}>: <FlattenBatchLayer 'data:data_flat' out_type=Data{[B&Packed{'time'},F|F'in'(3)]}>}
      layer = <local> <RepeatLayer 'repeat' out_type=Data{[B,T|'time'[B],F|F'in'(3)]}>
      _make_layer = <local> <function TFNetwork._flatten_layer_with_losses.<locals>._make_layer at 0x7fc98c07a310>
      type = <builtin> <class 'type'>
      layer.kwargs = <local> {'repetitions': 1, 'axis': Dim{'time'[B]}, 'out_dim': Dim{'time'[B]}, 'out_shape': {Dim{'time'[B]}, Dim{B}, Dim{F'in'(3)}}, '_network': <TFNetwork '' train=<tf.Tensor 'globals/train_
flag:0' shape=() dtype=bool>>, '_name': 'repeat', 'sources': [<SourceLayer 'data:data' out_type=Data{[B,T|'time'[B]..., len = 11
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/network.py", line 1378, in TFNetwork._flatten_layer_with_losses.<locals>._make_layer
    line: opts["output"] = layer_cls.get_out_data_from_opts(**opts)
    locals:
      opts = <local> {'repetitions': 1, 'axis': Dim{'time'[B]}, 'out_dim': Dim{'time'[B]}, 'out_shape': {Dim{'time'[B]}, Dim{B}, Dim{F'in'(3)}}, '_network': <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' 
shape=() dtype=bool>>, '_name': 'repeat', 'sources': [<FlattenBatchLayer 'data:data_flat' out_type=Data{[B&..., len = 10
      layer_cls = <local> <class 'returnn.tf.layers.basic.RepeatLayer'>
      layer_cls.get_out_data_from_opts = <local> <bound method RepeatLayer.get_out_data_from_opts of <class 'returnn.tf.layers.basic.RepeatLayer'>>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/layers/basic.py", line 4494, in RepeatLayer.get_out_data_from_opts
    line: original_axis = data.get_axis_from_description(axis, allow_int=False)
    locals:
      original_axis = <not found>
      data = <local> Data{'repeat_output', [B&Packed{'time'},F|F'in'(3)]}
      data.get_axis_from_description = <local> <bound method Data.get_axis_from_description of Data{'repeat_output', [B&Packed{'time'},F|F'in'(3)]}>
      axis = <local> Dim{'time'[B]}
      allow_int = <not found>
  File "/u/hilmes/experiments/tts_new_sis/recipe/returnn/tf/util/data.py", line 4777, in Data.get_axis_from_description
    line: assert axes, "%s: %r axis not found" % (self, axis)
    locals:
      axes = <local> []
      self = <local> Data{'repeat_output', [B&Packed{'time'},F|F'in'(3)]}
      axis = <local> Dim{'time'[B]}
AssertionError: Data{'repeat_output', [B&Packed{'time'},F|F'in'(3)]}: Dim{'time'[B]} axis not found
```


Discussion:

Resolving patch:
diff --git a/returnn/tf/network.py b/returnn/tf/network.py
--- a/returnn/tf/network.py
+++ b/returnn/tf/network.py
@@ -1457,6 +1457,11 @@ def _check_push_flattening_to_inputs_for_layer(layer_):
       deps = _layer_deps(layer_)
       if not deps:
         return False
+      layer_kwargs = layer_.kwargs.copy()
+      layer_kwargs.pop("out_shape", None)
+      layer_kwargs_flat_values = nest.flatten(layer_kwargs)
+      if any(dim in layer_kwargs_flat_values for dim in dims):  # e.g. to operate on the axis
+        return False
       valid_deps = all(
         set(dep_.output.dim_tags).issuperset(dims)
         or set(dep_.output.dim_tags).isdisjoint(dims)


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -4852,6 +4852,32 @@ def test_RepeatLayer_int():
     assert (in_seq_len * 5 == out_seq_len).all()
 
 
+def test_RepeatLayer_as_loss_flatten():
+  # https://github.com/rwth-i6/returnn_common/issues/201
+  # https://github.com/rwth-i6/returnn/issues/1115
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+  time_dim = SpatialDim('time')
+  in_dim = FeatureDim('in', 3)
+  config = Config(dict(extern_data={'data': {'dim_tags': (batch_dim, time_dim, in_dim)}}))
+  net_dict = {
+    'output': {
+      'class': 'repeat',
+      'from': 'data:data',
+      'repetitions': 1,
+      'axis': time_dim,
+      'out_dim': time_dim,
+      'loss': 'as_is',
+      'out_shape': {batch_dim, time_dim, in_dim}
+    },
+  }
+  with make_scope() as session:
+    net = TFNetwork(config=config, eval_flag=True)
+    net.construct_from_dict(net_dict)
+    feed_dict = make_feed_dict(net.extern_data)
+    fetches = net.get_fetches_dict(should_eval=True)
+    session.run(fetches, feed_dict=feed_dict)
+
+
 def test_TileLayer():
   with make_scope() as session:
     n_out = 5