returnn 1.20250108.104903__tar.gz → 1.20250109.114342__tar.gz

{returnn-1.20250108.104903/returnn.egg-info → returnn-1.20250109.114342}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250108.104903
+Version: 1.20250109.114342
 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.114342/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20250109.114342'
+long_version = '1.20250109.114342+git.bbbe0ae'

{returnn-1.20250108.104903 → returnn-1.20250109.114342}/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.20250108.104903 → returnn-1.20250109.114342}/returnn/frontend/attention.py

@@ -227,23 +227,28 @@ def _causal_self_att_step(
     state: Optional[CausalSelfAttentionState],
     self: rf.Module,
 ) -> Tuple[Tensor, Tensor, Dim, CausalSelfAttentionState]:
+    new_state = CausalSelfAttentionState()
     if axis == single_step_dim:
         assert state, f"{self}: need state for single step"
         k, hist_dim = rf.cum_concat_step(k, prev_accum=state.k_accum, axis=state.accum_axis)
         v, _ = rf.cum_concat_step(v, prev_accum=state.v_accum, out_spatial_dim=hist_dim, axis=state.accum_axis)
+        new_state.k_accum = k
+        new_state.v_accum = v
+        new_state.accum_axis = hist_dim
     else:
         if state and state.accum_axis.dimension != 0:
             raise NotImplementedError(  # need to concat ...
                 f"{self}: on sequence over {axis} with initial state {state} not implemented yet"
             )
+        # For further state usage with single_step_dim, keep using the original axis, not hist_dim,
+        # because the masking via hist_dim is not necessary for future single steps.
+        new_state.k_accum = k
+        new_state.v_accum = v
+        new_state.accum_axis = axis
         # See CumConcatLayer and https://github.com/rwth-i6/returnn/issues/391 for the idea.
         hist_dim = Dim(rf.range_over_dim(axis, device="cpu") + 1, name=f"{axis.description}:kv")
         k, _ = rf.replace_dim(k, in_dim=axis, out_dim=hist_dim)
         v, _ = rf.replace_dim(v, in_dim=axis, out_dim=hist_dim)
-    new_state = CausalSelfAttentionState()
-    new_state.k_accum = k
-    new_state.v_accum = v
-    new_state.accum_axis = hist_dim
     return k, v, hist_dim, new_state
 
 

{returnn-1.20250108.104903 → returnn-1.20250109.114342/returnn.egg-info}/PKG-INFO

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

{returnn-1.20250108.104903 → returnn-1.20250109.114342}/tests/test_rf_decoder_transformer.py

@@ -10,7 +10,7 @@ import unittest
 import torch
 from returnn.util import better_exchook
 from returnn.util.debug import PyTracer, check_py_traces_rf_to_pt_equal
-from returnn.tensor import Tensor, Dim
+from returnn.tensor import Tensor, Dim, single_step_dim
 import returnn.frontend as rf
 
 
@@ -26,6 +26,44 @@ def _setup():
 _setup()
 
 
+def test_transformer_prefix_single_step():
+    # We have some prefix we want to feed and then only get the logits of the last step,
+    # and also continue step-wise (single_step_dim) from that state.
+
+    from returnn.frontend.decoder.transformer import TransformerDecoder
+
+    rf.select_backend_torch()
+
+    vocab_dim = Dim(11, name="vocab")
+    lm = TransformerDecoder(
+        encoder_dim=None, vocab_dim=vocab_dim, model_dim=Dim(32, name="model"), num_layers=2, num_heads=2
+    )
+
+    # We currently need this for the case without batch dim.
+    # We might want to extend the test case later in the future.
+    state = lm.default_initial_state(batch_dims=[])
+
+    # Not sure if a static dim is reasonable?
+    # But for the single seq case (no batch dim), and eager mode (PyTorch), we know the size.
+    spatial_dim = Dim(7, name="spatial")
+    prefix_seq = rf.convert_to_tensor([3, 4, 5, 6, 7, 2, 1], dims=[spatial_dim], sparse_dim=vocab_dim)
+    logits, state = lm(prefix_seq, spatial_dim=spatial_dim, state=state)
+    print(logits, state["0"].self_att)
+    assert state["0"].self_att.accum_axis == spatial_dim
+    assert state["0"].self_att.accum_axis in state["0"].self_att.k_accum.dims
+
+    # Now feed some other labels step by step.
+    step = 1
+    for label_idx in [8, 9, 10]:
+        label = rf.constant(label_idx, dims=[], sparse_dim=vocab_dim)
+        logits, state = lm(label, spatial_dim=single_step_dim, state=state)
+        print(logits, state["0"].self_att)
+        # This assumes it's a static dim.
+        assert state["0"].self_att.accum_axis.dimension == spatial_dim.dimension + step
+        assert state["0"].self_att.accum_axis in state["0"].self_att.k_accum.dims
+        step += 1
+
+
 def test_llama():
     """
     Test that we can reproduce the Llama model.

returnn-1.20250108.104903/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250108.104903'
-long_version = '1.20250108.104903+git.fee68a8'