account if some sample within a batch does not have any msas or templates

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -780,6 +780,11 @@ def forward(
             if exists(msa_mask):
                 msa_mask = einx.get_at('b [s], b sampled -> b sampled', msa_mask, indices)
 
+        # account for no msa
+
+        if exists(msa_mask):
+            has_msa = reduce(msa_mask, 'b s -> b', 'any')
+
         # process msa
 
         msa = self.msa_init_proj(msa)
@@ -806,6 +811,12 @@ def forward(
 
             pairwise_repr = pairwise_block(pairwise_repr = pairwise_repr, mask = mask)
 
+        if exists(msa_mask):
+            pairwise_repr = einx.where(
+                'b, b ..., -> b ...',
+                has_msa, pairwise_repr, 0.
+            )
+
         return pairwise_repr
 
 # pairformer stack
@@ -1033,6 +1044,8 @@ def forward(
 
         v, merged_batch_ps = pack_one(v, '* i j d')
 
+        has_templates = reduce(template_mask, 'b t -> b', 'any')
+
         if exists(mask):
             mask = repeat(mask, 'b n -> (b t) n', t = num_templates)
 
@@ -1058,7 +1071,14 @@ def forward(
 
         avg_template_repr = einx.divide('b i j d, b -> b i j d', num, den.clamp(min = self.eps))
 
-        return self.to_out(avg_template_repr)
+        out = self.to_out(avg_template_repr)
+
+        out = einx.where(
+            'b, b ..., -> b ...',
+            has_templates, out, 0.
+        )
+
+        return out
 
 # diffusion related
 # both diffusion transformer as well as atom encoder / decoder

pyproject.toml

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