CumConcatLayer more (finished?)

Author: albertz

returnn/tf/layers/rec.py

@@ -8542,8 +8542,10 @@ def __init__(self, new_dim, **kwargs):
     :param DimensionTag new_dim:
     """
     super(CumConcatLayer, self).__init__(**kwargs)
-    assert self.network.is_inside_rec_layer()
+    rec_layer = self.network.get_rec_parent_layer(inside_loop=False)
+    assert rec_layer, "%r must be used inside a RecLayer" % self
     out_axis = self.output.get_axis_from_description(new_dim)
+    new_dim_ = self.output.dim_tags[out_axis]
 
     if self.network.is_inside_rec_layer(inside_loop=True):
       current_data = self.input_data.copy_compatible_to(self.output, unbroadcast=False)
@@ -8553,21 +8555,31 @@ def __init__(self, new_dim, **kwargs):
       self.rec_vars_outputs["state"] = concat_frames
       self.output.placeholder = concat_frames
 
-      dyn_size = tf.broadcast_to(self.network.get_rec_step_index() + 1, [data.get_batch_dim()])
+      if not new_dim_.dyn_size_ext:
+        # Unbroadcasting to [B] is not needed because any layers operating on this
+        # should be able to handle extended dyn sizes.
+        # Clipping it to the max length for sequences in the loop which are already ended
+        # (i.e. considering the end flag)
+        # is also not needed because any calculations after the end are irrelevant.
+        dyn_size = self.network.get_rec_step_index() + 1  # scalar
+        new_dim_.dyn_size_ext = Data(
+          name="%s:cum-concat:size-inside" % self.name,
+          dim_tags=[],  # scalar
+          placeholder=dyn_size)
 
     else:
-      # If not inside a rec loop, this layer is a no-op
-      self.output.placeholder = None  # TODO
-      data.size_placeholder = self.input_data.size_placeholder.copy()
-      dyn_size = tf.identity(data.get_dynamic_size(out_axis))
-
-    # We already set the size_placeholder to a dummy rec-history before, now do it properly
-    from returnn.tf.util.basic import DimensionTag
-    tag = DimensionTag(
-      description="rec-history:%s" % self.get_absolute_name(),
-      kind=DimensionTag.Types.Time)
-    data.size_placeholder[data.get_batch_axis_excluding_batch(out_axis)] = dyn_size
-    tag.set_tag_on_size_tensor(dyn_size)
+      # If not inside a rec loop, this layer is a no-op on the tensor.
+      self.output.placeholder = self.input_data.placeholder
+
+      # However, we used new dim tags, which were already prepared.
+      # We now must fill in the extended dynamic size information.
+      if not new_dim_.dyn_size_ext:
+        # This must match the logic above for inside the loop.
+        dyn_size = tf.range(tf.math.reduce_max(rec_layer.time_dim_tag.dyn_size)) + 1  # [T]
+        new_dim_.dyn_size_ext = Data(
+          name="%s:cum-concat:size-outside" % self.name,
+          dim_tags=[rec_layer.time_dim_tag],
+          placeholder=dyn_size)
 
   @classmethod
   def get_out_data_from_opts(cls, name, network, sources, new_dim, **kwargs):