Repository: rwth-i6/returnn

Issue:
Data.copy_compatible_to does not add feature dim
I have some tensor without a feature-dim, e.g. `[T]`, and want to copy it compatible to one with a feature dim, e.g. `[T,F]`, like in this test case here (`test_Data_copy_compatible_to_add_feature_dim`):

```python3
common_data = Data(name='energy', shape=(None, 3), batch_dim_axis=None)  # [T,F]
d1 = Data(name='mask', shape=(None,), batch_dim_axis=None, time_dim_axis=0, feature_dim_axis=None)  # [T]
d2 = d1.copy_compatible_to(common_data)
```

That however fails:
```
  File "/u/petrick/software/returnn/official/returnn/tf/util/data.py", line 1278, in copy_compatible_to
    line: assert data.feature_dim_axis is not None and v.feature_dim_axis is not None
    locals:
      data = <local> Data(name='energy', shape=(None, 3), batch_dim_axis=None, time_dim_axis=None, batch_shape_meta=[?,F|3])
      data.feature_dim_axis = <local> 1
      v = <local> Data(name='mask', shape=(1, None), batch_dim_axis=None, time_dim_axis=1, feature_dim_axis=None, batch_shape_meta=[1,T|?])
      v.feature_dim_axis = <local> None
```

I tried figuring out what went wrong in `Data.copy_compatible_to`, but I can't figure it out.
It would be great if somebody could take a look!

I pushed three test cases to 03416cdc3f0fe1b34bd20506ea0d88456b69528a. Note that `test_Data_copy_compatible_to_add_time_dim` already passes, this issue only occurs with the feature dim but not when adding a time dim.


Discussion:

Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -100,9 +100,7 @@ def concat_sources(src_layers):
     return network.concat_sources_dropout_cache[cache_key].copy()
   data = get_concat_sources_data_template(src_layers)
   # Currently we assume that get_concat_sources_data_template will match Data.get_common_data (besides the dim).
-  data_dyn_shape = []
-  common_source = Data.get_common_data(
-    [s.output for s in src_layers], warnings_out=log.v4, out_shape=data_dyn_shape)
+  common_source = Data.get_common_data([s.output for s in src_layers], warnings_out=log.v4)
   data.size_placeholder = common_source.size_placeholder.copy()  # to get right dimension tags
   layers_data = []
   with _name_scope_for_concat_src_layers(src_layers, "concat_sources"):
@@ -110,8 +108,7 @@ def concat_sources(src_layers):
       assert not layer.output.sparse, "sparse concat not supported"
       assert layer.output.dtype == data.dtype, "incompatible dtype with layer %r" % layer
       # unbroadcast is needed for tf.concat.
-      layers_data.append(layer.output.copy_compatible_to(
-        data, unbroadcast=True, except_feature=True, data_dyn_shape=data_dyn_shape))
+      layers_data.append(layer.output.copy_compatible_to(data, unbroadcast=True, except_feature=True))
     data.placeholder = tf.concat(
       axis=data.feature_dim_axis,
       values=[layer_data.placeholder for layer_data in layers_data])
