Repository: rwth-i6/returnn

Issue:
`MergeDimsLayer` should use `keep_order=True` (and disallow `keep_order=False`)
This is to follow the principle that the order of axes should never matter.
(`keep_order` is about using the user specified order in the layer `axes` argument, and ignoring the incoming data axes order.)

This probably needs a new behavior version (#508).



Discussion:
I now realize that this will also require consistent and unique dim tags (#632) because sth like `"axes": "static"` would merge all static axes, but again the order is not well defined here. So sth like this would be disallowed. You would have to use dim tags and specify it like `"axes": [dim_tag_1, dim_tag_2]` or so.

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
@@ -2838,10 +2838,10 @@ class MergeDimsLayer(_ConcatInputLayer):
   """
   layer_class = "merge_dims"
 
-  def __init__(self, axes, keep_order=False, n_out=None, **kwargs):
+  def __init__(self, axes, keep_order=NotSpecified, n_out=None, **kwargs):
     """
-    :param str|list[str]|list[int] axes: see Data.get_axes_from_description(), e.g. "except_time"
-    :param bool keep_order: By default (for historical reasons), the axes are sorted, and then merged.
+    :param str|list[DimensionTag|str] axes: see :func:`Data.get_axis_from_description`
+    :param bool|NotSpecified keep_order: The old default was: the axes are sorted, and then merged.
       Thus, the order of incoming axes will influence the result.
       E.g. inputs [B,S,F] and [B,F,S], with ``axes=["S","F"]``, will get different results,
       although the output shape is [B,S*F] in both cases.
@@ -2849,15 +2849,23 @@ def __init__(self, axes, keep_order=False, n_out=None, **kwargs):
       and all layers should behave in the same way, no matter the order.
       It is recommended to set ``keep_order=True``, such that the order defined in ``axes`` defines the behavior,
       and not the incoming axis order.
+      Since behavior version 6, this is already the case.
     :param int|None n_out:
     """
+    from returnn.util import BehaviorVersion
     super(MergeDimsLayer, self).__init__(**kwargs)
+    if keep_order is NotSpecified:
+      keep_order = True if BehaviorVersion.get() >= 6 else False
+    BehaviorVersion.require(
+      condition=keep_order, message="MergeDimsLayer, only keep_order=True is allowed", version=6)
     if keep_order:
-      assert isinstance(axes, (tuple, list)), "%s: unique axes %r required" % (self, axes)
+      assert isinstance(axes, (tuple, list)), (
+        "%s: axes %r must be a list or tuple, to have a well defined order in input %s" % (self, axes, self.input_data))
       axes_ = []
       for axis in axes:
         axis_ = self.input_data.get_axes_from_description(axis, allow_int=False)
-        assert len(axis_) <= 1, "%s: unique axes %r required, but got %r -> %r" % (self, axes, axis, axis_)
+        assert len(axis_) <= 1, (
+          "%s: unique axes %r required in input %s, but got %r -> %r" % (self, axes, self.input_data, axis, axis_))
         axes_ += axis_
       axes = axes_
     else:
@@ -2971,18 +2979,20 @@ def _set_output_sizes(self, merge_axes):
     target_tag.dyn_size_ext = out_size
 
   @classmethod
-  def get_out_data_from_opts(cls, name, axes, keep_order=False,
+  def get_out_data_from_opts(cls, name, axes, keep_order=NotSpecified,
                              sources=(), n_out=NotSpecified, out_type=None, **kwargs):
     """
     :param str name:
     :param str|list[str] axes:
-    :param bool keep_order:
+    :param bool|NotSpecified keep_order:
     :param list[LayerBase] sources:
     :param int|None|NotSpecified n_out:
     :param None|dict[str] out_type:
     :rtype: Data
     """
-    from ..util.data import DimensionTag
+    from returnn.util import BehaviorVersion
+    if keep_order is NotSpecified:
+      keep_order = True if BehaviorVersion.get() >= 6 else False
     assert not out_type, "currently ignored"
     input_data = get_concat_sources_data_template(sources)
     data = input_data.copy(name="%s_output" % name)
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
@@ -3530,6 +3530,11 @@ def get_axes_from_description(self, axes, allow_int=NotSpecified):
           s += len(static_axes)
         assert 0 <= s < len(static_axes), "%s get_axes_from_description: %r invalid" % (self, axes)
         return [static_axes[s]]
+      elif re.match("(dim):\\d+$", axes):
+        s = int(axes.split(":")[1])
+        dims = [a for a in range(self.batch_ndim) if self.batch_shape[a] == s]
+        assert dims, "%s get_axes_from_description: no dim %i found" % (self, s)
+        return dims
       elif axes in ["f", "feature", "non_spatial"]:
         return self.get_feature_batch_axes()
       elif all([a in "btf" for a in axes]):
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -209,7 +209,7 @@ class BehaviorVersion:
   The version will be set after the config is defined at __main__.init_config() or Engine.__init__()
   """
 
-  _latest_behavior_version = 5
+  _latest_behavior_version = 6
   _behavior_version = None  # type: typing.Optional[int]
 
   @classmethod


Test patch:
diff --git a/tests/test_TFEngine.py b/tests/test_TFEngine.py
--- a/tests/test_TFEngine.py
+++ b/tests/test_TFEngine.py
@@ -1233,7 +1233,7 @@ def test_attention_no_encoder_dependency():
                             'n_out': 4, 'padding': 'same'},
           "location_feedback": {'class': 'linear', 'from': ['convolved_att'], 'n_out': 6, 'activation': None},
           "att_energy_in": {'class': 'combine', 'kind': 'add', 'from': ['location_feedback', 's_transformed']},
-          "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights"},
+          "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights", "auto_squeeze": True},
         },
       },
       "decision": {"class": "decide", "from": ["output"], "loss": "edit_distance"}
