returnn 1.20241209.230715__tar.gz → 1.20241210.111636__tar.gz

{returnn-1.20241209.230715 → returnn-1.20241210.111636}/PKG-INFO

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

returnn-1.20241210.111636/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241210.111636'
+long_version = '1.20241210.111636+git.b091bd6'

{returnn-1.20241209.230715 → returnn-1.20241210.111636}/returnn/frontend/attention.py

@@ -325,7 +325,7 @@ class RotaryPosCausalSelfAttention(CausalSelfAttention):
         q = _apply_rope(
             q,
             (
-                rf.gather(pos_enc, axis=hist_dim, indices=hist_dim.dyn_size_ext - 1)
+                rf.gather(pos_enc, axis=hist_dim, indices=hist_dim.get_size_tensor() - 1)
                 if axis == single_step_dim
                 else rf.replace_dim(pos_enc, in_dim=hist_dim, out_dim=axis)[0]
             ),
@@ -503,6 +503,8 @@ def _rel_pos_enc_shift(x: Tensor, axis: Dim, pos_emb_spatial_dim: Dim, hist_dim:
     :param hist_dim: T' (equal to T but separate dim)
     :return: [B,H,T,T']
     """
+    if pos_emb_spatial_dim == hist_dim:  # happens for single_step_dim
+        return x  # no shift needed
     batch_dims = x.remaining_dims((axis, pos_emb_spatial_dim))
     x_padded, (pos_emb_spatial_dim_,) = rf.pad(
         x, axes=[pos_emb_spatial_dim], padding=[(1, 0)], value=0.0
@@ -604,6 +606,7 @@ class RelPosCausalSelfAttention(CausalSelfAttention):
             pos_emb, pos_emb_spatial_dim = relative_positional_encoding(
                 query_spatial_dim=axis, key_value_spatial_dim=hist_dim, feat_dim=self.pos_emb_feat_dim
             )
+        # pos_emb_spatial_dim is 2*time1-1 if axis!=single_step_dim, else time1
         if self.pos_emb_dropout:
             pos_emb = rf.dropout(pos_emb, self.pos_emb_dropout)
         if self.linear_pos is not None:
@@ -850,7 +853,8 @@ def _make_indices(
     if query_spatial_dim == single_step_dim:
         indices = kv_pos_vec
         out_spatial_dim = key_value_spatial_dim
-        assert query_offset is None  # not sure if any custom query offset makes sense?
+        # not sure if any custom query offset makes sense?
+        assert query_offset is None or (isinstance(query_offset, int) and query_offset == 0)
         # Assume the kv are the accumulated history, and query is cur frame of it,
         # corresponding to the last frame of the kv.
         query_offset = key_value_spatial_dim.get_size_tensor() - 1

{returnn-1.20241209.230715 → returnn-1.20241210.111636}/returnn.egg-info/PKG-INFO

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

{returnn-1.20241209.230715 → returnn-1.20241210.111636}/tests/test_rf_attention.py

@@ -3,7 +3,7 @@ RETURNN frontend (returnn.frontend) tests
 """
 
 from __future__ import annotations
-from typing import Tuple
+from typing import Union, Tuple
 import numpy as np
 import numpy.testing
 import _setup_test_env  # noqa
@@ -440,6 +440,84 @@ def test_rope_causal_self_att():
     print("  all matched!")
 
 
+def test_causal_self_att_variants_single_step_vs_full_seq():
+    from returnn.tensor import single_step_dim
+
+    time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
+    in_dim = Dim(7 * 2, name="in")
+    extern_data = TensorDict(
+        {
+            "data": Tensor("data", [batch_dim, time_dim, in_dim], dtype="float32"),
+        }
+    )
+
+    # noinspection PyShadowingNames
+    def _forward_step(*, model: Union[rf.CausalSelfAttention], extern_data: TensorDict):
+        x = extern_data["data"]
+
+        out_seq_level, _ = model(x, axis=time_dim)
+        out_seq_level.mark_as_output("out_seq_level", shape=[batch_dim, time_dim, model.out_dim])
+
+        out_seq_level_explicit_initial_state, _ = model(
+            x, axis=time_dim, state=model.default_initial_state(batch_dims=[batch_dim])
+        )
+        out_seq_level_explicit_initial_state.mark_as_output(
+            "out_seq_level_explicit_initial_state", shape=[batch_dim, time_dim, model.out_dim]
+        )
+
+        def _body(
+            _x: Tensor, _state: Union[rf.CausalSelfAttentionState]
+        ) -> Tuple[Tensor, Union[rf.CausalSelfAttentionState]]:
+            return model(_x, axis=single_step_dim, state=_state)
+
+        out_single_steps, _, _ = rf.scan(
+            spatial_dim=time_dim,
+            xs=x,
+            body=_body,
+            ys=Tensor("y", dims=[batch_dim, model.out_dim], dtype="float32"),
+            initial=model.default_initial_state(batch_dims=[batch_dim]),
+        )
+        out_single_steps.mark_as_output("out_single_steps", shape=[batch_dim, time_dim, model.out_dim])
+
+    common_opts = dict(
+        in_dim=in_dim,
+        proj_dim=Dim(5, name="out"),
+        key_dim_total=Dim(21 * 2, name="key-dim-total"),
+        value_dim_total=Dim(33, name="value-dim-total"),
+        num_heads=3,
+    )
+
+    def _make_causal_self_att(**_kwargs):
+        return rf.CausalSelfAttention(**common_opts)
+
+    def _make_rope_causal_self_att(**_kwargs):
+        return rf.RotaryPosCausalSelfAttention(**common_opts)
+
+    def _make_rel_pos_causal_self_att(**_kwargs):
+        return rf.RelPosCausalSelfAttention(**common_opts)
+
+    models = [_make_causal_self_att, _make_rope_causal_self_att, _make_rel_pos_causal_self_att]
+
+    for get_model in models:
+        print("> Testing model:", get_model.__name__)
+        res = run_model(
+            extern_data,
+            get_model,
+            _forward_step,
+            # TF needs TensorArray unstack, not implemented yet
+            test_tensorflow=False,
+        )
+
+        # Check that the single-step and the seq-level output are the same.
+        res_seq_level = res.data["out_seq_level"].raw_tensor
+        for key in ["out_seq_level_explicit_initial_state", "out_single_steps"]:
+            res_other = res.data[key].raw_tensor
+            assert res_seq_level.shape == res_other.shape
+            numpy.testing.assert_allclose(
+                res_other, res_seq_level, atol=1e-5, rtol=1e-5, err_msg=f"output {key} differs"
+            )
+
+
 def test_relative_positional_encoding():
     time_dim = Dim(Tensor("time", [batch_dim], dtype="int32"))
     in_dim = Dim(8, name="in")

returnn-1.20241209.230715/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241209.230715'
-long_version = '1.20241209.230715+git.55c6a02'