add an assert to make sure total sequence length of MSA or AA does no…

…t exceed the set length of the sparse attention kernel

alphafold2_pytorch/alphafold2.py

@@ -174,13 +174,15 @@ def __init__(
 
     def forward(self, x, mask = None):
         device, h = x.device, self.heads
-        n = x.shape[1]
+        b, n = x.shape[:2]
+        assert n <= self.seq_len, f'either the AA sequence length {n} or the total MSA length {n} exceeds the allowable sequence length {self.seq_len} for sparse attention, set by `max_seq_len`'
 
         remainder = x.shape[1] % self.block_size
 
         if remainder > 0:
             padding = self.block_size - remainder
             x = F.pad(x, (0, 0, 0, padding), value = 0)
+            mask = torch.ones(b, n, device = device).bool()
             mask = F.pad(mask, (0, padding), value = False)
 
         q, k, v = (self.to_q(x), *self.to_kv(x).chunk(2, dim = -1))

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'alphafold2-pytorch',
   packages = find_packages(),
-  version = '0.0.17',
+  version = '0.0.18',
   license='MIT',
   description = 'AlphaFold2 - Pytorch',
   author = 'Phil Wang, Eric Alcaide',