@@ -1345,7 +1345,7 @@ def test_attention_convolutional_feedback_variant1():
     "location_feedback": {'class': 'linear', 'from': ['convolved_att'], 'n_out': 6, 'activation': None},
     "att_energy_in": {'class': 'combine', 'kind': 'add', 'from': [
       'base:enc_transformed', 'location_feedback', 's_transformed']},
-    "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights"},
+    "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights", "auto_squeeze": True},
   }
 
   check_attention_variant(recurrent_unit_dict)
@@ -1373,7 +1373,7 @@ def test_attention_convolutional_feedback_variant2():
     "location_feedback": {'class': 'linear', 'from': ['convolved_att'], 'n_out': 6, 'activation': None},
     "att_energy_in": {'class': 'combine', 'kind': 'add', 'from': [
       'base:enc_transformed', 'location_feedback', 's_transformed']},
-    "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights"},
+    "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights", "auto_squeeze": True},
   }
 
   check_attention_variant(recurrent_unit_dict)
@@ -1412,7 +1412,7 @@ def test_attention_convolutional_feedback_variant3():
     "location_feedback": {'class': 'linear', 'from': ['convolved_att'], 'n_out': 6, 'activation': None},
     "att_energy_in": {'class': 'combine', 'kind': 'add', 'from': [
       'base:enc_transformed', 'location_feedback', 's_transformed']},
-    "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights"},
+    "c": {"class": "generic_attention", "base": "base:encoder", "weights": "att_weights", "auto_squeeze": True},
   }
 
   check_attention_variant(recurrent_unit_dict)
@@ -2135,7 +2135,7 @@ def test_rec_subnet_construct_1():
       "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": 1, "shape": (None, 1)}},
-      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["target_embed", "att"], "n_out": 10},
       "s2": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["s"], "n_out": 10},
       "readout_in": {"class": "linear", "from": ["prev:s2", "prev:target_embed", "att"], "activation": None, "n_out": 10},
@@ -2192,7 +2192,7 @@ def test_rec_subnet_construct_2():
       "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": 1, "shape": (None, 1)}},
-      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["target_embed", "att"], "n_out": 10},
       "s2": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["s"], "n_out": 10},
       "readout_in": {"class": "linear", "from": ["prev:s2", "prev:target_embed", "att"], "activation": None, "n_out": 10},
@@ -2255,7 +2255,7 @@ def test_rec_subnet_construct_3():
       "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": 1, "shape": (None, 1)}},
-      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["target_embed", "att"], "n_out": 10},
       "s2": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:s", "prev:target_embed", "att"], "n_out": 10},
       "readout_in": {"class": "linear", "from": ["s2"], "activation": None, "n_out": 10},
@@ -2288,9 +2288,9 @@ def test_rec_subnet_eval_init_out_apply0():
   # (also defined by num_inputs & num_outputs)
   beam_size = 3
   AttNumHeads = 2
-  EncKeyTotalDim = AttNumHeads * 2
+  EncKeyTotalDim = AttNumHeads * 5
   EncKeyPerHeadDim = EncKeyTotalDim // AttNumHeads
