returnn 1.20250106.235006__tar.gz → 1.20250109.111908__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20250106.235006
+Version: 1.20250109.111908
 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

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

{returnn-1.20250106.235006 → returnn-1.20250109.111908}/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.20250106.235006 → returnn-1.20250109.111908}/returnn/frontend/parameter.py

@@ -4,12 +4,13 @@ Parameter / variable
 
 from __future__ import annotations
 from typing import Optional, Union, TypeVar, Sequence
+from contextlib import contextmanager
 from returnn.tensor import Tensor, Dim
 import returnn.frontend as rf
 from ._backend import global_backend as _global_backend
 
 
-__all__ = ["Parameter"]
+__all__ = ["Parameter", "set_parameter_copy_behavior_ctx"]
 
 
 T = TypeVar("T")
@@ -97,8 +98,11 @@ class Parameter(Tensor[T]):
         self.initial = initial  # use setter
 
     def __copy__(self):
+        if _copy_behavior == "share":
+            return self
         # Should return new copy. https://github.com/rwth-i6/returnn_common/pull/215#issuecomment-1269651064
         # Note that the values are *not* copied, but rather it will use the same param init scheme.
+        assert _copy_behavior == "copy_init"
         res = type(self)(
             dims=self.dims,
             dtype=self.dtype,
@@ -111,10 +115,14 @@ class Parameter(Tensor[T]):
         return res
 
     def __deepcopy__(self, memo=None):
+        if _copy_behavior == "share":
+            return self
+
         # Should return new copy. https://github.com/rwth-i6/returnn_common/pull/215#issuecomment-1269651064
         # Note that the values are *not* copied, but rather it will use the same param init scheme.
         from copy import deepcopy
 
+        assert _copy_behavior == "copy_init"
         res = self.__copy__()
         if isinstance(self.initial, rf.init.ParamInit):
             res.initial = deepcopy(self.initial, memo=memo)  # noqa
@@ -241,3 +249,27 @@ class Parameter(Tensor[T]):
     @non_critical_for_restore.setter
     def non_critical_for_restore(self, value: bool):
         self._non_critical_for_restore = value
+
+
+_copy_behavior = "copy_init"
+
+
+@contextmanager
+def set_parameter_copy_behavior_ctx(mode: str):
+    """
+    Set the copy behavior for parameters.
+
+    :param mode: "copy_init" or "share".
+        "copy_init" (default when not set): When copying a parameter, a new parameter will be created.
+            It is not a shared parameter.
+            It does not copy the values, but it just copies the param init scheme.
+        "share": The parameter will not be copied, it will just return itself.
+    """
+    global _copy_behavior
+    assert mode in ("copy_init", "share")
+    old = _copy_behavior
+    try:
+        _copy_behavior = mode
+        yield
+    finally:
+        _copy_behavior = old

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

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

{returnn-1.20250106.235006 → returnn-1.20250109.111908}/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.20250106.235006/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20250106.235006'
-long_version = '1.20250106.235006+git.7b8fe93'