x-transformers 1.27.10__py3-none-any.whl → 1.27.12__py3-none-any.whl

x_transformers/x_transformers.py

@@ -57,6 +57,9 @@ def maybe(fn):
         return fn(x, *args, **kwargs)
     return inner
 
+def at_most_one_of(*bools):
+    return sum(map(int, bools)) <= 1
+
 class always():
     def __init__(self, val):
         self.val = val
@@ -435,19 +438,15 @@ class RotaryEmbedding(nn.Module):
 
     @autocast(enabled = False)
     def forward(self, t):
-        device, seq_len = self.inv_freq.device, t.shape[-1]
-
-        t = t.type_as(self.inv_freq)
-
-        t = t / self.interpolation_factor
+        max_pos = t.max()+1
 
-        freqs = torch.einsum('i , j -> i j', t, self.inv_freq)
+        freqs = torch.einsum('i , j -> i j', t.typ_as(self.inv_freq), self.inv_freq) / self.interpolation_factor
         freqs = torch.cat((freqs, freqs), dim = -1)
 
         if not exists(self.scale):
             return freqs, 1.
 
-        power = (torch.arange(seq_len, device = device) - (seq_len // 2)) / self.scale_base
+        power = (t - (max_pos // 2)) / self.scale_base
         scale = self.scale ** rearrange(power, 'n -> n 1')
         scale = torch.cat((scale, scale), dim = -1)
 
@@ -463,6 +462,7 @@ def rotate_half(x):
 def apply_rotary_pos_emb(t, freqs, scale = 1):
     rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
     freqs = freqs[-seq_len:, :]
+    scale = scale[-seq_len:, :] if isinstance(scale, torch.Tensor) else scale
 
     if t.ndim == 4 and freqs.ndim == 3:
         freqs = rearrange(freqs, 'b n d -> b 1 n d')
@@ -1264,8 +1264,6 @@ class AttentionLayers(nn.Module):
         # rotary positions
 
         if not exists(rotary_pos_emb) and exists(self.rotary_pos_emb):
-            max_rotary_emb_length = max(list(map(lambda m: (m.shape[1] if exists(m) else 0) + x.shape[1], mems)))
-
             maybe_mem = mems[0] # todo - handle edge case where different layers get different memory lengths. don't think this will ever come up but who knows
             mem_len = maybe_mem.shape[1] if exists(maybe_mem) else 0
 
@@ -1467,7 +1465,8 @@ class ViTransformerWrapper(nn.Module):
     def forward(
         self,
         img,
-        return_embeddings = False
+        return_embeddings = False,
+        return_logits_and_embeddings = False
     ):
         b, p = img.shape[0], self.patch_size
 
@@ -1484,16 +1483,23 @@ class ViTransformerWrapper(nn.Module):
             r = repeat(self.register_tokens, 'n d -> b n d', b = b)
             x, ps = pack((x, r), 'b * d')
 
-        x = self.attn_layers(x)
+        embed = self.attn_layers(x)
 
         if self.has_register_tokens:
-            x, _ = unpack(x, ps, 'b * d')
+            embed, _ = unpack(embed, ps, 'b * d')
+
+        assert at_most_one_of(return_embeddings, return_logits_and_embeddings)
 
         if not exists(self.mlp_head) or return_embeddings:
-            return x
+            return embed
+
+        pooled = embed.mean(dim = -2)
+        logits = self.mlp_head(pooled)
+
+        if not return_logits_and_embeddings:
+            return logits
 
-        x = x.mean(dim = -2)
-        return self.mlp_head(x)
+        return logits, embed
 
 class TransformerWrapper(nn.Module):
     def __init__(

{x_transformers-1.27.10.dist-info → x_transformers-1.27.12.dist-info}/METADATA

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

{x_transformers-1.27.10.dist-info → x_transformers-1.27.12.dist-info}/RECORD

@@ -4,11 +4,11 @@ x_transformers/autoregressive_wrapper.py,sha256=gYKIN5Rm8dMYSTX5yHpg9sPYyZf9rsRT
 x_transformers/continuous.py,sha256=92Wczoaz6dJalix-e3mdIzW0xyRIx3GlBSgsSQOsJeI,6123
 x_transformers/dpo.py,sha256=ek9dgiSs05xeCn8ORceOgKy6LJOnNDw-OJDqxAVLecM,2243
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=TPH5PitIzIBWTQdnO8nlctB8poSMvHkBPWcWFolgZAM,63429
+x_transformers/x_transformers.py,sha256=_mb1-TUDTmnuF-WKwg5VlTYYJthFkFC8Q9OfrRD59sQ,63618
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=ulEPep6i5Hl7H-H9vGfdsmHdprUmK8ajB306jViyV2c,8147
-x_transformers-1.27.10.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.27.10.dist-info/METADATA,sha256=v2ZVeG1yd-HPYFbBWYjNL-q4s74asgt7U8VWg4f9Leg,662
-x_transformers-1.27.10.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-x_transformers-1.27.10.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.27.10.dist-info/RECORD,,
+x_transformers-1.27.12.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.27.12.dist-info/METADATA,sha256=vpgAkW1OgQIzDRe2mZaTCHmIMO5-PUv8DO1FSXCLqy0,662
+x_transformers-1.27.12.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+x_transformers-1.27.12.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.27.12.dist-info/RECORD,,