-  EncValueTotalDim = AttNumHeads * 2
+  EncValueTotalDim = AttNumHeads * 5
   EncValuePerHeadDim = EncValueTotalDim // AttNumHeads
   network = {
     "lstm0_fw": {"class": "rec", "unit": "nativelstm2", "n_out": 2, "direction": 1, "from": "data:data"},
@@ -2334,7 +2334,7 @@ def test_rec_subnet_eval_init_out_apply0():
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": 1, "shape": (None, 1)}, "initial_output": "apply(0)"},  # (B, enc-T, 1)
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "except_batch", "from": ["att0"]},  # (B, H*V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "att0"},  # (B, H*V)
 
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["target_embed", "att"], "n_out": 2},  # transform
       "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -2768,13 +2768,13 @@ def custom_construction_algo(idx, net_dict):
 def test_net_safe_log_to_log_softmax():
   n_out = 5
   net_dict = {
-    "ff_in_window": {"class": "window", "window_size": 3, "from": "data:data"},  # (B,T,3,3)
-    "ff_in": {"class": "merge_dims", "axes": "except_time", "from": ["ff_in_window"]},  # (B,T,9)
-    "ff0": {"class": "hidden", "activation": "relu", "n_out": 8, "L2": 0.01, "from": ["ff_in"]},  # (B,T,8)
-    "ff_out": {"class": "softmax", "n_out": n_out, "from": ["ff0"]},  # (B,T,5)
+    "ff_in_window": {"class": "window", "window_size": 4, "from": "data:data"},  # (B,T,4,3)
+    "ff_in": {"class": "merge_dims", "axes": ["dim:3", "dim:4"], "from": "ff_in_window"},  # (B,T,9)
+    "ff0": {"class": "hidden", "activation": "relu", "n_out": 8, "L2": 0.01, "from": "ff_in"},  # (B,T,8)
+    "ff_out": {"class": "softmax", "n_out": n_out, "from": "ff0"},  # (B,T,5)
     "ff_out_prior": {
       "class": "accumulate_mean", "exp_average": 0.001,
-      "is_prob_distribution": True, "from": ["ff_out"]},  # (5,)
+      "is_prob_distribution": True, "from": "ff_out"},  # (5,)
     "output": {
       "class": "combine", "kind": "eval", "from": ["ff_out", "ff_out_prior"],
       "eval": "safe_log(source(0)) - safe_log(source(1))",
@@ -2826,7 +2826,7 @@ def test_preload_from_files():
           "class": "linear", "activation": None, "n_out": n_hidden, "from": "data:data",
           'bias_init': 1.0, 'forward_weights_init': 'orthogonal'},
         "output": {
-          "class": "linear", "activation": None, "n_out": n_out, "from": ["l1"],
+          "class": "linear", "activation": None, "n_out": n_out, "from": "l1",
           'bias_init': 2.0, 'forward_weights_init': 'orthogonal'}
       }
     })
@@ -3366,7 +3366,7 @@ def test_attention_forward_hdf_then_unflatten_2d():
       # (B, enc-T, 1)
       "energy": {"class": "linear", "activation": None, "with_bias": False, "from": ["energy_tanh"], "n_out": 1},
       "att_weights": {"class": "softmax_over_spatial", "from": ["energy"], "is_output_layer": True},  # (B, enc-T, 1)
-      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+      "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:target_embed", "prev:att"], "n_out": 10},
       "readout_in": {"class": "linear", "from": ["s", "prev:target_embed", "att"], "activation": None, "n_out": 10},
       "readout": {"class": "reduce_out", "mode": "max", "num_pieces": 2, "from": ["readout_in"]},
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -1226,9 +1226,9 @@ def test_dot_layer_shuffled_remaining_dims_static():
       "b": {"class": "transpose", "from": ["a"], "perm": {"static:0": "static:1", "static:1": "static:0"}},
       "dot": {
         "class": "dot", "from": ["a", "b"],
-        "red1": "static:-1", "red2": "static:-1", "var1": None, "var2": None,
+        "red1": "dim:5", "red2": "dim:5", "var1": None, "var2": None,
         "debug": True},
