x-transformers 2.11.10__tar.gz → 2.11.12__tar.gz

{x_transformers-2.11.10 → x_transformers-2.11.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.10
+Version: 2.11.12
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.11.10 → x_transformers-2.11.12}/pyproject.toml

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

{x_transformers-2.11.10 → x_transformers-2.11.12}/x_transformers/x_transformers.py

@@ -740,11 +740,14 @@ def apply_rotary_pos_emb(t, freqs, scale = 1):
     rot_dim, seq_len, orig_dtype = freqs.shape[-1], t.shape[-2], t.dtype
 
     freqs = freqs[:, -seq_len:, :]
-    scale = scale[:, -seq_len:, :] if isinstance(scale, torch.Tensor) else scale
+    scale = scale[:, -seq_len:, :] if is_tensor(scale) else scale
 
     if t.ndim == 4 and freqs.ndim == 3:
         freqs = rearrange(freqs, 'b n d -> b 1 n d')
 
+        if is_tensor(scale):
+            scale = rearrange(scale, 'b n d -> b 1 n d')
+
     # partial rotary embeddings, Wang et al. GPT-J
     t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
     t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
@@ -3438,6 +3441,7 @@ class TransformerWrapper(Module):
 
         kwargs = dict(
             **kwargs,
+            pos = pos,
             seq_pos_offset = seq_pos_offset,
             seq_start_pos = seq_start_pos,
             input_not_include_cache = input_not_include_cache