always use mask if sparse axial or conv attention needs to be padded

Author: lucidrains

dalle_pytorch/attention.py

@@ -97,9 +97,9 @@ def forward(self, x, mask = None):
 
         if n < seq_len:
             padding = seq_len - n
+            mask = default(mask, lambda: torch.ones(b, n, device = device).bool())
             x = F.pad(x, (0, 0, 0, padding), value = 0)
-            if exists(mask):
-                mask = F.pad(x, (0, padding), value = False)
+            mask = F.pad(x, (0, padding), value = False)
 
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), qkv)
@@ -196,10 +196,9 @@ def forward(self, x, mask = None):
 
         if n < seq_len:
             padding = seq_len - n
+            mask = default(mask, lambda: torch.ones(b, n, device = device).bool())
             x = F.pad(x, (0, 0, 0, padding), value = 0)
-
-            if exists(mask):
-                mask = F.pad(x, (0, padding), value = False)
+            mask = F.pad(x, (0, padding), value = False)
 
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), qkv)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'dalle-pytorch',
   packages = find_packages(),
-  version = '0.0.55',
+  version = '0.0.56',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',