-      "output": {"class": "merge_dims", "axes": "static", "from": ["dot"]}
+      "output": {"class": "merge_dims", "axes": ["dim:2", "dim:3"], "from": "dot"}
     }
     config = Config()
     config.update({
@@ -1839,17 +1839,17 @@ def _check_MergeDimsLayer(session, in_data_opts, in_static_shape, opts, out_data
 
 def test_MergeDimsLayer_basic():
   with make_scope() as session:
-    _check_MergeDimsLayer(session, {"shape": (4, 7), "time_dim_axis": None}, (2, 4, 7), {"axes": "except_batch"}, (4 * 7,), (2, 4 * 7))
-    _check_MergeDimsLayer(session, {"shape": (4, None, 7), "time_dim_axis": None}, (2, 4, 3, 7), {"axes": "static"}, (None, 4 * 7), (2, 3, 4 * 7))
-    _check_MergeDimsLayer(session, {"shape": (4, None, 7), "time_dim_axis": 2}, (2, 4, 3, 7), {"axes": "static"}, (None, 4 * 7), (2, 3, 4 * 7))
-    _check_MergeDimsLayer(session, {"shape": (1, None), "time_dim_axis": 2, "feature_dim_axis": 1}, (2, 1, 4), {"axes": "except_batch"}, (None,), (2, 4))
+    _check_MergeDimsLayer(session, {"shape": (4, 7), "time_dim_axis": None}, (2, 4, 7), {"axes": ["dim:4", "dim:7"]}, (4 * 7,), (2, 4 * 7))
+    _check_MergeDimsLayer(session, {"shape": (4, None, 7), "time_dim_axis": None}, (2, 4, 3, 7), {"axes": ["dim:4", "dim:7"]}, (None, 4 * 7), (2, 3, 4 * 7))
+    _check_MergeDimsLayer(session, {"shape": (4, None, 7), "time_dim_axis": 2}, (2, 4, 3, 7), {"axes": ["dim:4", "dim:7"]}, (None, 4 * 7), (2, 3, 4 * 7))
+    _check_MergeDimsLayer(session, {"shape": (1, None), "time_dim_axis": 2, "feature_dim_axis": 1}, (2, 1, 4), {"axes": ["F", "T"]}, (None,), (2, 4))
 
 
 def test_MergeDimsLayer_size_placeholder():
   with make_scope() as session:
     _check_MergeDimsLayer(
       session,
-      {"shape": (None, 2), "time_dim_axis": 1, "feature_dim_axis": 2}, (3, 4, 2), {"axes": "except_batch"}, (None,), (3, 8),
+      {"shape": (None, 2), "time_dim_axis": 1, "feature_dim_axis": 2}, (3, 4, 2), {"axes": ["T", "F"]}, (None,), (3, 8),
       in_sizes={0: (4, 2, 1)}, out_sizes={0: (8, 4, 2)})
 
 
@@ -1858,7 +1858,7 @@ def test_MergeDimsLayer_batch_time_ext():
     n_batch = 11
     n_time = 13
     _check_MergeDimsLayer(
-      session, {"shape": (None, 5, 3)}, (n_batch, n_time, 5, 3), {"axes": "BT"}, (5, 3), (n_batch * n_time, 5, 3))
+      session, {"shape": (None, 5, 3)}, (n_batch, n_time, 5, 3), {"axes": ["B", "T"]}, (5, 3), (n_batch * n_time, 5, 3))
 
 
 def test_MergeDimsLayer_batch_time_time_major():
@@ -1868,7 +1868,7 @@ def test_MergeDimsLayer_batch_time_time_major():
     layer = _check_MergeDimsLayer(
       session,
       {"shape": (None, 5), "time_dim_axis": 0, "batch_dim_axis": 1}, (n_time, n_batch, 5),
-      {"axes": "BT"}, (5,), (n_time * n_batch, 5))
+      {"axes": ["B", "T"]}, (5,), (n_time * n_batch, 5))
     assert layer.output.batch_dim_axis == 0
     assert layer.output.time_dim_axis is None
 
@@ -1880,7 +1880,7 @@ def test_MergeDimsLayer_batch_time_time_major_ext():
     layer = _check_MergeDimsLayer(
       session,
       {"shape": (None, 5, 3), "time_dim_axis": 0, "batch_dim_axis": 1}, (n_time, n_batch, 5, 3),
-      {"axes": "BT"}, (5, 3), (n_time * n_batch, 5, 3))
+      {"axes": ["B", "T"]}, (5, 3), (n_time * n_batch, 5, 3))
     assert layer.output.batch_dim_axis == 0
     assert layer.output.time_dim_axis is None  # Note: This behavior was changed.
 
@@ -1892,7 +1892,7 @@ def test_MergeDimsLayer_except_time_ext():
     layer = _check_MergeDimsLayer(
       session,
       {"shape": (3, None, 5), "time_dim_axis": 2}, (n_batch, 3, n_time, 5),
-      {"axes": "except_time"}, (None, 15), (n_batch, n_time, 15))
+      {"axes": ["dim:3", "dim:5"]}, (None, 15), (n_batch, n_time, 15))
     assert layer.output.batch_dim_axis == 0 and layer.output.time_dim_axis == 1
 
 
@@ -1902,7 +1902,7 @@ def test_MergeDimsLayer_static_time():
     layer = _check_MergeDimsLayer(
       session,
       {"shape": (3, 5), "time_dim_axis": 1}, (n_batch, 3, 5),
-      {"axes": "static"}, (15,), (n_batch, 15))
+      {"axes": ["dim:3", "dim:5"]}, (15,), (n_batch, 15))
     assert layer.output.batch_dim_axis == 0 and layer.output.feature_dim_axis == 1
     assert layer.output.time_dim_axis is None
 
