fix msa_mask gathering when subsampling

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1142,13 +1142,12 @@ def forward(
             # msa = einx.get_at('b [s] n dm, b sampled -> b sampled n dm', msa, indices)
 
             msa, unpack_one = pack_one(msa, 'b s *')
-            indices = repeat(indices, 'b sampled -> b sampled d', d = msa.shape[-1])
-            msa = msa.gather(1, indices)
+            msa_indices = repeat(indices, 'b sampled -> b sampled d', d = msa.shape[-1])
+            msa = msa.gather(1, msa_indices)
             msa = unpack_one(msa)
 
             if exists(msa_mask):
                 # msa_mask = einx.get_at('b [s], b sampled -> b sampled', msa_mask, indices)
-
                 msa_mask = msa_mask.gather(1, indices)
 
         # account for no msa

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.105"
+version = "0.2.106"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },

tests/test_af3.py

@@ -228,6 +228,7 @@ def test_msa_module(
     pairwise = torch.randn(2, 16, 16, 128).requires_grad_()
     msa = torch.randn(2, 7, 16, 64)
     mask = torch.randint(0, 2, (2, 16)).bool()
+    msa_mask = torch.randint(0, 2, (2, 7)).bool()
 
     msa_module = MSAModule(
         checkpoint = checkpoint,
@@ -238,7 +239,8 @@ def test_msa_module(
         msa = msa,
         single_repr = single,
         pairwise_repr = pairwise,
-        mask = mask
+        mask = mask,
+        msa_mask = msa_mask
     )
 
     assert pairwise.shape == pairwise_out.shape