Repository: rwth-i6/returnn

Issue:
`ConcatLayer` should be implemented
`CopyLayer` is a special case which concatenates in the feature dim.
However, following the discussion in #586, maybe that implicit dimension can be problematic in some cases.
Also, it is not generic to work for any other dim.

The newly proposed `ConcatLayer` should make this explicit.
For every input, it should specify also the axis to use.
The axis would be via `get_axis_from_description` (but disallowing ints), so `"F"` is still possible, but that behavior could change in the future (#586).

Proposed signature:
```
sources: List[Tuple[LayerBase, Union[str,DimensionTag]]]
```



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
@@ -307,6 +307,85 @@ def transform_config_dict(cls, d, network, get_layer):
       d["extra_deps"] = [get_layer(src_name) for src_name in extra_deps]
 
 
+class ConcatLayer(LayerBase):
+  """
+  Concatenates the inputs in specified axes.
+  This generalizes :class:`CopyLayer` which concatenates in the feature dim.
+  """
+  layer_class = "concat"
+
+  def __init__(self, sources, allow_broadcast=False, **kwargs):
+    """
+    :param list[(LayerBase,str|DimensionTag)] sources:
+    :param bool allow_broadcast:
+    """
+    if allow_broadcast:
+      raise NotImplementedError
+    sources, axes = zip(*sources)  # unzip
+    super(ConcatLayer, self).__init__(sources=sources, **kwargs)
+    sources_data = [layer.output for layer in sources]  # type: typing.List[Data]
+    axes_int = [src.get_axis_from_description(axis) for (src, axis) in zip(sources_data, axes)]
+    # Currently, assume that the output format matches the first input. See get_out_data_from_opts.
+    assert self.output.batch_ndim == sources_data[0].batch_ndim
+    out_concat_axis = axes_int[0]
+
+    def _copy_compatible(x, axis):
+      """
+      :param Data x: input
+      :param int axis:
+      :rtype: Data
+      """
+      dummy_ref = self.output.copy_template()
+      dummy_ref = dummy_ref.copy_template_replace_dim_tag(
+        axis=out_concat_axis, new_dim_tag=x.dim_tags[axis])
+      return x.copy_compatible_to(dummy_ref, add_dims=allow_broadcast, unbroadcast=False)
+
+    sources_data = [_copy_compatible(src, axis) for (src, axis) in zip(sources_data, axes_int)]
+    self.output.placeholder = tf_util.concat_with_opt_broadcast(
+      [src.placeholder for src in sources_data], axis=out_concat_axis,
+      allow_broadcast=[allow_broadcast] * len(sources_data))
+
+  @classmethod
+  def get_out_data_from_opts(cls, name, sources, allow_broadcast=False, **kwargs):
+    """
+    :param str name:
+    :param list[(LayerBase,str|DimensionTag)] sources:
+    :param bool allow_broadcast:
+    :rtype: Data
+    """
+    assert sources
+    sources, axes = zip(*sources)  # unzip
+    axes_int = [layer.output.get_axis_from_description(axis) for (layer, axis) in zip(sources, axes)]
+    concat_dim_tags = [
+      layer.output.dim_tags[axis] for (layer, axis) in zip(sources, axes_int)]  # type: typing.List[DimensionTag]
+    if any(tag.dimension is None for tag in concat_dim_tags):
+      dimension = None
+    else:
+      dimension = 0
+      for tag in concat_dim_tags:
+        dimension += tag.dimension
+    # We ignore allow_broadcast here... Anyway not currently implemented.
+    # Just overtake the first input format.
+    concat_res_dim_tag = DimensionTag(
+      kind=concat_dim_tags[0].kind, description="%s_concat" % name, dimension=dimension,
+      derived_from_tag=concat_dim_tags[0])
+    res_dim_tags = list(sources[0].output.dim_tags)
+    res_dim_tags[axes_int[0]] = concat_res_dim_tag
+    return Data(name="%s_output" % name, dim_tags=res_dim_tags, dtype=sources[0].output.dtype)
+
+  @classmethod
+  def transform_config_dict(cls, d, network, get_layer):
+    """
+    :param dict[str] d: will modify inplace
+    :param returnn.tf.network.TFNetwork network:
+    :param ((str) -> LayerBase) get_layer: function to get or construct another layer
+    """
+    sources_and_axes = d.pop("from")
+    d["from"], axes = zip(*sources_and_axes)
+    super(ConcatLayer, cls).transform_config_dict(d, network=network, get_layer=get_layer)
+    d["sources"] = list(zip(d["sources"], axes))
+
+
 class DropoutLayer(CopyLayer):
   """
   Just the same as :class:`CopyLayer`, because that one already supports dropout.
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
@@ -2337,12 +2337,14 @@ def copy_extend_batch(self, batch):
         data.placeholder = x
     return data
 
-  def copy_compatible_to(self, data, unbroadcast=False, except_feature=False, check_sparse=True, check_dtype=True):
+  def copy_compatible_to(self, data, add_dims=True, 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 would add any missing axes with a dim 1 axis for automatic broadcasting (with add_dims=True).
       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 add_dims: whether to add (broadcast, or unbroadcasted) dims. throws error if missing dim
+    :param bool unbroadcast: if True, all added 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 bool check_sparse:
     :param bool check_dtype:
@@ -2374,6 +2376,10 @@ def copy_compatible_to(self, data, unbroadcast=False, except_feature=False, chec
     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():
+        if not add_dims:
+          raise ValueError(
+            "%s.copy_compatible_to(%s) not allowed, axis %i (%s) not in source" % (
+              self, data, target_axis, data.dim_tags[target_axis]))
         # 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)


Test patch:
diff --git a/tests/test_TFNetworkLayer.py b/tests/test_TFNetworkLayer.py
--- a/tests/test_TFNetworkLayer.py
+++ b/tests/test_TFNetworkLayer.py
@@ -289,6 +289,23 @@ def test_concat_sources_dim1():
     session.run(out.output.placeholder, feed_dict=feed_dict)
 
 
+def test_ConcatLayer():
+  with make_scope() as session:
+    net_dict = {
+      "lin1": {"class": "linear", "activation": "sigmoid", "n_out": 5, "from": "data:data"},
+      "lin2": {"class": "linear", "activation": "sigmoid", "n_out": 3, "from": "data:data"},
+      "output": {"class": "concat", "from": [("lin1", "F"), ("lin2", "F"), ("data", "F")]},
+    }
+    config = Config({"extern_data": {"data": {"dim": 2}}})
+    network = TFNetwork(config=config)
+    network.construct_from_dict(net_dict)
+    out = network.get_default_output_layer()
+    assert_equal(out.output.shape, (None, 10))
+    feed_dict = make_feed_dict(network.extern_data, same_time=True)
+    session.run(tf_compat.v1.global_variables_initializer())
+    session.run(out.output.placeholder, feed_dict=feed_dict)
+
+
 def test_LinearLayer_batch_feature_major():
   with make_scope() as session:
     network = TFNetwork(config=Config(), extern_data=ExternData(), train_flag=True)