returnn 1.20250109.111908__tar.gz → 1.20250109.145311__tar.gz

{returnn-1.20250109.111908/returnn.egg-info → returnn-1.20250109.145311}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250109.111908
+Version: 1.20250109.145311
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20250109.145311/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250109.145311'
+long_version = '1.20250109.145311+git.855b3ae'

{returnn-1.20250109.111908 → returnn-1.20250109.145311}/returnn/datasets/postprocessing.py

@@ -22,9 +22,9 @@ __all__ = ["PostprocessingDataset", "LaplaceOrdering", "Sequential"]
 
 class PostprocessingDataset(CachedDataset2):
     """
-    A dataset that allows for generic post-processing of data from another dataset
-    using a function on the segment level and on the level of multiple segments via
-    an iterator.
+    A dataset that allows for generic online post-processing of data from another
+    dataset using a function on the segment level and on the level of multiple
+    segments via an iterator.
 
     This allows integrating various data augmentation techniques like e.g. Mixup,
     SpecAugment or speed perturbation into the data loading pipeline.

{returnn-1.20250109.111908 → returnn-1.20250109.145311}/returnn/frontend/decoder/transformer.py

@@ -203,6 +203,7 @@ class TransformerDecoder(rf.Module):
         state: rf.State,
         encoder: Optional[rf.State] = None,
         collected_outputs: Optional[Dict[str, Tensor]] = None,
+        output_only_last_frame: bool = False,
     ) -> Tuple[Tensor, rf.State]:
         """
         forward, single step or whole sequence.
@@ -212,6 +213,8 @@ class TransformerDecoder(rf.Module):
         :param state: e.g. via :func:`default_initial_state`
         :param encoder: via :func:`transform_encoder`
         :param collected_outputs:
+        :param output_only_last_frame: if True, and spatial_dim is not single_step_dim,
+            the returned logits will only be for the last frame
         :return: logits, new state
         """
         new_state = rf.State()
@@ -236,6 +239,9 @@ class TransformerDecoder(rf.Module):
             if collected_outputs is not None:
                 collected_outputs[layer_name] = decoded
 
+        if output_only_last_frame and spatial_dim != single_step_dim:
+            decoded = rf.gather(decoded, axis=spatial_dim, indices=rf.last_frame_position_of_dim(spatial_dim))
+
         decoded = self.final_layer_norm(decoded)
         logits = self.logits(decoded)
 

{returnn-1.20250109.111908 → returnn-1.20250109.145311}/returnn/frontend/dims.py

@@ -19,6 +19,7 @@ __all__ = [
     "dim_match_priority_when_needed",
     "num_elements_of_shape",
     "masked_fraction_of_shape",
+    "last_frame_position_of_dim",
 ]
 
 
@@ -195,3 +196,28 @@ def masked_fraction_of_shape(dims: Union[Dim, Sequence[Dim]], *, inverse: bool =
     for dim in dims:
         num_elems_total *= dim.get_dim_value_tensor()
     return (num_elems_masked / num_elems_total) if not inverse else (num_elems_total / num_elems_masked)
+
+
+def last_frame_position_of_dim(
+    dim: Dim, *, device: Optional[str] = None, allow_scalar_on_cpu: bool = True
+) -> Union[int, Tensor]:
+    """
+    :param dim:
+    :param device:
+    :param allow_scalar_on_cpu: if device is not given, we would use rf.get_default_device() for dynamic sizes.
+        If this is True, we would allow to return a scalar dynamic size on CPU.
+    :return: last frame position of dim
+    """
+    if dim.size is not None:
+        if dim.size > 0:
+            return dim.size - 1
+        return 0
+    if device is None:
+        if allow_scalar_on_cpu and not dim.dyn_size_ext.dims:
+            device = "cpu"
+        else:
+            device = rf.get_default_device()
+    pos = dim.get_dyn_size_ext_for_device(device) - 1
+    pos = rf.maximum(pos, 0)
+    pos.sparse_dim = dim
+    return pos

{returnn-1.20250109.111908 → returnn-1.20250109.145311/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250109.111908
+Version: 1.20250109.145311
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20250109.111908/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250109.111908'
-long_version = '1.20250109.111908+git.dfcd25c'