Repository: rwth-i6/returnn

Issue:
Reuse of layer for two losses with as_is flattening
Problem which was supposed to be fixed in #1078 , but the fix seems to have unwanted side effects, so for now this issue is for potential discussion.

```python
def test_double_flatten_loss():
  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
  time_dim = SpatialDim("time")
  _repeat_out_dim = time_dim * 5
  feature_dim = FeatureDim("feat", 1)
  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim), "dtype": "int32"}}})
  network = {
    'sub': {
      'class': 'combine',
      'from': ['data:data', 'data:data'],
      'kind': 'sub',
      'loss': 'as_is',
      'out_shape': {batch_dim, time_dim, feature_dim}
    },
    'repeat': {
      'class': 'repeat',
      'from': 'data:data',
      'repetitions': 5,
      'axis': time_dim,
      'out_dim': _repeat_out_dim,
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    },
    'sub_0': {
      'class': 'combine',
      'from': ['repeat', 'repeat'],
      'kind': 'sub',
      'loss': 'as_is',
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    },
    'output': {
      'class': 'copy',
      'from': 'sub_0',
      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
    }
  }
  with make_scope():
    net = TFNetwork(config=config, train_flag=True)
    net.construct_from_dict(network)
    net.get_total_loss()
```

with error:

```
layer /'data:data': [B,T|'time'[B],F|F'feat'(1)] int32
layer /'repeat': [B,T|'time*5'[B],F|F'feat'(1)] int32
layer /'sub_0': [B,T|'time*5'[B],F|F'feat'(1)] int32
layer /'output': [B,T|'time*5'[B],F|F'feat'(1)] int32
layer /'sub': [B,T|'time'[B],F|F'feat'(1)] int32
Loss flattened layer /'repeat_flat' output: Data{'repeat_flat_output', [B&Packed{'time*5'},F|F'feat'(1)], dtype='int32'}
Loss flattened layer /'sub_0' output: Data{'sub_0_output', [B&Packed{'time*5'},F|F'feat'(1)], dtype='int32'}
EXCEPTION
Traceback (most recent call last):
  File "tests/test_TFNetworkLayer.py", line 9137, in <module>
    line: globals()[arg]()  # assume function and execute
    locals:
      globals = <builtin> <built-in function globals>
      arg = <local> 'test_double_flatten_loss', len = 24
  File "tests/test_TFNetworkLayer.py", line 7321, in test_double_flatten_loss
    line: net.get_total_loss()
    locals:
      net = <local> <TFNetwork '' train>
      net.get_total_loss = <local> <bound method TFNetwork.get_total_loss of <TFNetwork '' train>>
  File "/u/hilmes/returnn/returnn/tf/network.py", line 1282, in TFNetwork.get_total_loss
    line: self.maybe_construct_objective()
    locals:
      self = <local> <TFNetwork '' train>
      self.maybe_construct_objective = <local> <bound method TFNetwork.maybe_construct_objective of <TFNetwork '' train>>
  File "/u/hilmes/returnn/returnn/tf/network.py", line 1267, in TFNetwork.maybe_construct_objective
    line: self._construct_objective()
    locals:
      self = <local> <TFNetwork '' train>
      self._construct_objective = <local> <bound method TFNetwork._construct_objective of <TFNetwork '' train>>
  File "/u/hilmes/returnn/returnn/tf/network.py", line 1249, in TFNetwork._construct_objective
    line: self._flatten_layer_with_losses()
    locals:
      self = <local> <TFNetwork '' train>
      self._flatten_layer_with_losses = <local> <bound method TFNetwork._flatten_layer_with_losses of <TFNetwork '' train>>
  File "/u/hilmes/returnn/returnn/tf/network.py", line 1567, in TFNetwork._flatten_layer_with_losses
    line: assert layer in mapped_layers
    locals:
      layer = <local> <CombineLayer 'sub' out_type=Data{[B,T|'time'[B],F|F'feat'(1)], dtype='int32'}>
      mapped_layers = <local> {<RepeatLayer 'repeat' out_type=Data{[B,T|'time*5'[B],F|F'feat'(1)], dtype='int32'}>: <FlattenBatchLayer 'repeat_flat' out_type=Data{[B&Packed{'time*5'},F|F'feat'(1)], dtype='int32'}>, <CombineLayer 'sub_0' out_type=Data{[B,T|'time*5'[B],F|F'feat'(1)], dtype='int32'}>: <CombineLayer 'sub_0' out_t...
AssertionError
```
`mapped_layers = {<RepeatLayer 'repeat' out_type=Data{[B,T|'time*5'[B],F|F'feat'(1)], dtype='int32'}>: <FlattenBatchLayer 'repeat_flat' out_type=Data{[B&Packed{'time*5'},F|F'feat'(1)], dtype='int32'}>, <CombineLayer 'sub_0' out_type=Data{[B,T|'time*5'[B],F|F'feat'(1)], dtype='int32'}>: <CombineLayer 'sub_0' out_type=Data{[B&Packed{'time*5'},F|F'feat'(1)], dtype='int32'}>}`

