Repository: rwth-i6/returnn

Issue:
Error when using segmental attention with state vector and separate pooling over the same segments
When using the config from: https://gist.github.com/robin-p-schmitt/3c4ee50f90b6d53755ade6f8523a0f2b I get the following error:

```
TensorFlow exception: 2 root error(s) found.
  (0) Invalid argument: Incompatible shapes: [8,80,12,1024] vs. [1,640,1,1]
	 [[node output/rec/pooled_segment/truediv (defined at u/schmitt/src/returnn/returnn/tf/layers/basic.py:5883) ]]
	 [[objective/loss/error/loss_ctc_error/sparse_labels/idxs/Cast_1/_1069]]
  (1) Invalid argument: Incompatible shapes: [8,80,12,1024] vs. [1,640,1,1]
	 [[node output/rec/pooled_segment/truediv (defined at u/schmitt/src/returnn/returnn/tf/layers/basic.py:5883) ]]
0 successful operations.
0 derived errors ignored.

Errors may have originated from an input operation.
Input Source operations connected to node output/rec/pooled_segment/truediv:
 output/rec/pooled_segment/x_masked_axis_2/SelectV2 (defined at u/schmitt/src/returnn/returnn/tf/util/basic.py:3662)

Input Source operations connected to node output/rec/pooled_segment/truediv:
 output/rec/pooled_segment/x_masked_axis_2/SelectV2 (defined at u/schmitt/src/returnn/returnn/tf/util/basic.py:3662)

Original stack trace for 'output/rec/pooled_segment/truediv':
  File "u/schmitt/src/returnn/rnn.py", line 11, in <module>
    main()
  File "u/schmitt/src/returnn/returnn/__main__.py", line 657, in main
    execute_main_task()
  File "u/schmitt/src/returnn/returnn/__main__.py", line 455, in execute_main_task
    engine.init_train_from_config(config, train_data, dev_data, eval_data)
  File "u/schmitt/src/returnn/returnn/tf/engine.py", line 1056, in init_train_from_config
    self.init_network_from_config(config)
  File "u/schmitt/src/returnn/returnn/tf/engine.py", line 1121, in init_network_from_config
    self._init_network(net_desc=net_dict, epoch=self.epoch)
  File "u/schmitt/src/returnn/returnn/tf/engine.py", line 1301, in _init_network
    self.network, self.updater = self.create_network(
  File "u/schmitt/src/returnn/returnn/tf/engine.py", line 1342, in create_network
    network.construct_from_dict(net_dict)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 641, in construct_from_dict
    self.construct_layer(net_dict, name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 961, in construct_layer
    return add_layer(name=name_with_prefix, layer_class=layer_class, **layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 1123, in add_layer
    layer = self._create_layer(name=name, layer_class=layer_class, **layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 1037, in _create_layer
    layer = layer_class(**layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 259, in __init__
    y = self._get_output_subnet_unit(self.cell)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 1070, in _get_output_subnet_unit
    output = cell.get_output()
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 2958, in get_output
    self._construct_output_layers_moved_out(
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3762, in _construct_output_layers_moved_out
    self.output_layers_net.layers[layer_name] = get_layer(layer_name)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3747, in get_layer
    return self.output_layers_net.construct_layer(self.net_dict, name=name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 954, in construct_layer
    layer_class.transform_config_dict(layer_desc, network=net, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 616, in transform_config_dict
    d["sources"] = [
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 617, in <listcomp>
    get_layer(src_name)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3747, in get_layer
    return self.output_layers_net.construct_layer(self.net_dict, name=name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 954, in construct_layer
    layer_class.transform_config_dict(layer_desc, network=net, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 616, in transform_config_dict
    d["sources"] = [
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 617, in <listcomp>
    get_layer(src_name)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3747, in get_layer
    return self.output_layers_net.construct_layer(self.net_dict, name=name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 954, in construct_layer
    layer_class.transform_config_dict(layer_desc, network=net, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 616, in transform_config_dict
    d["sources"] = [
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 617, in <listcomp>
    get_layer(src_name)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3747, in get_layer
    return self.output_layers_net.construct_layer(self.net_dict, name=name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 954, in construct_layer
    layer_class.transform_config_dict(layer_desc, network=net, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 616, in transform_config_dict
    d["sources"] = [
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 617, in <listcomp>
    get_layer(src_name)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3747, in get_layer
    return self.output_layers_net.construct_layer(self.net_dict, name=name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 954, in construct_layer
    layer_class.transform_config_dict(layer_desc, network=net, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 616, in transform_config_dict
    d["sources"] = [
  File "u/schmitt/src/returnn/returnn/tf/layers/base.py", line 617, in <listcomp>
    get_layer(src_name)
  File "u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 3747, in get_layer
    return self.output_layers_net.construct_layer(self.net_dict, name=name, get_layer=get_layer)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 961, in construct_layer
    return add_layer(name=name_with_prefix, layer_class=layer_class, **layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 1123, in add_layer
    layer = self._create_layer(name=name, layer_class=layer_class, **layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 1037, in _create_layer
    layer = layer_class(**layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/layers/basic.py", line 5806, in __init__
    self.output.placeholder = self.reduce(
  File "u/schmitt/src/returnn/returnn/tf/layers/basic.py", line 5883, in reduce
    x_ = x_ / tf.cast(seq_len_bc, tf.float32)
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1125, in binary_op_wrapper
    return func(x, y, name=name)
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
    return target(*args, **kwargs)
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1297, in truediv
    return _truediv_python3(x, y, name)
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1236, in _truediv_python3
    return gen_math_ops.real_div(x, y, name=name)
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 7460, in real_div
    _, _, _op, _outputs = _op_def_library._apply_op_helper(
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/framework/op_def_library.py", line 742, in _apply_op_helper
    op = g._create_op_internal(op_type_name, inputs, dtypes=None,
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 3477, in _create_op_internal
    ret = Operation(
  File "work/tools/asr/python/3.8.0_tf_2.3-v1-generic+cuda10.1/lib/python3.8/site-packages/tensorflow/python/framework/ops.py", line 1949, in __init__
    self._traceback = tf_stack.extract_stack()

Exception InvalidArgumentError() in step 0. (pid 16843)
Failing op: <tf.Operation 'output/rec/pooled_segment/truediv' type=RealDiv>
```
The full error log can be seen in the same Gist from above.