@@ -1282,7 +1279,7 @@ def __init__(self, position, position_axis, output_dim_via_time_from, filter_inv
     pos_shape = [position.output.get_dim(i) for i in range(pos_ndim)]
     output_dim = output_dim_via_time_from.output.time_dimension()
     input_shape = pos_shape + [self.input_data.get_dim(i) for i in input_extra_axes]
-    input_expanded = self.input_data.copy_compatible_to(common, unbroadcast=True, data_dyn_shape=input_shape)
+    input_expanded = self.input_data.copy_compatible_to(common, unbroadcast=True)
     input_v = input_expanded.placeholder
     if filter_invalid_indices:
       mask = tf.logical_or(tf.less(pos_v, 0), tf.greater_equal(pos_v, output_dim))
@@ -4738,25 +4735,23 @@ def __init__(self, axis=None, **kwargs):
       i.e. some reasonable default for a new spatial axis.
     """
     super(StackLayer, self).__init__(**kwargs)
-    data_dyn_shape = []
-    axis_, common_source = self._get_axis_and_common(self.sources, data_dyn_shape=data_dyn_shape)
+    axis_, common_source = self._get_axis_and_common(self.sources)
     if axis is None:
       axis = axis_
     assert self.output.batch_shape[axis] == len(self.sources)
     sources_ = [
-      src.output.copy_compatible_to(common_source, unbroadcast=True, data_dyn_shape=data_dyn_shape)
+      src.output.copy_compatible_to(common_source, unbroadcast=True)
       for src in self.sources]
     self.output.placeholder = tf.stack([src.placeholder for src in sources_], axis=axis)
 
   @classmethod
-  def _get_axis_and_common(cls, sources, data_dyn_shape=None):
+  def _get_axis_and_common(cls, sources):
     """
     :param list[LayerBase] sources:
-    :param list[int]|None data_dyn_shape:
     :rtype: (int,Data)
     """
     from returnn.tf.util.basic import DimensionTag
-    common_source = Data.get_common_data([src.output for src in sources], out_shape=data_dyn_shape).copy_template()
+    common_source = Data.get_common_data([src.output for src in sources]).copy_template()
     tag = DimensionTag(kind=DimensionTag.Types.Spatial, dimension=1)
     return common_source.get_default_new_axis_for_dim_tag(tag), common_source
 
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
@@ -167,7 +167,7 @@ def is_equal(self, other, ignore_feature_dim=False, allow_same_feature_dim=False
           return True
         if broadcast_matches and (self.dimension == 1 or other.dimension == 1):
           return True
-      if unknown_spatial_matches and ((self.dyn_size is None) != (other.dyn_size is None)):
+      if unknown_spatial_matches and ((self.dyn_size is None) or (other.dyn_size is None)):
         return True
     if self.description == other.description:
       return True
@@ -824,48 +824,35 @@ def copy_with_time_dim_axis(self, time_dim_axis):
     assert self.time_dim_axis is not None
     return self.copy_move_axis(self.time_dim_axis, time_dim_axis)
 
-  def copy_move_axis(self, old_axis, new_axis):
+  def copy_transpose(self, perm):
     """
-    :param int old_axis: counted with batch-dim
-    :param int new_axis: counted with batch-dim
-    :return: copy of myself with moved axis (see :func:`move_axis`)
+    :param list[int] perm: permutation of the axes, counted with batch-dim.
+      Maps the new axes to the old axes
+    :return: copy of myself with permuted axes
     :rtype: Data
     """
-    from .basic import move_axis
-    if old_axis < 0:
-      old_axis += self.batch_ndim
-      assert old_axis >= 0
-    assert 0 <= old_axis < self.batch_ndim
-    if new_axis < 0:
-      new_axis += self.batch_ndim
-      assert new_axis >= 0
-    assert 0 <= new_axis < self.batch_ndim
-    if old_axis == new_axis:
+    assert len(perm) == self.batch_ndim
+    assert set(perm) == set(range(self.batch_ndim))
+    if all(perm[axis] == axis for axis in range(self.batch_ndim)):
       return self.copy()
+    inv_perm_ = {j: i for (i, j) in enumerate(perm)}
+    inv_perm = [j for (i, j) in sorted(inv_perm_.items())]
 
     def translate_axis(axis):
       """
-      :param int|None axis:
-      :return: axis after move_axis
+      :param int|None axis: counted with batch-dim
+      :return: translated axis (if not None)
       :rtype: int|None
       """
       if axis is None:
         return None
-      if old_axis == new_axis:
-        return axis
-      if axis < min(old_axis, new_axis) or axis > max(old_axis, new_axis):
-        return axis
-      if axis == old_axis:
-        return new_axis
-      if old_axis < new_axis:
-        assert old_axis < axis <= new_axis
-        return axis - 1
-      assert new_axis <= axis < old_axis
-      return axis + 1
+      return inv_perm[axis]
 
     data = self.copy()
-    if data.placeholder is not None:
-      data.placeholder = move_axis(data.placeholder, old_axis, new_axis)
+    if self.placeholder is not None:
+      from returnn.tf.util.basic import get_valid_scope_name_from_str
+      data.placeholder = tf.transpose(
+        self.placeholder, perm, name="%s_transpose" % get_valid_scope_name_from_str(self.name))
     data.batch_dim_axis = translate_axis(self.batch_dim_axis)
     new_feature_dim_axis = translate_axis(self.feature_dim_axis)
     if new_feature_dim_axis != data.feature_dim_axis:
@@ -884,6 +871,50 @@ def translate_axis(axis):
     data.sanity_check()
     return data
 
+  def copy_move_axis(self, old_axis, new_axis):
+    """
+    :param int old_axis: counted with batch-dim
+    :param int new_axis: counted with batch-dim
+    :return: copy of myself with moved axis (see :func:`move_axis`)
+    :rtype: Data
+    """
+    if old_axis < 0:
+      old_axis += self.batch_ndim
+      assert old_axis >= 0
+    assert 0 <= old_axis < self.batch_ndim
+    if new_axis < 0:
+      new_axis += self.batch_ndim
+      assert new_axis >= 0
+    assert 0 <= new_axis < self.batch_ndim
+    if old_axis == new_axis:
+      return self.copy()
+
+    perm = list(range(self.batch_ndim))
+    old = perm.pop(old_axis)
+    perm.insert(new_axis, old)
+    return self.copy_transpose(perm)
+
+  def copy_swap_axes(self, axis1, axis2):
+    """
+    Like :func:`Data.copy_move_axis`, but keeps all other axes unchanged.
+    :param int axis1: counted with batch-dim
+    :param int axis2: counted with batch-dim
+    :return: copy of myself with moved axis (see :func:`swapaxes`)
+    :rtype: Data
+    """
+    if axis1 < 0:
+      axis1 += self.batch_ndim
+    assert 0 <= axis1 < self.batch_ndim
+    if axis2 < 0:
+      axis2 += self.batch_ndim
+    assert 0 <= axis2 < self.batch_ndim
+    if axis1 == axis2:
+      return self.copy()
+
+    perm = list(range(self.batch_ndim))
+    perm[axis1], perm[axis2] = perm[axis2], perm[axis1]
+    return self.copy_transpose(perm)
+
   def copy_as_bt_or_tb_major(self):
     """
     :rtype: Data
@@ -1091,22 +1122,47 @@ def get_default_new_axis_for_dim_tag(self, dim_tag):
   def copy_add_dim_by_tag(self, dim_tag, unbroadcast=False, axis=None):
     """
     :param DimensionTag dim_tag:
-    :param bool unbroadcast:
+    :param bool unbroadcast: If True unbroadcast the newly added axis.
+      If `dim_tag.src_data` has a placeholder, will use the shape from there.
+      Otherwise, uses `dim_tag.dimension` (if static) or `dim_tag.dyn_size` (if dynamic) to infer shape.
     :param int|None axis:
     :rtype: Data
     """
     if axis is None:
       axis = self.get_default_new_axis_for_dim_tag(dim_tag=dim_tag)
+
+    def maybe_unbroadcast_placeholder():
+      """
+      If a placeholder is set, unbroadcasts the newly added axis
+      """
+      if res.placeholder is None:
+        return
+      assert res.batch_shape[axis] == 1
+      with tf.name_scope("copy_add_dim_by_tag_unbroadcast"):
+        dyn_size = None
+        if dim_tag.src_data is not None and dim_tag.src_axis is not None and dim_tag.src_data.placeholder is not None:
+          dyn_size = tf.shape(dim_tag.src_data.placeholder)[dim_tag.src_axis]
+        if dyn_size is None:
+          dyn_size = dim_tag.dimension
+        if dyn_size is None and dim_tag.dyn_size is not None:
+          dyn_size = tf.math.reduce_max(dim_tag.dyn_size)
+        if dyn_size is None and dim_tag.kind == DimensionTag.Types.Batch:
+          from returnn.tf.layers.base import LayerBase
+          dyn_size = LayerBase.get_recent_layer().get_batch_dim()
+        assert dyn_size is not None, 'need placeholder, dim_tag.dimension or dim_tag.dyn_size for unbroadcast'
+        tiles = [1] * axis + [dyn_size] + [1] * (res.batch_ndim - axis - 1)
+        res.placeholder = tf.tile(res.placeholder, tiles)
+
     if dim_tag.kind == DimensionTag.Types.Batch:
       res = self.copy_add_batch_dim(batch_dim_axis=axis)
       if unbroadcast:
-        assert res.placeholder is None  # not implemented yet...
+        maybe_unbroadcast_placeholder()
       return res
     # Note: if dim_tag is feature, but we are sparse, we just treat is as spatial, handled below.
     if dim_tag.kind == DimensionTag.Types.Feature and not self.sparse:
       res = self.copy_add_feature_dim(axis=axis)
       if unbroadcast:
-        assert res.placeholder is None  # not implemented yet...
+        maybe_unbroadcast_placeholder()
         res.dim = dim_tag.dimension
         shape = list(res.shape)
         shape[res.get_batch_axis_excluding_batch(res.feature_dim_axis)] = dim_tag.dimension
@@ -1117,7 +1173,7 @@ def copy_add_dim_by_tag(self, dim_tag, unbroadcast=False, axis=None):
     res = self.copy_add_spatial_dim(spatial_dim_axis=axis, dim=1)
     assert res.batch_shape[axis] == 1
     if unbroadcast:
-      assert res.placeholder is None  # not implemented yet...
+      maybe_unbroadcast_placeholder()
       shape = list(res.shape)
       shape[res.get_batch_axis_excluding_batch(axis)] = dim_tag.dimension
       res.shape = tuple(shape)
@@ -1165,16 +1221,13 @@ def copy_split_feature_dim(self, new_feature_dim):
     v.sanity_check()
     return v
 
-  def copy_compatible_to(self, data, unbroadcast=False, except_feature=False,
-                         data_dyn_shape=None, check_sparse=True, check_dtype=True):
+  def copy_compatible_to(self, data, unbroadcast=False, except_feature=False, check_sparse=True, check_dtype=True):
     """
     :param Data data: other data which the returned tensor should be compatible to
       It would add any missing axes with a dim 1 axis for automatic broadcasting.
       It currently does not check whether existing dims match.
     :param bool unbroadcast: if True, all broadcast axes (axes with dim 1) will be tiled such that they match
     :param bool except_feature: if unbroadcast, do not unbroadcast the feature dim
-    :param tf.Tensor|list[tf.Tensor|int]|tuple[tf.Tensor|int]|None data_dyn_shape:
-      For unbroadcast, if we do not want to rely on tf.shape(data.placeholder).
     :param bool check_sparse:
     :param bool check_dtype:
     :returns: Data, might add broadcast dimensions
@@ -1190,123 +1243,54 @@ def copy_compatible_to(self, data, unbroadcast=False, except_feature=False,
         v.dim = v.batch_shape[v.feature_dim_axis]
       else:
         v.dim = None
-    if data.batch_dim_axis is not None and v.batch_dim_axis is None:
-      v = v.copy_add_batch_dim(0)  # later we might move the axis
     if v.batch_dim_axis is not None and data.batch_dim_axis is None:
       raise ValueError("copy_compatible_to: self %r has batch-dim, but target data %r has not" % (self, data))
     if data.batch_ndim < v.batch_ndim:
       raise ValueError("copy_compatible_to: self %r already has more dims than target data %r" % (self, data))
-    start = v
-    _, dim_tags = DimensionTag.get_all_dimension_tags([start, data], dict(allow_same_feature_dim=True))
-    assert len(dim_tags[start]) == start.batch_ndim
-    assert len(dim_tags[data]) == data.batch_ndim
-    # This sets it explicitly. We will later make it NotSpecified if needed.
-    # This avoids unexpected behavior after copy_add_spatial_dim and simplifies the logic.
-    v.feature_dim_axis = v.feature_dim_axis
-    # Add dims, in case we miss any.
-    for axis in range(data.batch_ndim):
-      if axis == data.batch_dim_axis:
-        continue
-      axis_wo_batch = data.get_batch_axis_excluding_batch(axis)
-      v_axis = v.get_batch_axis(axis_wo_batch)
-      existing_axis = None
-      if dim_tags[data][axis] in dim_tags[start]:
-        existing_axis = dim_tags[start].index(dim_tags[data][axis])
-      if existing_axis is None:
-        # Try a bit harder to find an existing.
-        if axis == data.feature_dim_axis and v.feature_dim_axis is not None:
-          if v.batch_shape[v.feature_dim_axis] == data.batch_shape[axis]:
-            if v.batch_shape[v.feature_dim_axis] is not None:
-              existing_axis = v.feature_dim_axis  # There might be cases that the dim_tags did not match.
-          if v.batch_shape[v.feature_dim_axis] == 1:
-            existing_axis = v.feature_dim_axis  # Interpret the existing as broadcast dim.
-        if axis == data.time_dim_axis and v.time_dim_axis is not None:
-          if v.batch_shape[v.time_dim_axis] == data.batch_shape[axis]:
-            if v.batch_shape[v.time_dim_axis] is not None:
-              existing_axis = v.time_dim_axis  # There might be cases that the dim_tags did not match.
-          if v.batch_shape[v.time_dim_axis] == 1:
-            existing_axis = v.time_dim_axis  # Interpret the existing as broadcast dim.
-      if existing_axis is not None:
-        # We go from left to right, so we should have moved it already.
-        # However, it could be that we confused some other axis earlier.
-        if existing_axis > v_axis:
-          v = v.copy_move_axis(old_axis=existing_axis, new_axis=v_axis)
-          dim_tags[start].insert(v_axis, dim_tags[start].pop(existing_axis))  # keep consistent
-        continue
-      if data.batch_ndim > v.batch_ndim:
-        if axis == data.feature_dim_axis:
-          v = v.copy_add_feature_dim(v_axis)
-        else:
-          v = v.copy_add_spatial_dim(v_axis, auto_time_dim_axis=False)  # time-dim would be set later
-        dim_tags[start].insert(v_axis, v.get_dim_tag(v_axis))  # keep consistent
-        if axis == data.time_dim_axis and v.time_dim_axis != v_axis:
-          v.time_dim_axis = v_axis
-        if axis == data.feature_dim_axis and v.feature_dim_axis != v_axis:
-          v.feature_dim_axis = v_axis
-    # Now we assume that we have all missing axes added,
-    # but they might still be in a wrong order.
+
+    is_equal_opts = dict(
+      allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True, ignore_feature_dim=True)
+    mapped_axes = data.find_matching_dim_map(v, list(range(v.batch_ndim)), is_equal_opts)  # maps v -> data
+    assert len(mapped_axes) == v.batch_ndim
+
+    for target_axis in range(data.batch_ndim):
+      new_v_axis = min(target_axis, v.batch_ndim)
+      if target_axis not in mapped_axes.values():
+        # Dim in data, but not in v
+        unbroadcast_axis = unbroadcast and not (except_feature and data.feature_dim_axis == target_axis)
+        v = v.copy_add_dim_by_tag(data.get_dim_tag(target_axis), axis=new_v_axis, unbroadcast=unbroadcast_axis)
+        # Keep mapped_axes consistent
+        mapped_axes = {v_ax + (1 if v_ax >= new_v_axis else 0): trg_ax for v_ax, trg_ax in mapped_axes.items()}
+        mapped_axes[new_v_axis] = target_axis
+      else:
+        # Dim exists in both data and in v. Maybe order is wrong.
+        matching_v_axes = [v_ax for v_ax, trg_ax in mapped_axes.items() if trg_ax == target_axis]
+        assert len(matching_v_axes) == 1
+        matching_v_axis = matching_v_axes[0]
+        if target_axis != matching_v_axis:
+          # Order was wrong
+          v = v.copy_swap_axes(matching_v_axis, new_v_axis)
+          # Keep mapped_axes consistent
+          mapped_axes[matching_v_axis], mapped_axes[new_v_axis] = mapped_axes[new_v_axis], mapped_axes[matching_v_axis]
+
     assert v.batch_ndim == data.batch_ndim
-    # Now maybe move batch/feature axis.
-    # We might do multiple iterations here, depending on which axis comes first.
-    # This is a bit ugly, but the code is simpler.
-    num_iterations = 0
-    while True:
-      num_iterations += 1
-      assert num_iterations <= 4
-      if v.batch_dim_axis != data.batch_dim_axis:
-        assert data.batch_dim_axis is not None and v.batch_dim_axis is not None
-        v = v.copy_with_batch_dim_axis(data.batch_dim_axis)
-        assert v.batch_dim_axis == data.batch_dim_axis
-        continue
-      if v.feature_dim_axis != data.feature_dim_axis:
-        assert data.feature_dim_axis is not None and v.feature_dim_axis is not None
-        v = v.copy_with_feature_dim_axis(data.feature_dim_axis)
-        assert v.feature_dim_axis == data.feature_dim_axis
-        if data.feature_dim_axis_or_unspecified is NotSpecified:
-          v.feature_dim_axis = NotSpecified
-          assert v.feature_dim_axis == data.feature_dim_axis
-        continue
-      # Now we have both equal.
-      break
-    if data.feature_dim_axis_or_unspecified is NotSpecified and v.feature_dim_axis_or_unspecified is not NotSpecified:
-      if v._default_feature_dim_axis() == v.feature_dim_axis:
-        v.feature_dim_axis = NotSpecified
+    assert all(mapped_axes[ax] == ax for ax in range(v.batch_ndim))
+
+    # Ensure time_dim_axis and feature_dim_axis is same as in data
+    assert v.batch_dim_axis == data.batch_dim_axis  # there is only at most one batch_dim_axis
+    v.time_dim_axis = data.time_dim_axis
+    v.feature_dim_axis = data.feature_dim_axis_or_unspecified
+    if v.feature_dim_axis is not None:
+      v.dim = v.batch_shape[v.feature_dim_axis]
+    else:
+      v.dim = None
+
+    # Reset sparse
     if self.sparse:
       v.feature_dim_axis = NotSpecified
       v.sparse = True  # reset
       v.dim = self.dim  # reset
-    if unbroadcast and any([d1 != 1 and d2 == 1 for (d1, d2) in zip(data.batch_shape, v.batch_shape)]):
-      v.size_placeholder.update(data.size_placeholder or {})
-      if v.placeholder is not None:
-        with tf.name_scope("copy_compatible_to_unbroadcast"):
-          tiles = [1] * v.batch_ndim
-          for axis in range(v.batch_ndim):
-            if v.batch_shape[axis] != 1:
-              continue
-            if except_feature and axis == v.feature_dim_axis:
-              continue
-            if data.batch_shape[axis] is not None:
-              tiles[axis] = data.batch_shape[axis]
-            elif data_dyn_shape is not None:
-              tiles[axis] = data_dyn_shape[axis]
-            else:
-              assert data.placeholder, "need data.placeholder for unbroadcast (target data: %r)" % v
-              tiles[axis] = tf.shape(data.placeholder)[axis]
-          if set(tiles) != {1}:
-            v.placeholder = tf.tile(v.placeholder, tiles)
-      new_shape = list(v.batch_shape)
-      for axis in range(v.batch_ndim):
-        if except_feature and axis == data.feature_dim_axis:
-          continue
-        if data.batch_shape[axis] != 1 and new_shape[axis] == 1:
-          new_shape[axis] = data.batch_shape[axis]
-      if v.feature_dim_axis is not None:
-        v.dim = new_shape[v.feature_dim_axis]
-      if v.batch_dim_axis is not None:
-        del new_shape[v.batch_dim_axis]
-      v.shape = tuple(new_shape)
-      if v.placeholder is not None and not except_feature:
-        v.placeholder.set_shape(v.batch_shape)
+
     v.sanity_check()
     return v
 
@@ -2429,29 +2413,23 @@ def get_batch_shape_dim_tags(self):
     return tuple([self.get_dim_tag(i) for i in range(self.batch_ndim)])
 
   @classmethod
-  def get_common_data(cls, sources, warnings_out=None, out_shape=None):
+  def get_common_data(cls, sources, warnings_out=None):
     """
     :param list[Data] sources:
     :param io.TextIOBase|io.StringIO|typing.TextIO|None warnings_out:
-    :param list[int|tf.Tensor]|None out_shape: will insert the shape dynamically
     :return: some generic data where the sources should be compatible to (with copy_compatible_to),
       i.e. it contains the union of all axes from all sources (least common multiple).
     :rtype: Data|None
     """
-    assert not out_shape
     if not sources:
       return None
     assert sources
     if len(sources) == 1:
-      if out_shape is not None:
-        out_shape.extend(sources[0].get_dynamic_batch_shape())
       return sources[0]
     max_ndim = max([s.batch_ndim for s in sources])
     # Try with the (first) largest.
     common = [s for s in sources if s.batch_ndim == max_ndim][0]
     common = common.copy()
-    if out_shape is not None:
-      out_shape.extend(common.get_dynamic_batch_shape())
     if any([s.beam for s in sources]):
       # Note: we don't use copy_extend_with_beam because we don't want to create any ops in the TF graph at this point.
       common.beam = SearchBeam.get_combined_beam(*[s.beam for s in sources])
@@ -2485,12 +2463,6 @@ def get_common_data(cls, sources, warnings_out=None, out_shape=None):
       return common
     # Ok, there is some other axis (or multiple), or we cannot identify/compare them because of incomplete information.
     # Try something more complex: Make all axes unique.
-    # Note that this should also work at template construction time,
-    # where we do not have access to the size_placeholder,
-    # and thus the dimension tags are not reliable (in the current implementation).
-    tags_dict_ext = {
-      id(tag): [(data, tags_dict[data].index(tag)) for data in sources if tag in tags_dict[data]]
-      for tag in all_dim_tags}
     for dim_tag in all_dim_tags:
       if not dim_tag.can_compare():
         if warnings_out:
@@ -2504,10 +2476,6 @@ def get_common_data(cls, sources, warnings_out=None, out_shape=None):
         largest_dim_tags.append(dim_tag)
         axis = common.get_default_new_axis_for_dim_tag(dim_tag)
         common = common.copy_template().copy_add_dim_by_tag(dim_tag, unbroadcast=True, axis=axis)
-        if out_shape is not None:
-          tag_data, tag_data_axis = tags_dict_ext[id(dim_tag)][0]
-          assert isinstance(tag_data, Data)
-          out_shape.insert(axis, tag_data.get_dim(tag_data_axis))
     # Simple fallback: Use first with biggest batch_ndim.
     # Was even simpler before: Use first.
     return common
@@ -2522,15 +2490,20 @@ def find_matching_dims(self, dim_tag, is_equal_opts):
     """
     return [axis for axis in range(self.batch_ndim) if self.get_dim_tag(axis).is_equal(dim_tag, **is_equal_opts)]
 
-  def find_matching_dim_map(self, other, other_axes):
+  def find_matching_dim_map(self, other, other_axes, is_equal_opts=None):
     """
     Looks up all other_axes of another Data in this Data. Does not allow duplicates.
 
     :param Data other:
     :param list[int] other_axes: a list of axes of ``other``, counted with batch dim
     :return: a dict mapping other axes to own axes, all counted with batch dim
+    :param dict[str,bool]|None is_equal_opts: passed to DimensionTag.is_equal
     :rtype: dict[int,int]
     """
+    if is_equal_opts is None:
+      is_equal_opts = dict(
+        allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True)
+
     def map_other_axis_to_self(other_axis, taken_self_axes):
       """
       :param int other_axis: counted with batch dim
@@ -2538,20 +2511,26 @@ def map_other_axis_to_self(other_axis, taken_self_axes):
       :return: the axis of ``self`` that matches ``other_axis``, counted with batch dim
       :rtype: int
       """
-      is_equal_opts = dict(
-        allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True)
       other_axis_dim_tag = other.get_dim_tag(other_axis)
       matching = [
         self_axis for self_axis in self.find_matching_dims(other_axis_dim_tag, is_equal_opts)
         if self_axis not in taken_self_axes]
       if not matching:
-        # The DimensionTags do not match. Then we also allow one single dyn_size to be unknown
+        # Try harder by allowing broadcasting to match
+        is_equal_opts["broadcast_matches"] = True
+        matching = [
+          self_axis for self_axis in self.find_matching_dims(other_axis_dim_tag, is_equal_opts)
+          if self_axis not in taken_self_axes]
+      if not matching:
+        # If still not, then also allow one single dyn_size to be unknown
         is_equal_opts["unknown_spatial_matches"] = True
         matching = [
           self_axis for self_axis in self.find_matching_dims(other_axis_dim_tag, is_equal_opts)
           if self_axis not in taken_self_axes]
-        assert len(matching) == 1, 'cannot match the axes %s from %s to %s' % (other_axes, other, self)
-      assert matching, 'cannot match the axes %s from %s to %s' % (other_axes, other, self)
+        assert len(matching) == 1, 'cannot match the axes %s from %s to %s. Failing to match axis %s' % (
+          other_axes, other, self, other_axis)
+      assert matching, 'cannot match the axes %s from %s to %s. Failing at axis %s' % (
+        other_axes, other, self, other_axis)
       # If there are multiple matches (e.g. because two axes have the same feature dim), leave their order intact.
       # We do this by always choosing the first unused match which is the smallest axes
       return matching[0]


Test patch:
diff --git a/tests/test_TFUtil.py b/tests/test_TFUtil.py
--- a/tests/test_TFUtil.py
+++ b/tests/test_TFUtil.py
@@ -293,6 +293,59 @@ def test_Data_copy_compatible_to_time_major():
   assert d2a.feature_dim_axis == d1.feature_dim_axis
 
 
+def test_Data_find_matching_dim_map_different_static_dims():
+  d1 = Data(name='p1_output', shape=(5, 5, 3), batch_dim_axis=None, time_dim_axis=None)  # [5,5,F|3]
+  d2 = Data(name='p2_output', shape=(5, 1, 1), batch_dim_axis=None, time_dim_axis=None)  # [5,1,F|1]
+
+  # without broadcast_matches=False should not match
+  is_equal_opts = dict(allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True)
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(0), is_equal_opts=is_equal_opts), [0,1])
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(1), is_equal_opts=is_equal_opts), [])
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(2), is_equal_opts=is_equal_opts), [])
+
+  # with broadcast_matches=True should match
+  is_equal_opts = dict(
+    allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True,
+    broadcast_matches=True)
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(0), is_equal_opts=is_equal_opts), [0,1])
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(1), is_equal_opts=is_equal_opts), [0,1,2])
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(2), is_equal_opts=is_equal_opts), [0,1,2])
+
+  mapping = d1.find_matching_dim_map(d2, list(range(d2.batch_ndim)))  # maps d2 -> d1
+  assert len(mapping.values()) == d2.batch_ndim
+  assert all(mapping[i] == i for i in range(d2.batch_ndim))
+
+  d2_compatible = d2.copy_compatible_to(d1)
+  assert d2_compatible.batch_shape == d2.batch_shape
+  d1_compatible = d1.copy_compatible_to(d2)
+  assert d1_compatible.batch_shape == d1.batch_shape
+
+
+def test_Data_find_matching_dim_map_broadcast_matches():
+  d1 = Data(name='d1', shape=(5, None), time_dim_axis=2)  # [B,F|5,T]
+  d2 = Data(name='d2', shape=(5, 1), batch_dim_axis=None, time_dim_axis=None, feature_dim_axis=0)  # [F|5,1]
+  print('d1:', d1)
+  print('d2:', d2)
+
+  # default should not match
+  is_equal_opts = dict(allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True)
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(0), is_equal_opts=is_equal_opts), [1])
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(1), is_equal_opts=is_equal_opts), [])
+
+  # with broadcast_matches=True should match
+  is_equal_opts_match = dict(
+    allow_same_feature_dim=True, allow_same_spatial_dim=True, treat_feature_as_spatial=True, broadcast_matches=True)
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(0), is_equal_opts=is_equal_opts_match), [1])
+  assert_equal(d1.find_matching_dims(d2.get_dim_tag(1), is_equal_opts=is_equal_opts_match), [1,2])
+
+  mapping = d1.find_matching_dim_map(d2, list(range(d2.batch_ndim)), is_equal_opts)  # maps d2 -> d1
+  assert mapping[0] == 1 and mapping[1] == 2
+
+  copied = d2.copy_compatible_to(d1)
+  assert copied.batch_ndim == d1.batch_ndim and copied.batch_shape == (None, 5, 1)
+  print('copied compatible:', copied)
+
+
 def test_Data_sparse_int32_with_dim_kwargs_init():
   data = Data(name="classes_with_dim", shape=(None,), dim=10, sparse=True, dtype="int32")
   assert data.sparse and data.have_time_axis() and data.shape == (None,) and data.dim == 10
@@ -442,6 +495,27 @@ def test_Data_copy_compatible_to_add_batch_dim():
   assert d2.batch_dim_axis is not None and d2.batch_shape == (None, 1, 1)
 
 
+def test_Data_copy_compatible_to_add_feature_dim():
+  common_data = Data(name='energy', shape=(None, 3), time_dim_axis=0, batch_dim_axis=None)  # [T,F]
+  d1 = Data(name='mask', shape=(None,), batch_dim_axis=None, time_dim_axis=0, feature_dim_axis=None)  # [T]
+  d2 = d1.copy_compatible_to(common_data)
+  assert d2.batch_dim_axis is None and d2.time_dim_axis is not None and d2.batch_shape == (None, 1)
+
+
+def test_Data_copy_compatible_to_add_batch_feature_dim():
+  common_data = Data(name='energy', shape=(None, 3))  # [B,T,F]
+  d1 = Data(name='mask', shape=(None,), batch_dim_axis=None, time_dim_axis=0, feature_dim_axis=None)  # [T]
+  d2 = d1.copy_compatible_to(common_data)
+  assert d2.batch_dim_axis is not None and d2.time_dim_axis is not None and d2.batch_shape == (None, None, 1)
+
+
+def test_Data_copy_compatible_to_add_time_dim():
+  common_data = Data(name='energy', shape=(None, 3))  # [B,T,F]
+  d1 = Data(name='mask', shape=(3,))  # [B,F]
+  d2 = d1.copy_compatible_to(common_data)
+  assert d2.time_dim_axis is not None and d2.batch_shape == (None, 1, 3)
+
+
 def test_Data_copy_compatible_to_time_axis_at_end():
   data = Data(name='att_weights_output', shape=(1, None), time_dim_axis=2, feature_dim_axis=1)
   common_data = Data(name='accum_att_weights_output', shape=(None, 1))
@@ -1367,8 +1441,8 @@ def rel_embed(x, v, t):
     """
     import tensorflow as tf
     from returnn.tf.util.basic import nd_indices
-    v = v.copy_compatible_to(x)  # t_rel_var. (B, Ts, K)
-    assert v.dim == x.dim
+    v = v.copy_move_axis(v.batch_dim_axis, x.batch_dim_axis)  # t_rel_var. (B, Ts, K) or (Ts, B, K)
+    assert v.feature_dim_axis == x.feature_dim_axis and v.dim == x.dim
     t = t + 1  # shift by 1, because we init at -1
     # t = tf.Print(t, ["t:", t])
     time_dim = tf.shape(x.placeholder)[x.time_dim_axis]