potentially related error from a bigger network:
```python
EXCEPTION
Traceback (most recent call last):
  File "/u/hilmes/returnn/rnn.py", line 11, in <module>
    line: main()
    locals:
      main = <local> <function main at 0x7ff53aeb9ee0>
  File "/u/hilmes/returnn/returnn/__main__.py", line 657, in main
    line: execute_main_task()
    locals:
      execute_main_task = <global> <function execute_main_task at 0x7ff53aeb9dc0>
  File "/u/hilmes/returnn/returnn/__main__.py", line 455, in execute_main_task
    line: engine.init_train_from_config(config, train_data, dev_data, eval_data)
    locals:
      engine = <global> <returnn.tf.engine.Engine object at 0x7ff506588ac0>
      engine.init_train_from_config = <global> <bound method Engine.init_train_from_config of <returnn.tf.engine.Engine object at 0x7ff506588ac0>>
      config = <global> <returnn.config.Config object at 0x7ff550627d00>
      train_data = <global> <MetaDataset 'train' epoch=1>
      dev_data = <global> <MetaDataset 'dev' epoch=1>
      eval_data = <global> None
  File "/u/hilmes/returnn/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 0x7ff506588ac0>
      self.init_network_from_config = <local> <bound method Engine.init_network_from_config of <returnn.tf.engine.Engine object at 0x7ff506588ac0>>
      config = <local> <returnn.config.Config object at 0x7ff550627d00>
  File "/u/hilmes/returnn/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 0x7ff506588ac0>
      self._init_network = <local> <bound method Engine._init_network of <returnn.tf.engine.Engine object at 0x7ff506588ac0>>
      net_desc = <not found>
      net_dict = <local> {'ctc_model_mean_absolute_difference_reduce': {'class': 'copy', 'from': 'mean_absolute_difference', 'loss': 'as_is', 'out_shape': {Dim{'phonemes_time'[B(-1)]}, Dim{B}}}, 'ctc_model_mean_
absolute_difference_0_reduce': {'class': 'copy', 'from': 'mean_absolute_difference_0', 'loss': 'as_is', 'out_sh..., len = 16
      epoch = <local> None
      self.epoch = <local> 1
  File "/u/hilmes/returnn/returnn/tf/engine.py", line 1301, 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 0x7ff506588ac0>
      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 0x7ff550627d00>
      extern_data = <local> <ExternData data={'audio_features': Data{'audio_features', [B,T|'audio_features_time'[B],F|F'audio_features_feature'(80)]}, 'phonemes': Data{'phonemes', [B,T|'phonemes_time'[B]], dtyp
e='int32', sparse_dim=Dim{F'phonemes_indices'(43)}}, 'speaker_labels': Data{'speaker_labels', [B,T|'speaker_labels...
      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> {'ctc_model_mean_absolute_difference_reduce': {'class': 'copy', 'from': 'mean_absolute_difference', 'loss': 'as_is', 'out_shape': {Dim{'phonemes_time'[B(-1)]}, Dim{B}}}, 'ctc_model_mean_
absolute_difference_0_reduce': {'class': 'copy', 'from': 'mean_absolute_difference_0', 'loss': 'as_is', 'out_sh..., len = 16
  File "/u/hilmes/returnn/returnn/tf/engine.py", line 1351, 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 0x7ff550627d00>
      network = <local> <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>
      initial_learning_rate = <local> 1.0
  File "/u/hilmes/returnn/returnn/tf/updater.py", line 177, in Updater.__init__
    line: self.loss = network.get_objective()
    locals:
      self = <local> <returnn.tf.updater.Updater object at 0x7ff44cdcaaf0>
      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/returnn/returnn/tf/network.py", line 1274, 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/returnn/returnn/tf/network.py", line 1267, 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/returnn/returnn/tf/network.py", line 1249, 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/returnn/returnn/tf/network.py", line 1555, in TFNetwork._flatten_layer_with_losses
    line: assert layer_queue  # there must be others which we need to flatten
    locals:
      layer_queue = <local> []
AssertionError
```


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
@@ -1314,6 +1314,7 @@ def _flatten_layer_with_losses(self):
     from .util.data import BatchInfo
     from .layers.basic import SourceLayer, InternalLayer, SubnetworkLayer, CopyLayer, FlattenBatchLayer
     from tensorflow.python.util import nest