The relevant layers are:

```
    "label_model": {
        "back_prop": True,
        "class": "rec",
        "from": "data:label_ground_truth",
        "include_eos": True,
        "is_output_layer": True,
        "name_scope": "output/rec",
        "unit": {
            "att": {"axes": "except_time", "class": "merge_dims", "from": "att0"},
            "att0": {
                "add_var2_if_empty": False,
                "class": "dot",
                "from": ["att_val_split", "att_weights"],
                "reduce": "stag:att_t",
                "var1": "f",
                "var2": None,
            },
            "att_ctx": {
                "L2": None,
                "activation": None,
                "class": "linear",
                "dropout": 0.0,
                "from": "segments",
                "n_out": 1024,
                "with_bias": True,
            },
            "att_energy": {
                "class": "reinterpret_data",
                "from": "att_energy0",
                "is_output_layer": False,
                "set_dim_tags": {
                    "f": Dim(
                        kind=Dim.Types.Spatial, description="att_heads", dimension=1
                    )
                },
            },
            "att_energy0": {
                "activation": None,
                "class": "linear",
                "from": ["energy_tanh"],
                "n_out": 1,
                "with_bias": False,
            },
            "att_energy_in": {
                "class": "combine",
                "from": ["att_ctx", "att_query"],
                "kind": "add",
                "n_out": 1024,
            },
            "att_query": {
                "activation": None,
                "class": "linear",
                "from": "lm",
                "is_output_layer": False,
                "n_out": 1024,
                "with_bias": False,
            },
            "att_val": {"class": "copy", "from": "segments"},
            "att_val_split": {
                "class": "reinterpret_data",
                "from": "att_val_split0",
                "set_dim_tags": {
                    "dim:1": Dim(
                        kind=Dim.Types.Spatial, description="att_heads", dimension=1
                    )
                },
            },
            "att_val_split0": {
                "axis": "f",
                "class": "split_dims",
                "dims": (1, -1),
                "from": "att_val",
            },
            "att_weights": {
                "class": "dropout",
                "dropout": 0.1,
                "dropout_noise_shape": {"*": None},
                "from": "att_weights0",
                "is_output_layer": False,
            },
            "att_weights0": {
                "axis": "stag:att_t",
                "class": "softmax_over_spatial",
                "energy_factor": 0.03125,
                "from": "att_energy",
            },
            "energy_tanh": {
                "activation": "tanh",
                "class": "activation",
                "from": ["att_energy_in"],
            },
            "lm": {"class": "unmask", "from": "lm_masked", "mask": "prev:non_sil_mask"},
            "lm_masked": {
                "class": "masked_computation",
                "from": "prev_non_blank_embed",
                "mask": "prev:non_sil_mask",
                "unit": {
                    "class": "subnetwork",
                    "from": "data",
                    "subnetwork": {
                        "input_embed": {
                            "axes": "except_time",
                            "class": "merge_dims",
                            "from": "input_embed0",
                        },
                        "input_embed0": {
                            "class": "window",
                            "from": "data",
                            "window_left": 0,
                            "window_right": 0,
                            "window_size": 1,
                        },
                        "lm": {
                            "class": "rec",
                            "from": ["base:prev:att", "input_embed"],
                            "n_out": 1024,
                            "name_scope": "lm",
                            "unit": "nativelstm2",
                        },
                        "output": {"class": "copy", "from": "lm"},
                    },
                },
            },
            "pool_segments": {"class": "copy", "from": "segments"},
            "pooled_segment": {
                "axes": ["stag:att_t"],
                "class": "reduce",
                "from": "pool_segments",
                "mode": "mean",
            },
            "segments": {
                "class": "reinterpret_data",
                "from": "segments0",
                "set_dim_tags": {
                    "stag:sliced-time:segments": Dim(
                        kind=Dim.Types.Spatial, description="att_t"
                    )
                },
            },
            "segments0": {
                "class": "slice_nd",
                "from": "base:encoder",
                "size": "segment_lens",
                "start": "segment_starts",
            },
        },
    },
```

