x-transformers 1.28.0__py3-none-any.whl → 1.28.2__py3-none-any.whl

x_transformers/x_transformers.py

@@ -1,6 +1,5 @@
 import math
 from random import random
-from typing import Dict
 from packaging import version
 
 import torch
@@ -11,7 +10,7 @@ from torch.cuda.amp import autocast
 from functools import partial, wraps
 from collections import namedtuple
 from dataclasses import dataclass
-from typing import List, Callable, Optional, Union
+from typing import List, Dict, Tuple, Callable, Optional, Union
 
 from einops import rearrange, repeat, reduce, pack, unpack
 from einops.layers.torch import Rearrange
@@ -91,7 +90,10 @@ def l2norm(t, groups = 1):
     t = F.normalize(t, p = 2, dim = -1)
     return rearrange(t, '... g d -> ... (g d)')
 
-def pad_at_dim(t, pad, dim = -1, value = 0.):
+def pad_at_dim(t, pad: Tuple[int, int], dim = -1, value = 0.):
+    if pad == (0, 0):
+        return t
+
     dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
     zeros = ((0, 0) * dims_from_right)
     return F.pad(t, (*zeros, *pad), value = value)
@@ -816,7 +818,7 @@ class Attention(nn.Module):
         return_intermediates = False,
         cache: Optional[Intermediates] = None,
     ):
-        b, n, h, kv_h, head_scale, device, has_context = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, x.device, exists(context)
+        b, n, h, kv_h, head_scale, num_mem_kv, device, has_context = x.shape[0], x.shape[1], self.heads, self.kv_heads, self.head_scale, self.num_mem_kv, x.device, exists(context)
 
         kv_input = default(context, x)
 
@@ -895,9 +897,7 @@ class Attention(nn.Module):
 
         # maybe append memory key / values
 
-        has_mem_kv = self.num_mem_kv > 0
-
-        if has_mem_kv:
+        if num_mem_kv > 0:
             mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b = b), (self.mem_k, self.mem_v))
 
             if self.qk_norm:
@@ -933,6 +933,7 @@ class Attention(nn.Module):
             range_k = torch.arange(j, device = device)
             dist = rearrange(range_q, 'i -> 1 1 i 1') - rearrange(range_k, 'j -> 1 1 1 j')
             max_attend_past_mask = dist > self.max_attend_past
+            max_attend_past_mask = pad_at_dim(max_attend_past_mask, (num_mem_kv, 0), value = False, dim = -1) # handle memory key / values
             masks.append(max_attend_past_mask)
 
         if len(masks) > 0:
@@ -943,11 +944,7 @@ class Attention(nn.Module):
         attn_bias = None
         if exists(rel_pos):
             attn_bias = rel_pos(i, j)
-
-        # append with no bias for memory key / values
-
-        if has_mem_kv:
-            attn_bias = pad_at_dim(attn_bias, (self.num_mem_kv, 0), value = 0.)
+            attn_bias = pad_at_dim(attn_bias, (num_mem_kv, 0), value = 0.) # handle memory key / values
 
         # attention is all we need
 

{x_transformers-1.28.0.dist-info → x_transformers-1.28.2.dist-info}/METADATA

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

{x_transformers-1.28.0.dist-info → x_transformers-1.28.2.dist-info}/RECORD

@@ -4,11 +4,11 @@ x_transformers/autoregressive_wrapper.py,sha256=gYKIN5Rm8dMYSTX5yHpg9sPYyZf9rsRT
 x_transformers/continuous.py,sha256=dpHK4NSMDQAJQ_N3Uj9rip0fYGXyu0QCCO_OfEdbRGs,6192
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 x_transformers/nonautoregressive_wrapper.py,sha256=ys_p8obc7lTeeodCqvkRKxOXQ1C9T3j5Jwr-JbVgnXk,10432
-x_transformers/x_transformers.py,sha256=GvqVKQZRtIldnSWX4V6qE2sWOGruRvBhk4MVit7ZD_M,63897
+x_transformers/x_transformers.py,sha256=GhhRfzxOQoUAqEeT8VnSAtW7wIJ6aW_5DF4LnsqozdQ,64018
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=EN3hxxleTRGYeAz6i4x3U_PrOm9TjxMF3eDhMKGx59E,8575
-x_transformers-1.28.0.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.28.0.dist-info/METADATA,sha256=o1AbarRMIJY_R0gNaEm5SNUWm3YHEesLL2EEy_Uk6gA,661
-x_transformers-1.28.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-x_transformers-1.28.0.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.28.0.dist-info/RECORD,,
+x_transformers-1.28.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.28.2.dist-info/METADATA,sha256=oT95hrc_XiI7dMKF9ATWyUwir3cfSfeD1PFTZF2zpy4,661
+x_transformers-1.28.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+x_transformers-1.28.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.28.2.dist-info/RECORD,,