handle rotary and polar positional embeddings with caching when attention layers is not wrapped

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.15.2"
+version = "2.16.0"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

tests/test_x_transformers.py

@@ -1533,7 +1533,7 @@ def test_seq_start_pos_parity():
         input_not_include_cache = True,
         attn_layers = Decoder(
             dim = 32,
-            depth = 2
+            depth = 2,
         )
     )
 
@@ -1568,3 +1568,44 @@ def test_seq_start_pos_parity():
     is_not_masked = torch.arange(seq_len) >= seq_start_pos[:, None]
 
     assert torch.allclose(parallel_logits[is_not_masked], seq_logits[is_not_masked], atol = 1e-5)
+
+@param('pos_emb_type', ('rotary', 'polar'))
+def test_pos_emb_parity(pos_emb_type):
+    pos_emb_kwargs = {f'{pos_emb_type}_pos_emb': True}
+
+    model = Decoder(
+        dim = 128,
+        depth = 1,
+        heads = 4,
+        **pos_emb_kwargs
+    )
+
+    model.eval()
+
+    # parallel
+
+    seq = torch.randn(2, 15, 128)
+
+    parallel_logits = model(seq)
+
+    # prompt pass
+
+    prompt = seq[:, :10]
+    cache = None
+    all_seq_logits = []
+
+    logits, cache = model(prompt, cache = cache, return_hiddens = True)
+    all_seq_logits.append(logits[:, -1:])
+
+    # sequential
+
+    for i in range(4):
+        input_embeds = seq[:, 10 + i : 10 + i + 1]
+        logits, cache = model(input_embeds, cache = cache, return_hiddens = True)
+        all_seq_logits.append(logits[:, -1:])
+
+    seq_logits = torch.cat(all_seq_logits, dim = 1)
+
+    parallel_logits_without_prompt = parallel_logits[:, 9 : 14]
+
+    assert torch.allclose(seq_logits, parallel_logits_without_prompt, atol = 1e-5)

x_transformers/x_transformers.py

@@ -2710,7 +2710,7 @@ def forward(
         mems = None,
         mem_masks = None,
         seq_start_pos: Tensor | None = None,
-        seq_pos_offset: int = 0,
+        seq_pos_offset = None,
         cache: LayerIntermediates | None = None,
         input_not_include_cache = False,
         cache_age = 1,
@@ -2738,6 +2738,12 @@ def forward(
         assert not (self.cross_attend ^ exists(context)), 'context must be passed in if cross_attend is set to True'
         assert not (exists(condition) ^ self.need_condition), 'condition needs to be passed in if using adaptive layernorm or vice versa'
 
+        # handle seq pos offset if not passed in from wrapper
+        # default to 0, but if cache is detected, set appropriate for the relative positional embeddings
+
+        if not exists(seq_pos_offset):
+            seq_pos_offset = cache.cache_length if exists(cache) else 0
+
         # handle condition
 
         if exists(condition):