+    from pprint import pformat
 
     def _relevant_dims_for_layer(layer_):
       """
@@ -1416,6 +1417,7 @@ def _check_push_flattening_to_inputs_for_layer(layer_):
       Checks whether the inputs to the layer should be flattened aswell
 
       :param LayerBase layer_:
+      :return: False when we should stop here
       :rtype: bool
       """
       if not _check_push_flattening_to_inputs_for_layer_simple(layer_):
@@ -1436,6 +1438,7 @@ def _check_push_flattening_to_inputs_for_layer(layer_):
         for dep_ in deps)
       if not valid_deps:
         return False
+      have_any_deps_which_needs_flattening = False
       for dep_ in deps:
         if dep_.output.beam:
           return False
@@ -1443,7 +1446,8 @@ def _check_push_flattening_to_inputs_for_layer(layer_):
           if any(d.dimension is None for d in set(dep_.output.dim_tags).difference(dims)):  # any other dynamic?
             return False
           layer_queue.append(dep_)
-      return True
+          have_any_deps_which_needs_flattening = True
+      return have_any_deps_which_needs_flattening
 
     def _resolve_layer(layer_):
       """
@@ -1563,7 +1567,9 @@ def _resolve_layer(layer_):
 
     # All end points must be mapped now.
     for layer in end_points:
-      assert layer in mapped_layers
+      assert layer in mapped_layers, (
+        "end point %r not mapped.\n end points:\n%s\n mapped:\n%s\n blacklist:\n%s\n starting points:\n%s" % (
+          layer, pformat(end_points), pformat(mapped_layers), pformat(blacklist), pformat(starting_points)))
     # Assign flatten_with_seq_len_mask cache to mapped layers.
     for layer, new_layer in mapped_layers.items():
       if not _should_flatten_layer_output(layer):


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -7321,6 +7321,48 @@ def test_double_flatten_loss():
     net.get_total_loss()
 
 
+def test_double_flatten_loss_1079():
+  # https://github.com/rwth-i6/returnn/issues/1079
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+  time_dim = SpatialDim("time")
+  _repeat_out_dim = time_dim * 5
+  feature_dim = FeatureDim("feat", 1)
+  config = Config({"extern_data": {"data": {"dim_tags": (batch_dim, time_dim, feature_dim), "dtype": "int32"}}})
+  network = {
+    'sub': {
+      'class': 'combine',
+      'from': ['data:data', 'data:data'],
+      'kind': 'sub',
+      'loss': 'as_is',
+      'out_shape': {batch_dim, time_dim, feature_dim}
+    },
+    'repeat': {
+      'class': 'repeat',
+      'from': 'data:data',
+      'repetitions': 5,
+      'axis': time_dim,
+      'out_dim': _repeat_out_dim,
+      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
+    },
+    'sub_0': {
+      'class': 'combine',
+      'from': ['repeat', 'repeat'],
+      'kind': 'sub',
+      'loss': 'as_is',
+      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
+    },
+    'output': {
+      'class': 'copy',
+      'from': 'sub_0',
+      'out_shape': {batch_dim, feature_dim, _repeat_out_dim}
+    }
+  }
+  with make_scope():
+    net = TFNetwork(config=config, train_flag=True)
+    net.construct_from_dict(network)
+    net.get_total_loss()
+
+
 def test_LossLayer_sublayers():
   from returnn.tf.util.basic import Dim
   n_in, n_out = 7, 11