Repository: rwth-i6/returnn

Issue:
Way to get final last recurrent (sub) layer output
Either we excend `GetLastHiddenStateLayer` or (maybe cleaner) make a separate layer for that.

We also need to extend `_SubnetworkRecCell`. It is quite trivial for those layers which end up in `prev_layers_needed` and which are not moved out. Otherwise we need to think a bit how to do it.

Or maybe there could be some special layer for inside the loop, which would disable the accumulation logic. Some flag like `rec_no_accum`. And then users can simply use `CopyLayer` and also set `is_output_layer` and then access it from outside as usual.

Not sure.

But this is needed.



Discussion:

Resolving patch:
diff --git a/returnn/tf/layers/base.py b/returnn/tf/layers/base.py
--- a/returnn/tf/layers/base.py
+++ b/returnn/tf/layers/base.py
@@ -68,6 +68,7 @@ def __init__(self, name, network, output, n_out=NotSpecified, out_type=None, sou
                updater_opts=None,
                initial_output=None,
                state=None,
+               need_last=False,
                rec_previous_layer=None,
                collocate_with=None,
                trainable=None,
@@ -112,13 +113,18 @@ def __init__(self, name, network, output, n_out=NotSpecified, out_type=None, sou
     :param float|None spatial_smoothing: see :func:`returnn.tf.util.basic.spatial_smoothing_energy`
     :param float|None param_variational_noise: adds variational noise to the params during training
     :param dict[str]|None updater_opts: accepts similar opts as TFUpdater, e.g. "optimizer", "learning_rate", ...
-    :param bool|None is_output_layer:
+    :param bool|None is_output_layer: triggers the construction of this layer in the root net.
+      Inside a :class:`RecLayer`, it triggers the explicit accumulation of all frames.
+      Also see the ``need_last`` option.
     :param bool only_on_eval: if True, this layer will only be calculated in eval
     :param bool only_on_search: if True, this layer will only be calculated when search is done
     :param int|None copy_output_loss_from_source_idx: if set, will copy output_loss from this source
     :param bool|dict batch_norm: see self.batch_norm()
     :param str|float initial_output: used for recurrent layer, see self.get_rec_initial_output()
     :param state: explicitly defines the rec state. initial_state would define the initial state (in the first frame)
+    :param bool need_last: Inside :class:`RecLayer`, make sure that we can access the last frame.
+      Similar to ``is_output_layer, but this is specifically about the last frame,
+      i.e. it does not trigger accumulation.
     :param LayerBase|None rec_previous_layer: via the recurrent layer, layer (template) which represents the past of us.
       You would not explicitly set this in a config. This is automatically, internally, via :class:`RecLayer`.
     :param list[str]|None collocate_with: in the rec layer, collocate with the specified other layers
@@ -166,6 +172,7 @@ def __init__(self, name, network, output, n_out=NotSpecified, out_type=None, sou
     self.search_choices = None  # type: typing.Optional[SearchChoices]
     self._src_common_search_choices = _src_common_search_choices
     self._initial_output = initial_output
+    self.need_last = need_last
     self._rec_previous_layer = rec_previous_layer
     self.collocate_with = collocate_with or []
     self.post_init_hooks = []  # list of functions
diff --git a/returnn/tf/layers/rec.py b/returnn/tf/layers/rec.py
--- a/returnn/tf/layers/rec.py
+++ b/returnn/tf/layers/rec.py
@@ -1038,6 +1038,7 @@ def __init__(self, net_dict, source_data, time_dim_tag, rec_layer_name, parent_n
     self.prev_layer_templates = {}  # type: typing.Dict[str,_TemplateLayer]
     self._template_construction_exceptions = None  # type: typing.Optional[typing.List[str]]
     self._construct_template(parent_get_layer=parent_get_layer)
+    self._last_frames = {}  # type: typing.Dict[str,Data]
     self._initial_outputs = None  # type: typing.Optional[typing.Dict[str,tf.Tensor]]
     self._initial_extra_outputs = None  # type: typing.Optional[typing.Dict[str,typing.Dict[str,typing.Union[tf.Tensor,typing.Tuple[tf.Tensor,...]]]]]  # nopep8
 
@@ -1254,6 +1255,8 @@ def add_templated_layer(lself, name, layer_class, **layer_desc):
         layer_.kwargs["output"] = layer_class.fixup_out_data(layer_.kwargs["output"], network=self.net)
         layer_.kwargs["output"].sanity_check(ignore_placeholder=True)  # placeholder might be overwritten later
         layer_.init(layer_class=layer_class, **layer_.kwargs)
+        if layer_.need_last:
+          self.prev_layers_needed.add(name)
         if layer_ in ConstructCtx.partially_finished:
           if lself.got_uninitialized_deps_count == 0:  # in this case, we safely know that it is finished
             ConstructCtx.partially_finished.remove(layer_)
@@ -1678,6 +1681,15 @@ def get_layer(name):
                     layer_choices.beam_size == layer.output.beam.beam_size == layer_template.output.beam.beam_size), (
               "Layer %r has buggy search choices resolution." % layer,
               self.net.debug_search_choices(layer) or "see search choices debug output")
+        if name == "end":
+          # Special logic for the end layer, to always logical_or to the prev:end layer.
+          assert layer.output.shape == (), "%s: 'end' layer %r unexpected shape" % (self.parent_rec_layer, layer)
+          prev_end_layer = self.net.layers["prev:end"]
+          choices = layer.get_search_choices()
+          if choices:
+            prev_end_layer = choices.translate_to_this_search_beam(prev_end_layer)
+          with tf.name_scope("end_flag"):
+            layer.output.placeholder = tf.logical_or(prev_end_layer.output.placeholder, layer.output.placeholder)
         return layer
       except Exception as exc:
         self._handle_construct_exception(description="in-loop construction of layer %r" % name, exception=exc)
@@ -1730,6 +1742,14 @@ def get_layer_from_outside(self, layer_name):
          Add 'is_output_layer':True to the layer options." % (self.parent_rec_layer, layer_name))
     return None
 
+  def get_layer_last_frame(self, layer_name):
+    """
+    :param str layer_name:
+    :return: the output from the last loop frame. compatible to layer templates. masking correctly taken into account
+    :rtype: Data
+    """
+    return self._last_frames[layer_name]
+
   def get_sub_networks(self):
     """
     :rtype: list[returnn.tf.network.TFNetwork]
@@ -1880,6 +1900,44 @@ def get_layer_rec_var_from_loop_vars(self, loop_vars, layer_name, final_frame=Fa
           rec_vars_outputs, seq_len=seq_len)
     return rec_vars_outputs
 
+  def _assign_last_frames_from_loop_vars(self, loop_vars):
+    """
+    :param (list[tf.Tensor],list[tf.Tensor]) loop_vars: loop_vars like in self.get_next_loop_vars()
+    """
+    final_outputs_flat, final_extra_flat = loop_vars
+    assert len(final_outputs_flat) == len(self._initial_outputs)
+    assert len(final_extra_flat) == len(self._initial_extra_outputs)
+    for k, v in zip(sorted(self._initial_outputs), final_outputs_flat):
+      layer_template = self.layer_data_templates[k]
+      if not layer_template.need_last:
+        continue
+      out = layer_template.output.copy_template()
+      out.placeholder = v
+      assert k not in self._last_frames
+      self._last_frames[k] = out
+
+  def _assign_last_frames_from_optimized(self):
+    for k in self.prev_layers_needed:
+      if k in self.output_layers_moved_out:
+        layer = self.output_layers_net.layers[k]
+      elif k in self.input_layers_moved_out:
+        layer = self.input_layers_net.layers[k]
+      else:
+        continue
+      assert isinstance(layer, LayerBase)
+      if not layer.need_last:
+        continue
+      with tf.name_scope(layer.tf_scope_name):
+        indices = layer.output.get_sequence_lengths() - 1  # [B]
+        out = layer.output.copy_as_batch_major().copy_with_time_dim_axis(1)  # [B,T,...]
+        v = tf.gather(out.placeholder, indices=tf.maximum(indices, 0), batch_dims=1, axis=1)  # [B,...]
+        v = tf.where(tf.less(indices, 0), self._get_init_output(k), v)
+        layer_template = self.layer_data_templates[k]
+        out = layer_template.output.copy_template()
+        out.placeholder = v
+        assert k not in self._last_frames
+        self._last_frames[k] = out
+
   def get_parent_deps(self):
     """
     :return: list of dependencies to the parent network
@@ -2502,14 +2560,14 @@ def body(i, net_vars, acc_tas, seq_len_info=None):
           step_info_i += global_tensor(
             lambda: tf_compat.v1.placeholder(tf.int32, (), name="global_rec_step_offset"),
             name="global_rec_step_offset")
-        rec_step_info = dict(i=step_info_i, end_flag=None, seq_lens=fixed_seq_len)
-        end_flag, dyn_seq_len, prev_end_layer = None, None, None
+        rec_step_info = dict(i=step_info_i, prev_end_flag=None, seq_lens=fixed_seq_len)
+        prev_end_flag, prev_dyn_seq_len, prev_end_layer = None, None, None
         if seq_len_info:
-          end_flag, dyn_seq_len = seq_len_info
-          rec_step_info["end_flag"] = end_flag
-          prev_end_layer = self.layer_data_templates["end"].copy_as_prev_time_frame(prev_output=end_flag)
+          prev_end_flag, prev_dyn_seq_len = seq_len_info
+          rec_step_info["prev_end_flag"] = prev_end_flag
+          prev_end_layer = self.layer_data_templates["end"].copy_as_prev_time_frame(prev_output=prev_end_flag)
           self.net.layers["prev:end"] = prev_end_layer
-          rec_step_info["end_flag_source"] = prev_end_layer
+          rec_step_info["prev_end_layer"] = prev_end_layer
         self.net.set_rec_step_info(**rec_step_info)
         # get next loop vars (net_vars)
         from returnn.tf.util.basic import identity_op_nested
@@ -2571,6 +2629,46 @@ def maybe_transform(layer):
           transformed_cache[layer] = transformed_layer
           return transformed_layer
 
+        # Handle need_last.
+        for k in sorted(self._initial_outputs):
+          layer = self.net.layers[k]
+          if not layer.need_last:
+            continue
+          # We need to make sure that the output is correct also for the seqs which are already ended.
+          prev_layer = self.net.layers["prev:" + k]
+          # Last frame corresponds to the frame seq_len - 1.
+          # With include_eos=False, when "end" layer is True <=> we are behind the last frame.
+          # With include_eos=True, when "prev:end" layer is True <=> we are behind the last frame.
+          rel_end_layer = self.net.layers["prev:end"] if rec_layer.include_eos else self.net.layers["end"]
+          choices = layer.get_search_choices()
+          if choices:
+            prev_layer, rel_end_layer = choices.translate_to_this_search_beam([prev_layer, rel_end_layer])
+          from returnn.tf.util.basic import where_bc
+          layer.output.placeholder = where_bc(
+            condition=rel_end_layer.output.copy_compatible_to(layer.output).placeholder,
+            x=prev_layer.output.placeholder, y=layer.output.placeholder)
+
+        if seq_len_info is not None:
+          end_layer = maybe_transform(self.net.layers["end"])
+          assert end_layer.output.shape == (), "end layer %r unexpected shape" % end_layer
+          end_flag = end_layer.output.placeholder  # we handled the logical_or(prev_end...) in construct get_layer
+          end_flag.set_shape([None])
+          choices = end_layer.get_search_choices()
+          if choices:
+            from .basic import SelectSearchSourcesLayer
+            prev_end_layer = choices.translate_to_this_search_beam(prev_end_layer)
+            assert isinstance(prev_end_layer, SelectSearchSourcesLayer), (
+              "unexpected search choices: cur end %r, prev end %r" % (choices, prev_end_layer.get_search_choices()))
+            prev_end_flag = prev_end_layer.output.placeholder
+            with tf.name_scope("dyn_seq_len"):
+              prev_dyn_seq_len = prev_end_layer.transform_func(prev_dyn_seq_len)
+          with tf.name_scope("dyn_seq_len"):
+            dyn_seq_len = prev_dyn_seq_len + tf.where(
+              prev_end_flag if rec_layer.include_eos else end_flag,
+              constant_with_shape(0, shape=tf.shape(end_flag)),
+              constant_with_shape(1, shape=tf.shape(end_flag)))  # (batch * beam,)
+            seq_len_info = (end_flag, dyn_seq_len)
+
         outputs_flat = [
           maybe_transform(self.net.layers[k]).output.copy_compatible_to(
             self.layer_data_templates[k].output).placeholder
@@ -2582,42 +2680,6 @@ def maybe_transform(layer):
           extra_flat.append(sorted_values_from_dict(layer.rec_vars_outputs))
         net_vars = (outputs_flat, extra_flat)
 
-        if seq_len_info is not None:
-          assert self.net.layers["end"].output.shape == (), "end layer %r unexpected shape" % self.net.layers["end"]
-          end_layer = maybe_transform(self.net.layers["end"])
-          choices = end_layer.get_search_choices()
-          if choices:
-            from .basic import SelectSearchSourcesLayer
-            cur_end_layer = choices.translate_to_this_search_beam(prev_end_layer)
-            assert isinstance(cur_end_layer, SelectSearchSourcesLayer), (
-              "unexpected search choices: cur end %r, prev end %r" % (choices, prev_end_layer.get_search_choices()))
-            assert cur_end_layer.search_choices_seq, (
-              "unexpected search choices: cur end %r (via %r), prev end %r (via %r)" % (
-                choices, end_layer, prev_end_layer.get_search_choices(), prev_end_layer))
-            assert cur_end_layer.output.shape == (), "end layer %r unexpected shape" % cur_end_layer
-            with tf.name_scope("end_flag"):
-              end_flag = cur_end_layer.output.placeholder
-              end_flag = tf.logical_or(end_flag, self.net.layers["end"].output.placeholder)  # (batch * beam,)
-              end_flag.set_shape([None])
-            with tf.name_scope("dyn_seq_len"):
-              dyn_seq_len = cur_end_layer.transform_func(dyn_seq_len)
-              dyn_seq_len += tf.where(
-                cur_end_layer.output.placeholder if rec_layer.include_eos else end_flag,
-                constant_with_shape(0, shape=tf.shape(end_flag)),
-                constant_with_shape(1, shape=tf.shape(end_flag)))  # (batch * beam,)
-              seq_len_info = (end_flag, dyn_seq_len)
-          else:
-            assert self.net.layers["end"].output.shape == (), "end layer %r unexpected shape" % self.net.layers["end"]
-            with tf.name_scope("end_flag"):
-              end_flag = tf.logical_or(end_flag, self.net.layers["end"].output.placeholder)
-              end_flag.set_shape([None])
-            with tf.name_scope("dyn_seq_len"):
-              dyn_seq_len += tf.where(
-                end_flag,
-                constant_with_shape(0, shape=tf.shape(end_flag)),
-                constant_with_shape(1, shape=tf.shape(end_flag)))  # (batch * beam,)
-              seq_len_info = (end_flag, dyn_seq_len)
-
         assert len(acc_tas) == len(outputs_to_accumulate)
         acc_tas = [
           out.write_to_tensor_array(acc_ta, index=i)
@@ -2723,6 +2785,7 @@ def cond(i, net_vars, acc_tas, seq_len_info=None):
       self.final_acc_tas_dict = {
         out.name: out.get_final_tensor_array(final_acc_ta)
         for (final_acc_ta, out) in zip(final_acc_tas, outputs_to_accumulate)}  # type: typing.Dict[str,typing.Optional[tf.TensorArray]]  # nopep8
+      self._assign_last_frames_from_loop_vars(final_net_vars)
     else:  # no layers inside loop, all optimized out
       seq_len = None
       final_net_vars = None
@@ -2732,6 +2795,7 @@ def cond(i, net_vars, acc_tas, seq_len_info=None):
     self._construct_output_layers_moved_out(
       loop_accumulated=self.final_acc_tas_dict, seq_len=seq_len,
       extra_output_layers=extra_output_layers, final_net_vars=final_net_vars)
+    self._assign_last_frames_from_optimized()
 
     if layer_names_with_losses:
       from returnn.tf.network import LossHolder
@@ -2831,6 +2895,52 @@ def cond(i, net_vars, acc_tas, seq_len_info=None):
 
     return output
 
+  def _get_search_choice_seq(self, search_choices):
+    """
+    :param SearchChoices search_choices:
+    :rtype: list[LayerBase]
+    """
+    layer_choice = search_choices.owner
+    assert not layer_choice.name.startswith("prev:")
+    assert not isinstance(layer_choice, _TemplateLayer)
+
+    # There can be multiple choices in a single rec step. Collect them.
+    # Find next choice layer. Then iterate through its source choice layers through time
+    # and resolve the output over time to be in line with the final output search choices.
+
+    def get_choice_seq(choice_base):
+      """
+      :param LayerBase choice_base:
+      :return: choice_seq, prev_choice
+      :rtype: (list[LayerBase], _TemplateLayer)
+      """
+      choice_seq = [choice_base]
+      choice = choice_base
+      while True:
+        assert choice.network is self.net, "not yet implemented otherwise"
+        assert ("choice_%s" % choice.name) in self.final_acc_tas_dict
+        assert choice.search_choices
+        assert choice.search_choices.src_layer
+        choice = choice.search_choices.src_layer
+        if isinstance(choice, _TemplateLayer):
+          assert choice.is_prev_time_frame
+          return choice_seq, choice
+        choice_seq.append(choice)
+
+    # Get the very latest choice in the rec layer.
+    latest_layer_choice = layer_choice
+    _, latest_layer_choice = get_choice_seq(latest_layer_choice)
+    assert latest_layer_choice.name.startswith("prev:")
+    latest_layer_choice = self.net.layers[latest_layer_choice.name[len("prev:"):]]
+    assert latest_layer_choice.search_choices
+
+    # Get the whole choice sequence, starting from the latest to the first.
+    choice_seq_in_frame, prev_frame_choice = get_choice_seq(latest_layer_choice)
+    assert prev_frame_choice.name == "prev:%s" % latest_layer_choice.name
+    assert layer_choice in choice_seq_in_frame
+    assert choice_seq_in_frame[0] is latest_layer_choice
+    return choice_seq_in_frame
+
   def _opt_search_resolve(self, layer_name, acc_ta, final_net_vars, seq_len, search_choices_cache):
     """
     This assumes that we have frame-wise accumulated outputs of the specific layer (acc_ta).
@@ -2883,37 +2993,12 @@ def _opt_search_resolve(self, layer_name, acc_ta, final_net_vars, seq_len, searc
       assert layer_choice.is_prev_time_frame
       assert layer_choice.name.startswith("prev:")
       layer_choice = self.net.layers[layer_choice.name[len("prev:"):]]
+      assert layer_choice.search_choices
+      search_choices = layer_choice.search_choices
       is_prev_choice = True
 
-    # There can be multiple choices in a single rec step. Collect them.
-    # Find next choice layer. Then iterate through its source choice layers through time
-    # and resolve the output over time to be in line with the final output search choices.
-
-    def get_choice_seq(choice_base):
-      """
-      :param LayerBase choice_base:
-      :return: choice_seq, prev_choice
-      :rtype: (list[LayerBase], _TemplateLayer)
-      """
-      choice_seq = [choice_base]
-      choice = choice_base
-      while True:
-        assert choice.network is self.net, "not yet implemented otherwise"
-        assert ("choice_%s" % choice.name) in self.final_acc_tas_dict
-        assert choice.search_choices
-        assert choice.search_choices.src_layer
-        choice = choice.search_choices.src_layer
-        if isinstance(choice, _TemplateLayer):
-          assert choice.is_prev_time_frame
-          return choice_seq, choice
-        choice_seq.append(choice)
-
-    # Get the very latest choice in the rec layer.
-    latest_layer_choice = layer_choice
-    _, latest_layer_choice = get_choice_seq(latest_layer_choice)
-    assert latest_layer_choice.name.startswith("prev:")
-    latest_layer_choice = self.net.layers[latest_layer_choice.name[len("prev:"):]]
-    assert latest_layer_choice.search_choices
+    choice_seq_in_frame = self._get_search_choice_seq(search_choices)
+    latest_layer_choice = choice_seq_in_frame[0]
 
     # The max_seq_len might actually be one more, as it includes the EOS, but that does not matter;
     # we just want to create a new acc_ta with the same length.
@@ -2927,11 +3012,6 @@ def get_choice_seq(choice_base):
     initial_beam_choices = expand_dims_unbroadcast(
       initial_beam_choices, axis=0, dim=batch_dim)  # (batch, beam_out)
 
-    # Get the whole choice sequence, starting from the latest to the first.
-    choice_seq_in_frame, prev_frame_choice = get_choice_seq(latest_layer_choice)
-    assert prev_frame_choice.name == "prev:%s" % latest_layer_choice.name
-    assert layer_choice in choice_seq_in_frame
-
     # Translate the beams of seq_lens to correspond to the latest choice
     try:
       end_layer = self.net.get_layer("end")
@@ -3634,6 +3714,7 @@ def init(self, output, layer_class, template_type="template", **kwargs):
     self.kwargs.setdefault("name", self.name)
     self.kwargs["output"] = output
     self._is_output_layer = kwargs.get("is_output_layer", None)
+    self.need_last = kwargs.get("need_last", False)
     if self._has_search_choices():
       self.search_choices = SearchChoices(owner=self, beam_size=self._get_search_choices_beam_size())
     self.collocate_with = kwargs.get("collocate_with", None) or []
@@ -3945,12 +4026,12 @@ class RecStepInfoLayer(LayerBase):
 
   layer_class = ":i"
 
-  def __init__(self, i=None, end_flag=None, end_flag_source=None, seq_lens=None, **kwargs):
+  def __init__(self, i=None, prev_end_flag=None, prev_end_layer=None, seq_lens=None, **kwargs):
     """
     :param tf.Tensor|None i: scalar, int32, current step (time)
-    :param tf.Tensor|None end_flag: (batch,), bool, says that the current sequence has ended.
+    :param tf.Tensor|None prev_end_flag: (batch,), bool, says that the current sequence has ended.
       Can be with beam. In that case, end_flag_source should be "prev:end", and define the search choices.
-    :param LayerBase|None end_flag_source:
+    :param LayerBase|None prev_end_layer: corresponds to the "prev:end" layer if available
     :param tf.Tensor|None seq_lens: (batch,) int32, seq lens
     """
     if "output" not in kwargs:
@@ -3959,13 +4040,13 @@ def __init__(self, i=None, end_flag=None, end_flag_source=None, seq_lens=None, *
     super(RecStepInfoLayer, self).__init__(**kwargs)
     self.step = None
     self._end_flag = None
-    self.end_flag_source = None
+    self.prev_end_layer = None
     if not self.output.have_time_axis():  # the normal case
       assert i is not None and i.get_shape().ndims == 0
       self.output.placeholder = i
       self.step = i
-      self._end_flag = end_flag
-      self.end_flag_source = end_flag_source
+      self._end_flag = prev_end_flag
+      self.prev_end_layer = prev_end_layer
     else:
       # This only is valid if we are moved out from a RecLayer.
       assert self.output.size_placeholder and 0 in self.output.size_placeholder
@@ -3973,12 +4054,14 @@ def __init__(self, i=None, end_flag=None, end_flag_source=None, seq_lens=None, *
       self.output.placeholder = tf.range(tf.reduce_max(seq_lens))
     self._seq_lens = seq_lens
     if seq_lens is None:
-      assert end_flag_source
+      assert prev_end_layer
 
-  def get_end_flag(self, target_search_choices):
+  def get_prev_end_flag(self, target_search_choices):
     """
     :param SearchChoices|None target_search_choices:
-    :return: (batch,) of type bool. batch might include beam size
+    :return: (batch,) of type bool. batch might include beam size.
+      This returns the end flag corresponding to the last frame.
+      I.e. if the "end" layer exists and is used, this is "prev:end".
     :rtype: tf.Tensor
     """
     if self._end_flag is not None:
@@ -3987,21 +4070,28 @@ def get_end_flag(self, target_search_choices):
       assert self._seq_lens is not None
       from returnn.tf.util.basic import reuse_name_scope_of_tensor
       with reuse_name_scope_of_tensor(self.step, postfix="/end_flag"):
-        end_flag = tf.greater_equal(self.step, self._seq_lens)
+        rec_layer = self.network.get_rec_parent_layer()
+        assert rec_layer
+        end_flag = tf.greater_equal(
+          # Without include_eos, end_flag=True happens first in frame seq_lens.
+          # With include_eos, end_flag=True happens first in frame seq_lens - 1.
+          # We are supposed to return the prev end flag here.
+          self.step,
+          self._seq_lens if rec_layer.include_eos else (self._seq_lens + 1))
     source_search_choices = None
-    if self.end_flag_source:
-      source_search_choices = self.end_flag_source.get_search_choices()
+    if self.prev_end_layer:
+      source_search_choices = self.prev_end_layer.get_search_choices()
     if target_search_choices:
       if source_search_choices:
-        assert self.end_flag_source
-        end_flag_transformed_layer = target_search_choices.translate_to_this_search_beam(self.end_flag_source)
+        assert self.prev_end_layer
+        end_flag_transformed_layer = target_search_choices.translate_to_this_search_beam(self.prev_end_layer)
         assert isinstance(end_flag_transformed_layer, LayerBase)
         end_flag = end_flag_transformed_layer.output.placeholder
       else:
         from returnn.tf.util.basic import tile_transposed
         end_flag = tile_transposed(end_flag, axis=0, multiples=target_search_choices.beam_size)
     else:
-      assert not self.end_flag_source or not source_search_choices
+      assert not self.prev_end_layer or not source_search_choices
     return end_flag
 
   @classmethod
@@ -4035,6 +4125,50 @@ def get_out_data_from_opts(cls, network, **kwargs):
       size_placeholder={0: seq_lens})
 
 
+class RecLastOutputLayer(LayerBase):
+  """
+  Gets the last output from some sub layer inside a :class:`RecLayer`.
+  You should explicitly set ``need_last`` on the specific layer such that this information is available.
+  """
+  layer_class = "rec_last_output"
+
+  def __init__(self, rec_layer, sub_layer_name, **kwargs):
+    """
+    :param RecLayer rec_layer:
+    :param str sub_layer_name:
+    """
+    super(RecLastOutputLayer, self).__init__(**kwargs)
+    assert isinstance(rec_layer, RecLayer)
+    cell = rec_layer.cell
+    assert isinstance(cell, _SubnetworkRecCell)
+    out = cell.get_layer_last_frame(sub_layer_name)
+    assert out.dim_tags == self.output.dim_tags
+    self.output.placeholder = out.placeholder
+
+  @classmethod
+  def transform_config_dict(cls, d, network, get_layer):
+    """
+    :param dict[str] d:
+    :param returnn.tf.network.TFNetwork network:
+    :param get_layer:
+    """
+    d.setdefault("from", ())
+    super(RecLastOutputLayer, cls).transform_config_dict(d, network=network, get_layer=get_layer)
+    d["rec_layer"] = get_layer(d["rec_layer"])
+
+  @classmethod
+  def get_out_data_from_opts(cls, rec_layer, sub_layer_name, **kwargs):
+    """
+    :param RecLayer rec_layer:
+    :param str sub_layer_name:
+    :rtype: Data
+    """
+    assert isinstance(rec_layer, RecLayer)
+    cell = rec_layer.cell
+    assert isinstance(cell, _SubnetworkRecCell)
+    return cell.layer_data_templates[sub_layer_name].output
+
+
 class RnnCellLayer(_ConcatInputLayer):
   """
   Wrapper around tf.contrib.rnn.RNNCell.
@@ -4878,7 +5012,7 @@ def __init__(self, beam_size, keep_beams=False,
         if self.length_normalization:
           assert self.network.have_rec_step_info()
           t = self.network.get_rec_step_index()  # scalar
-          end_flags_flat = self.network.get_rec_step_info().get_end_flag(
+          end_flags_flat = self.network.get_rec_step_info().get_prev_end_flag(
             target_search_choices=base_search_choices)  # (batch * beam_in,)
           with tf.name_scope("length_normalization"):
             end_flags = tf.reshape(end_flags_flat, [net_batch_dim, beam_in])  # (batch, beam_in)
@@ -4896,7 +5030,7 @@ def __init__(self, beam_size, keep_beams=False,
         if self.network.have_rec_step_info():
           scores_in = filter_ended_scores(
             scores_in,
-            end_flags=self.network.get_rec_step_info().get_end_flag(target_search_choices=base_search_choices),
+            end_flags=self.network.get_rec_step_info().get_prev_end_flag(target_search_choices=base_search_choices),
             dim=scores_in_dim, batch_dim=net_batch_dim * scores_beam_in)  # (batch * beam_in, dim)
         scores_in = tf.reshape(scores_in, [net_batch_dim, scores_beam_in, scores_in_dim])  # (batch, beam_in, dim)
         with tf.control_dependencies([
@@ -4917,7 +5051,7 @@ def __init__(self, beam_size, keep_beams=False,
           with tf.name_scope("custom_score_combine"):
             if self.network.have_rec_step_info():
               t = self.network.get_rec_step_index()  # scalar
-              end_flags_flat = self.network.get_rec_step_info().get_end_flag(
+              end_flags_flat = self.network.get_rec_step_info().get_prev_end_flag(
                 target_search_choices=base_search_choices)  # (batch * beam_in,)
               end_flags = tf.reshape(end_flags_flat, [net_batch_dim, beam_size])  # (batch, beam_in)
               end_flags = end_flags[:, :base_beam_in]  # see scores_in
@@ -5086,7 +5220,7 @@ def __init__(self, beam_size, keep_beams=False,
           scores_in_dim = tf.shape(self.sources[0].output.placeholder)[self.sources[0].output.feature_dim_axis]
         scores_in = filter_ended_scores(
           scores_in,
-          end_flags=self.network.get_rec_step_info().get_end_flag(target_search_choices=base_search_choices),
+          end_flags=self.network.get_rec_step_info().get_prev_end_flag(target_search_choices=base_search_choices),
           dim=scores_in_dim, batch_dim=tf.shape(scores_in)[0])  # (batch * beam_in, dim)
         # We also assume that the ground truth output are 0 when the seq ended.
       scores_in_ = batch_gather(scores_in, self.output.placeholder)  # (batch*beam_in,)
@@ -6910,7 +7044,7 @@ def __init__(self, debug=False, blank_idx=None, **kwargs):
     self._last_row = self._rec_previous_layer.rec_vars_outputs["state"]
     rec_step_info = self.network.get_rec_step_info()
     batch_dim = self.get_batch_dim()
-    mask_flag = rec_step_info.get_end_flag(target_search_choices=self.get_search_choices())
+    mask_flag = rec_step_info.get_prev_end_flag(target_search_choices=self.get_search_choices())
     source = source_data.placeholder
     if blank_idx is None:
       from returnn.tf.util.basic import expand_dims_unbroadcast
@@ -6936,7 +7070,7 @@ def __init__(self, debug=False, blank_idx=None, **kwargs):
           "src_beams", choice.src_beams if choice.src_beams is not None else "None"]
       print_out += [
         "a_n", rec_step_info.step,
-        "a_ended", rec_step_info.get_end_flag(target_search_choices=self.get_search_choices()),
+        "a_ended", rec_step_info.get_prev_end_flag(target_search_choices=self.get_search_choices()),
         "a", vocab_idx_repr(source_data.placeholder, target_data),
         "b", vocab_idx_repr(target_data.placeholder, target_data),
         "b_len", target_data.get_sequence_lengths(),
@@ -7067,7 +7201,7 @@ def __init__(self, debug=False, blank_idx=None, **kwargs):
     reduce_out = next_edit_distance_reduce(
       last_row=last_row,
       a=successors, a_n=src_len,
-      a_ended=rec_step_info.get_end_flag(target_search_choices=self.get_search_choices()),
+      a_ended=rec_step_info.get_prev_end_flag(target_search_choices=self.get_search_choices()),
       b=target_data.placeholder, b_len=target_data.get_sequence_lengths(),
       optimal_completion=True, a_blank_idx=blank_idx)
     reduce_out.set_shape((None, self.output.dim))
@@ -7083,7 +7217,7 @@ def __init__(self, debug=False, blank_idx=None, **kwargs):
       top_values = -top_values
       print_out += [
         "a_n", rec_step_info.step,
-        "a_ended", rec_step_info.get_end_flag(target_search_choices=self.get_search_choices()),
+        "a_ended", rec_step_info.get_prev_end_flag(target_search_choices=self.get_search_choices()),
         "a best", vocab_idx_repr(top_indices, target_data), top_values,
         "b", vocab_idx_repr(target_data.placeholder, target_data),
         "b_len", target_data.get_sequence_lengths(),
diff --git a/returnn/tf/network.py b/returnn/tf/network.py
--- a/returnn/tf/network.py
+++ b/returnn/tf/network.py
@@ -2168,18 +2168,19 @@ def get_global_batch_info(self):
         return layer.output.batch.get_global_base()
     raise Exception("%s: Cannot get global batch info" % root)
 
-  def set_rec_step_info(self, i, end_flag=None, end_flag_source=None, seq_lens=None):
+  def set_rec_step_info(self, i, prev_end_flag=None, prev_end_layer=None, seq_lens=None):
     """
     Used by _SubnetworkRecCell.
 
     :param tf.Tensor i: scalar, int32, current step (time)
-    :param tf.Tensor|None end_flag: (batch,), bool, says that the current sequence has ended
-    :param LayerBase|None end_flag_source:
+    :param tf.Tensor|None prev_end_flag: (batch,), bool, says that the current sequence has ended.
+     This is about the last frame, not the current!
+    :param LayerBase|None prev_end_layer:
     :param tf.Tensor|None seq_lens: (batch,) int32, seq lens
     """
     from returnn.tf.layers.rec import RecStepInfoLayer
     self.layers[":i"] = RecStepInfoLayer(
-      name=":i", network=self, i=i, end_flag=end_flag, end_flag_source=end_flag_source, seq_lens=seq_lens)
+      name=":i", network=self, i=i, prev_end_flag=prev_end_flag, prev_end_layer=prev_end_layer, seq_lens=seq_lens)
 
   def is_inside_rec_layer(self, inside_loop=True):
     """


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -6169,7 +6169,7 @@ def test_OptimalCompletionsLayer():
     # Fake that we are inside a rec layer.
     net.set_rec_step_info(
       i=tf.convert_to_tensor(0, name="i"),
-      end_flag=expand_dims_unbroadcast(tf.convert_to_tensor(False), 0, n_batch),
+      prev_end_flag=expand_dims_unbroadcast(tf.convert_to_tensor(False), 0, n_batch),
       seq_lens=target.get_sequence_lengths())
     kwargs = dict(
       name="opt_completions", network=net, debug=True, target="target",