fix masking for memory key/values exist

lucidrains · Nov 8, 2020 · 499296e · 499296e
1 parent dc8ec28
commit 499296e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 9 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '0.0.26',
+  version = '0.0.27',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

diff --git a/x_transformers/x_transformers.py b/x_transformers/x_transformers.py
@@ -193,10 +193,21 @@ def forward(self, x, context = None, mask = None, context_mask = None, rel_pos =
 
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q_, *kv))
 
+        input_mask = None
+        if any(map(exists, (mask, context_mask))):
+            q_mask = default(mask, lambda: torch.ones((b, n), device = device).bool())
+            k_mask = q_mask if not exists(context) else context_mask
+            k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device = device).bool())
+            q_mask = rearrange(q_mask, 'b i -> b () i ()')
+            k_mask = rearrange(k_mask, 'b j -> b () () j')
+            input_mask = q_mask * k_mask
+
         if self.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))
             k = torch.cat((mem_k, k), dim = -2)
             v = torch.cat((mem_v, v), dim = -2)
+            if exists(input_mask):
+                input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value = True)
 
         dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
 
@@ -206,14 +217,9 @@ def forward(self, x, context = None, mask = None, context_mask = None, rel_pos =
         if exists(rel_pos):
             dots = rel_pos(dots)
 
-        if any(map(exists, (mask, context_mask))):
-            q_mask = default(mask, lambda: torch.ones((b, dots.shape[-2]), device = device).bool())
-            k_mask = default(context_mask, lambda: torch.ones((b, dots.shape[-1]), device = device).bool())
-            q_mask = rearrange(q_mask, 'b i -> b () i ()')
-            k_mask = rearrange(k_mask, 'b j -> b () () j')
-            mask = q_mask * k_mask
-            dots.masked_fill_(mask, float('-inf'))
-            del mask
+        if exists(input_mask):
+            dots.masked_fill_(input_mask, float('-inf'))
+            del input_mask
 
         if self.causal:
             i, j = dots.shape[-2:]