complete the sample without replacement logic for msa

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -696,10 +696,13 @@ def __init__(
         msa_pwa_dropout_row_prob = 0.15,
         msa_pwa_heads = 8,
         msa_pwa_dim_head = 32,
-        pairwise_block_kwargs: dict = dict()
+        pairwise_block_kwargs: dict = dict(),
+        max_num_msa: int | None = None
     ):
         super().__init__()
 
+        self.max_num_msa = default(max_num_msa, float('inf'))  # cap the number of MSAs, will do sample without replacement if exceeds
+
         self.msa_init_proj = LinearNoBias(dim_msa_input, dim_msa) if exists(dim_msa_input) else nn.Identity()
 
         self.single_to_msa_feats = LinearNoBias(dim_single, dim_msa)
@@ -752,6 +755,24 @@ def forward(
         msa_mask: Bool['b s'] | None = None,
     ) -> Float['b n n dp']:
 
+        batch, num_msa, device = *msa.shape[:2], msa.device
+
+        # sample without replacement
+
+        if num_msa > self.max_num_msa:
+            rand = torch.randn((batch, num_msa), device = device)
+
+            if exists(msa_mask):
+                rand.masked_fill_(~msa_mask, max_neg_value(msa))
+
+            indices = rand.topk(self.max_num_msa, dim = -1).indices
+
+            msa = einx.get_at('b [s] n dm, b sampled -> b sampled n dm', msa, indices)
+
+            if exists(msa_mask):
+                msa_mask = einx.get_at('b [s], b sampled -> b sampled', msa_mask, indices)
+
+        # process msa
 
         msa = self.msa_init_proj(msa)
 

pyproject.toml

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

tests/test_af3.py

@@ -132,7 +132,9 @@ def test_msa_module():
     msa = torch.randn(2, 7, 16, 64)
     mask = torch.randint(0, 2, (2, 16)).bool()
 
-    msa_module = MSAModule()
+    msa_module = MSAModule(
+        max_num_msa = 3 # will randomly select 3 out of the MSAs, accounting for mask, using sample without replacement
+    )
 
     pairwise_out = msa_module(
         msa = msa,