@@ -1913,7 +1913,7 @@ def test_MergeDimsLayer_feat_static_static():
     layer = _check_MergeDimsLayer(
       session,
       {"shape": (None, 8, 2, 3), "feature_dim_axis": 2}, (n_batch, 7, 8, 2, 3),
-      {"axes": ["F", "static:1"]}, (None, 16, 3), (n_batch, 7, 16, 3))
+      {"axes": ["F", "dim:2"]}, (None, 16, 3), (n_batch, 7, 16, 3))
     assert (layer.output.batch_dim_axis, layer.output.time_dim_axis, layer.output.feature_dim_axis) == (0, 1, 2)
 
 
@@ -1977,7 +1977,7 @@ def test_MergeDimsLayer_SplitBatchTimeLayer_time_major():
         size_placeholder={0: tf.constant(seq_lens)}))
     assert input_layer.output.is_time_major
     net.construct_from_dict({
-      "merge_dims": {"class": "merge_dims", "from": "input", "axes": "BT"},
+      "merge_dims": {"class": "merge_dims", "from": "input", "axes": ["B", "T"]},
       "split_dims": {"class": "split_batch_time", "from": "merge_dims", "base": "input"},
       "output": {"class": "copy", "from": "split_dims"}
     })
@@ -1998,7 +1998,7 @@ def test_MergeDimsLayer_SplitBatchTimeLayer_two_time_axes():
         "debug_print_layer_output_template": True}))
     feed_dict = make_feed_dict(net.extern_data)
     net.construct_from_dict({
-      "merge_dims": {"class": "merge_dims", "from": "data", "axes": "BT"},
+      "merge_dims": {"class": "merge_dims", "from": "data", "axes": ["B", "T"]},
       "split_dims": {"class": "split_batch_time", "from": "merge_dims", "base": "data"},
       "output": {"class": "copy", "from": "split_dims"}
     })
@@ -2048,7 +2048,7 @@ def test_MergeDimsLayer_simple_feat():
   with make_scope() as session:
     net = TFNetwork(config=config)
     net.construct_from_dict({
-      "output": {"class": "merge_dims", "axes": "static", "from": "data:data"}})
+      "output": {"class": "merge_dims", "axes": ["dim:%i" % n_in1, "dim:%i" % n_in2], "from": "data:data"}})
     out_t = net.get_default_output_layer().output.placeholder
     assert out_t.shape.as_list() == [None, None, n_in1 * n_in2]
     in_v = numpy.arange(0, n_batch * n_time * n_in1 * n_in2).astype("float32").reshape((n_batch, n_time, n_in1, n_in2))
@@ -2836,7 +2836,7 @@ def test_reuse_params_map_custom_dep_loop():
                               "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                               "eval": "source(0) + source(1) * source(2) * 0.5",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
                        "n_out": 2 * 6},
@@ -3441,14 +3441,14 @@ def test_conv_window_merge_dims():
                'n_out': 12,
                'padding': 'valid',
                'strides': 16},
-    'flatten_conv': {'axes': 'except_time',
+    'flatten_conv': {'axes': ['dim:4', 'dim:12'],
                      'class': 'merge_dims',
                      'from': ['conv_2'],
                      'n_out': 48},
     'window_1': {'class': 'window',
                  'from': ['flatten_conv'],
                  'window_size': 17},
-    'flatten_window': {'axes': 'except_time',
+    'flatten_window': {'axes': ['dim:17', 'dim:48'],
                        'class': 'merge_dims',
                        'from': ['window_1']},
     'output': {'activation': None,
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -2183,7 +2183,9 @@ def test_target_with_beam():
                                     "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}, "trainable": False},
       "teacher_att0": {"class": "generic_attention", "weights": "teacher_att_weights", "base": "base:teacher_enc_value",
                        "trainable": False},  # (B, H, V)
-      "teacher_att": {"class": "merge_dims", "axes": "except_batch", "from": ["teacher_att0"], "trainable": False},
+      "teacher_att": {
+        "class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim],
+        "from": "teacher_att0", "trainable": False},
       # (B, H*V)
       "teacher_s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:teacher_target_embed", "prev:teacher_att"],
                     "n_out": 5, "dropout": 0.3, "trainable": False},  # transform
@@ -2248,7 +2250,7 @@ def test_target_with_beam():
                                     "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}},
       "student_att0": {"class": "generic_attention", "weights": "student_att_weights",
                        "base": "base:student_enc_value"},  # (B, H, V)
-      "student_att": {"class": "merge_dims", "axes": "except_batch", "from": ["student_att0"]},  # (B, H*V)
+      "student_att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "student_att0"},  # (B, H*V)
       "student_s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:student_target_embed", "prev:student_att"],
                     "n_out": 5, "dropout": 0.3},  # transform
       "student_readout_in": {"class": "linear", "from": ["student_s", "prev:student_target_embed", "student_att"],
@@ -2336,7 +2338,7 @@ def get_net_dict():
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -2502,7 +2504,7 @@ def get_net_dict(l2_target_embed=0.0, l2_readout_in=0.0):
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -2611,7 +2613,7 @@ def test_rec_layer_move_out_of_loop_ref_att_generic_att():
                                      "eval": "source(0) * source(1) * 0.5"},
       "accum_att_weights": {"class": "cumsum", "from": "att_weights_with_fertility"},
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "except_time", "from": ["att0"]},  # (B, H*V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "att0"},  # (B, H*V)
       "s": {"class": "rnn_cell", "unit": "standardlstm",
             "from": ["target_embed", "att"], "n_out": 10},
       "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -2818,7 +2820,7 @@ def test_rec_layer_rnn_train_and_search():
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}},
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "except_batch", "from": ["att0"]},  # (B, H*V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "att0"},  # (B, H*V)
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:target_embed", "prev:att"], "n_out": 10},
       "readout_in": {"class": "linear", "from": ["s", "prev:target_embed", "att"], "activation": None, "n_out": 10},
       "readout": {"class": "reduce_out", "mode": "max", "num_pieces": 2, "from": ["readout_in"]},