When I remove `base:prev:att` from the `lm_masked/lm` layer, the error does not happen.


Discussion:
```
  File "u/schmitt/src/returnn/returnn/tf/network.py", line 1037, in _create_layer
    layer = layer_class(**layer_desc)
  File "u/schmitt/src/returnn/returnn/tf/layers/basic.py", line 5806, in __init__
    self.output.placeholder = self.reduce(
  File "u/schmitt/src/returnn/returnn/tf/layers/basic.py", line 5883, in reduce
    x_ = x_ / tf.cast(seq_len_bc, tf.float32)
```
This is in `ReduceLayer.reduce`.
The code there is:
```python
x = input_data
...
x_ = x.placeholder
...
mask = x.get_sequence_mask_broadcast(axis=axis)
...
x_ = tf_util.where_bc(mask, x_, replacement_value, name="x_masked_axis_%i" % axis)
if f == tf.reduce_mean:
  seq_len_bc = tf.reshape(
    x.get_sequence_lengths(),
    [1 if (i != x.batch_dim_axis) else -1 for i in range(x.batch_ndim)])  # (1,..B..,1)
  x_ = x_ / tf.cast(seq_len_bc, tf.float32)
  f = tf.reduce_sum
```

So, from the error `[8,80,12,1024] vs. [1,640,1,1]`, it means:
- `x_` has shape [8,80,12,1024]
- `seq_len_bc` has shape [1,640,1,1]

The code itself looks broken. It assumes that the seq len has shape [B].

I will see how to fix this.

> The code itself looks broken. It assumes that the seq len has shape [B].

Oh okay. In my case the `segments` layer has seq lens of shape [B,T]

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
@@ -5892,9 +5892,7 @@ def reduce(cls, input_data, mode, axes=None, keep_dims=False, enforce_batch_dim_
 
           x_ = tf_util.where_bc(mask, x_, replacement_value, name="x_masked_axis_%i" % axis)
           if f == tf.reduce_mean:
-            seq_len_bc = tf.reshape(
-              x.get_sequence_lengths(),
-              [1 if (i != x.batch_dim_axis) else -1 for i in range(x.batch_ndim)])  # (1,..B..,1)
+            seq_len_bc = x.get_sequence_lengths_broadcast(axis=axis)
             x_ = x_ / tf.cast(seq_len_bc, tf.float32)
             f = tf.reduce_sum
       elif f == tf.reduce_mean:
diff --git a/returnn/tf/util/data.py b/returnn/tf/util/data.py
--- a/returnn/tf/util/data.py
+++ b/returnn/tf/util/data.py
@@ -5105,6 +5105,26 @@ def get_sequence_mask_broadcast(self, axis=None):
         assert seq_mask.get_shape().ndims == self.batch_ndim
     return seq_mask
 
+  def get_sequence_lengths_broadcast(self, axis=None):
+    """
+    :param int|None axis:
+    :return: seq len of some shape which is broadcastable to self.placeholder.
+      Note that this is not always possible, e.g. when the seq len has shape [B]
+      but the tensor has just shape [T]. We currently throw an error then.
+    :rtype: tf.Tensor
+    """
+    if axis is None:
+      assert self.time_dim_axis is not None
+      axis = self.time_dim_axis
+    if axis < 0:
+      assert axis + self.batch_ndim > 0
+      axis += self.batch_ndim
+    assert 0 <= axis < self.batch_ndim
+    assert axis != self.batch_dim_axis
+    tag = self.dim_tags[axis]
+    assert tag.dyn_size_ext
+    return tag.dyn_size_ext.copy_compatible_to(self, check_dtype=False, check_sparse=False).placeholder
+
   def copy_masked(self, mask_value):
     """
     :param float|int|tf.Tensor mask_value:


Test patch: