x-transformers 1.22.10__py3-none-any.whl → 1.22.12__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -28,8 +28,7 @@ def eval_decorator(fn):
 
 # for variable lengthed prefixes
 
-def align(t, lens, pad_id = 0, left = False, right = False):
-    assert left ^ right
+def align_right(t, lens, pad_id = 0):
     batch, seq_len, device, dtype = *t.shape, t.device, t.dtype
 
     assert lens.ndim == 1 and lens.shape[0] == batch
@@ -41,14 +40,9 @@ def align(t, lens, pad_id = 0, left = False, right = False):
     batch_arange = torch.arange(batch, device = device, dtype = torch.long)[..., None]
     prompt_len_arange = torch.arange(seq_len, device = device, dtype = torch.long)
 
-    if left:
-        padding = (0, max_pad_len)
-        offset = pad_lens
-    elif right:
-        padding = (max_pad_len, 0)
-        offset = max_pad_len - pad_lens
+    t = F.pad(t, (max_pad_len, 0), value = 0)
+    offset = max_pad_len - pad_lens
 
-    t = F.pad(t, padding, value = 0)
     aligned = t[batch_arange, prompt_len_arange + offset[..., None]]
     return aligned
 
@@ -157,7 +151,7 @@ class AutoregressiveWrapper(Module):
 
         seq_start_pos = None
         if exists(prompt_lens):
-            prompts = align(prompts, prompt_lens, pad_id = self.pad_value, right = True)
+            prompts = align_right(prompts, prompt_lens, pad_id = self.pad_value)
             seq_start_pos = t - prompt_lens
 
         # output from which sampled tokens appended to
@@ -244,11 +238,6 @@ class AutoregressiveWrapper(Module):
                     out = out.masked_fill(mask, self.pad_value)
                     break
 
-        # if variable lengthed, needs to realign
-
-        if exists(prompt_lens):
-            out = align(out, prompt_lens, pad_id = self.pad_value, left = True)
-
         out = out[:, t:]
 
         out, = unpack(out, ps, '* n')

x_transformers/x_transformers.py

@@ -434,7 +434,6 @@ class RotaryEmbedding(nn.Module):
 
         freqs = torch.einsum('... i , j -> ... i j', t, self.inv_freq)
         freqs = torch.cat((freqs, freqs), dim = -1)
-        freqs = rearrange(freqs, '... i j -> ... 1 i j')
 
         if not exists(self.scale):
             return freqs, 1.
@@ -455,6 +454,9 @@ def apply_rotary_pos_emb(t, freqs, scale = 1):
     rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
     freqs = freqs[..., -seq_len:, :]
 
+    if t.ndim == 4 and freqs.ndim == 3:
+        freqs = rearrange(freqs, '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)

{x_transformers-1.22.10.dist-info → x_transformers-1.22.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.22.10
+Version: 1.22.12
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.22.10.dist-info → x_transformers-1.22.12.dist-info}/RECORD

@@ -1,12 +1,12 @@
 x_transformers/__init__.py,sha256=FDb654rUx8FpXRd76B8q0diH8I7q-ZjTWEtEJ4UM21Y,701
 x_transformers/attend.py,sha256=xPa6RjnMDsc1jKliQdThETMTQeRX3ycmAlw5pgzLIf4,12605
-x_transformers/autoregressive_wrapper.py,sha256=hz3qp_Hmt4oR2MfVw-0ctFBLHW4Liu8XHlkd-rr8O48,8818
+x_transformers/autoregressive_wrapper.py,sha256=uUnwXP2uZ4oJSN4EVXfWQormKWv8c6yzrE5tDZUjSag,8480
 x_transformers/continuous_autoregressive_wrapper.py,sha256=pTiDqu6JRUlnQJQp_xHATYHy0lgSd6ERLqyiFO3pC-4,1575
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=0DeRGThrBtsepw-5B6kvv1_ntA6LCkCcIMU3uMedMm4,59038
+x_transformers/x_transformers.py,sha256=mYveA7PqRUZg9-82ALFBpuhTfQirfCF5rxL6EUCdU5I,59075
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
-x_transformers-1.22.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.22.10.dist-info/METADATA,sha256=k5npGKZe99UhSIVByOIYBusEljaaocy17rszyQP7om0,662
-x_transformers-1.22.10.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.22.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.22.10.dist-info/RECORD,,
+x_transformers-1.22.12.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.22.12.dist-info/METADATA,sha256=lissCQf2eUs5Oalsk0PRw0gwTiybipzr9P5YRybgZdQ,662
+x_transformers-1.22.12.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.22.12.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.22.12.dist-info/RECORD,,