@@ -2981,7 +2983,7 @@ def test_rec_layer_local_att_train_and_search():
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}},
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "except_batch", "from": ["att0"]},  # (B, H*V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "att0"},  # (B, H*V)
       "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:target_embed", "prev:att"], "n_out": 10},
       # transform
       "readout_in": {"class": "linear", "from": ["s", "prev:target_embed", "att"], "activation": None, "n_out": 20},
@@ -3169,7 +3171,7 @@ def get_net_dict():
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -3635,7 +3637,7 @@ def test_reclayer_optimize_out_dot():
       # energy outside the loop will be (B, H, enc-T, dec-T). I.e. enc-T is still the first time axis.
       "att_weights": {"class": "softmax_over_spatial", "from": ["energy"]},  # (B, enc-T, H, 1)
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": ["att0"]},  # (B, H*V); Use "static" here.
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "att0"},  # (B, H*V); Use "static" here.
       },
     shared_base_net={
       "encoder": {"class": "copy", "from": ["data"]},
@@ -3681,7 +3683,7 @@ def test_reclayer_optimize_out_dot_consistent_axes():
       # It still works due to time-dim-axis being set to the first source.
       "att_weights": {"class": "softmax_over_spatial", "from": "energy"},  # (B, T, H)
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": "att0"},  # (B, H*V); Use "static" here.
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % n_heads, "dim:%i" % n_value], "from": "att0"},  # (B, H*V); Use "static" here.
       },
     shared_base_net={
       "encoder": {"class": "copy", "from": "data"},
@@ -3725,7 +3727,7 @@ def test_reclayer_optimize_out_dot_consistent_axes_enc_dec():
       # then to the respective var axes.
       "att_weights": {"class": "softmax_over_spatial", "from": "energy"},  # (B, enc-T, H)
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": "att0"},  # (B, H*V); Use "static" here.
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % n_heads, "dim:%i" % n_value], "from": "att0"},  # (B, H*V); Use "static" here.
     },
     shared_base_net={
       # Use conv with padding valid to make sure we get another time dim,
@@ -3770,7 +3772,7 @@ def test_reclayer_optimize_out_dot_kv_in_rec():
                  "from": ["enc_ctx", "att_query"]},
       "att_weights": {"class": "softmax_over_spatial", "from": ["energy"]},  # (B, enc-T, H, 1)
       "att0": {"class": "generic_attention", "weights": "att_weights", "base": "enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": ["att0"]},  # (B, H*V); Use "static" here.
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncValuePerHeadDim], "from": "att0"},  # (B, H*V); Use "static" here.
       },
     shared_base_net={
       # need to mark 'encoder' as output layer, otherwise it will not be constructed.
@@ -3932,7 +3934,7 @@ def test_reclayer_att_with_kv_in_rec():
         'att_weights': {
           'axis': 'T', 'class': 'softmax_over_spatial', 'from': ['att_energy']},
         'att_output': {'class': 'generic_attention', 'weights': 'att_weights', 'base': 'att_value'},
-        'att_att': {'axes': 'static', 'class': 'merge_dims', 'from': ['att_output']},
+        'att_att': {'axes': ["dim:2", "dim:3"], 'class': 'merge_dims', 'from': ['att_output']},
         'end': {'class': 'compare', 'from': ['output'], 'value': 0},
         'output': {
           'beam_size': 4, 'class': 'choice', 'from': ['output_prob'], 'initial_output': 'zeros',
@@ -4314,7 +4316,8 @@ def add_trafo_dec_layer(self, db, d, inp, output):
 
     d[output + '_att0'] = {"class": "generic_attention", "weights": output + '_att_weights_drop',
                            "base": 'base:' + output + '_att_value'}  # (B, H, V)
-    d[output + '_att_att'] = {"class": "merge_dims", "axes": "static",
+    d[output + '_att_att'] = {"class": "merge_dims",
+                              "axes": ["dim:%i" % self.AttNumHeads, "dim:%i" % self.EncValuePerHeadDim],
                               "from": [output + '_att0']}  # (B, H*V) except_batch
     d[output + '_att_lin'] = {"class": "linear", "activation": None, "with_bias": False,
                               "from": [output + '_att_att'],
@@ -5034,7 +5037,7 @@ def get_net_dict():
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -5111,7 +5114,7 @@ def add_lstm(i, direction, src):
       return name
     network["%s_win" % name] = {
       "class": "window", "window_size": lstm_window, "window_right": lstm_window - 1, "from": src}  # (B,T,W,D)
-    network["%s_mdims" % name] = {"class": "merge_dims", "axes": "BT", "from": ["%s_win" % name]}  # (B*T,W,D)
+    network["%s_mdims" % name] = {"class": "merge_dims", "axes": ["B", "T"], "from": ["%s_win" % name]}  # (B*T,W,D)
     network["%s_rdims" % name] = {
       "class": "reinterpret_data", "enforce_batch_major": True, "set_axes": {"T": "spatial"},
       "from": ["%s_mdims" % name]}  # (B*T,W,D)
@@ -5146,6 +5149,7 @@ def test_GenericAttentionLayer_basic0():
   time = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
   kwargs = dict(
     name="att", network=net,
+    auto_squeeze=True,
     weights=InternalLayer(
       name="att_weights", network=net,
       output=Data(
@@ -5171,6 +5175,7 @@ def test_GenericAttentionLayer_basic():
   time = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
   kwargs = dict(
     name="att", network=net,
+    auto_squeeze=True,
     weights=InternalLayer(
       name="att_weights", network=net,
       output=Data(
@@ -5217,6 +5222,7 @@ def test_GenericAttentionLayer_weights_auto_squeeze_time_end():
   time = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
   kwargs = dict(
     name="att", network=net,
+    auto_squeeze=True,
     weights=InternalLayer(
       name="att_weights", network=net,
       output=Data(
@@ -5242,6 +5248,7 @@ def test_GenericAttentionLayer_weights_static_time_axis():
   time = DimensionTag(kind=DimensionTag.Types.Spatial, description="time")
   kwargs = dict(
     name="att", network=net,
+    auto_squeeze=True,
     weights=InternalLayer(
       name="att_weights", network=net,
       output=Data(
@@ -5294,6 +5301,7 @@ def test_GenericAttentionLayer_weights_heads_auto_squeeze_time_end():
   num_heads = 8
   kwargs = dict(
     name="att", network=net,
+    auto_squeeze=True,
     weights=InternalLayer(
       name="att_weights", network=net,
       output=Data(
@@ -5484,7 +5492,7 @@ def test_att_train_search_loss_prev_beam():
 def test_MaskedComputationLayer_search_choices_resolution():
   beam_size = 3
   EncKeyTotalDim = 10
-  AttNumHeads = 1
+  AttNumHeads = 2
   target = "classes"
   num_classes = 13
   blank_idx = num_classes - 2
@@ -5516,8 +5524,8 @@ def test_MaskedComputationLayer_search_choices_resolution():
       "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}},
-      "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": "att0"},  # (B, H*V)
+      "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value", "auto_squeeze": False},  # (B, H, V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncKeyTotalDim], "from": "att0"},  # (B, H*V)
 
       'not_blank_mask': {'class': 'compare', 'from': ['output'], 'value': blank_idx, 'kind': 'not_equal',
                          'initial_output': True},
@@ -5566,7 +5574,7 @@ def test_MaskedComputationLayer_search_choices_resolution():
 def test_MaskedComputationLayer_subnet_search_choices_resolution():
   beam_size = 3
   EncKeyTotalDim = 10
-  AttNumHeads = 1
+  AttNumHeads = 2
   target = "classes"
   num_classes = 13
   blank_idx = num_classes - 2
@@ -5598,8 +5606,8 @@ def test_MaskedComputationLayer_subnet_search_choices_resolution():
       "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}},
-      "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": "att0"},  # (B, H*V)
+      "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value", "auto_squeeze": False},  # (B, H, V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncKeyTotalDim], "from": "att0"},  # (B, H*V)
 
       'not_blank_mask': {'class': 'compare', 'from': ['output'], 'value': blank_idx, 'kind': 'not_equal',
                          'initial_output': True},
@@ -5657,7 +5665,7 @@ def test_MaskedComputationLayer_subnet_search_choices_resolution():
 def test_MaskedComputationLayer_subnet_rec_search():
   beam_size = 3
   EncKeyTotalDim = 10
-  AttNumHeads = 1
+  AttNumHeads = 2
   target = "classes"
   num_classes = 13
   blank_idx = num_classes - 2
@@ -5689,8 +5697,8 @@ def test_MaskedComputationLayer_subnet_rec_search():
       "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:inv_fertility"],
                             "eval": "source(0) + source(1) * source(2) * 0.5",
                             "out_type": {"dim": AttNumHeads, "shape": (None, AttNumHeads)}},
-      "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value"},  # (B, H, V)
-      "att": {"class": "merge_dims", "axes": "static", "from": "att0"},  # (B, H*V)
+      "att0": {"class": "generic_attention", "weights": "att_weights", "base": "base:enc_value", "auto_squeeze": False},  # (B, H, V)
+      "att": {"class": "merge_dims", "axes": ["dim:%i" % AttNumHeads, "dim:%i" % EncKeyTotalDim], "from": "att0"},  # (B, H*V)
 
       'not_blank_mask': {'class': 'compare', 'from': ['output'], 'value': blank_idx, 'kind': 'not_equal',
                          'initial_output': True},
@@ -5907,7 +5915,7 @@ def test_MaskedComputationLayer_outside():
 def test_subnet_deps_search():
   beam_size = 3
   EncKeyTotalDim = 10
-  AttNumHeads = 1
+  AttNumHeads = 2
   target = "classes"
   num_classes = 13
   from test_TFNetworkLayer import make_feed_dict
@@ -5934,8 +5942,8 @@ def test_subnet_deps_search():
                               'eval': 'source(0) + source(1) * source(2) * 0.5',
                               'from': ['prev:accum_att_weights', 'att_weights', 'base:inv_fertility'],
                               'out_type': {'dim': AttNumHeads, 'shape': (None, AttNumHeads)}},
-        'att': {'axes': 'except_batch', 'class': 'merge_dims', 'from': 'att0'},
-        'att0': {'base': 'base:enc_value', 'class': 'generic_attention', 'weights': 'att_weights'},
+        'att': {'axes': ["dim:%i" % AttNumHeads, "dim:%i" % EncKeyTotalDim], 'class': 'merge_dims', 'from': 'att0'},
+        'att0': {'base': 'base:enc_value', 'class': 'generic_attention', 'weights': 'att_weights', "auto_squeeze": False},
         'att_weights': {'class': 'softmax_over_spatial', 'from': 'energy'},
         'combo_output_prob': {'class': 'eval',
                               'eval': 'safe_log(source(0)) - 0.22 * safe_log(source(1))',
@@ -6052,7 +6060,7 @@ def test_untrainable_sublayers():
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -6117,7 +6125,7 @@ def test_untrainable_reclayer():
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -6182,7 +6190,7 @@ def test_trainable_sublayers():
         "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
                               "eval": "source(0) + source(1) / (2.0 * source(2))",
                               "out_type": {"dim": 1, "shape": (None, 1)}},
-        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder", "auto_squeeze": True},
         "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
               "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
         "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
@@ -6779,7 +6787,7 @@ def test_generalized_non_rec_self_attention():
       "red1": new_dim, "red2": new_dim,
       "var1": time_dim, "var2": "static:-1"},  # [B,H,T,V]
     "att_new": {
-      "class": "merge_dims", "from": "att", "axes": "static",
+      "class": "merge_dims", "from": "att", "axes": ["dim:%i" % n_heads, "dim:%i" % n_value_dim_per_head],
       "is_output_layer": True},  # [B,T,V']
   }
   with make_scope() as session:
@@ -7124,7 +7132,9 @@ def _build_self_attention_layer(d, input, output, inside_rec_layer, query_axis,
     'class': 'dot', 'from': [output + '_weights_drop', output + '_value_accum'],
     'red1': key_axis, 'red2': key_axis,
     "var1": None if inside_rec_layer else query_axis, "var2": "static:-1"}  # [B,n,T?,F|d_v]
-  d[output + '_att'] = {'class': 'merge_dims', 'axes': 'static', 'from': [output + '_output']}  # [B,T?,F|n*d_v]
+  d[output + '_att'] = {
+    'class': 'merge_dims', 'axes': ["dim:%i" % num_heads, "dim:%i" % value_dim],
+    'from': output + '_output'}  # [B,T?,F|n*d_v]
 
 
 def test_CumConcatLayer_self_attention_equal_to_SelfAttentionLayer():