make returning confidence head logits correct when sampling

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3639,19 +3639,17 @@ def forward(
             if exists(atom_mask):
                 sampled_atom_pos = einx.where('b m, b m c, -> b m c', atom_mask, sampled_atom_pos, 0.)
 
+            if return_confidence_head_logits:
+                assert exists(molecule_atom_indices)
+                pred_atom_pos = einx.get_at('b [m] c, b n -> b n c', sampled_atom_pos, molecule_atom_indices)
+
             if exists(missing_atom_mask) and return_present_sampled_atoms:
                 sampled_atom_pos = sampled_atom_pos[~missing_atom_mask]
 
             if not return_confidence_head_logits:
                 return sampled_atom_pos
 
-            # todo - handle missing atoms
-
-            assert exists(molecule_atom_indices)
-
-            pred_atom_pos = einx.get_at('b [m] c, b n -> b n c', sampled_atom_pos, molecule_atom_indices)
-
-            logits = self.confidence_head(
+            confidence_head_logits = self.confidence_head(
                 single_repr = single.detach(),
                 single_inputs_repr = single_inputs.detach(),
                 pairwise_repr = pairwise.detach(),
@@ -3660,7 +3658,7 @@ def forward(
                 return_pae_logits = True
             )
 
-            return sampled_atom_pos, logits
+            return sampled_atom_pos, confidence_head_logits
 
         # if being forced to return loss, but do not have sufficient information to return losses, just return 0
 
@@ -3826,7 +3824,7 @@ def forward(
 
             pred_atom_pos = einx.get_at('b [m] c, b n -> b n c', denoised_atom_pos, molecule_atom_indices)
 
-            logits = self.confidence_head(
+            ch_logits = self.confidence_head(
                 single_repr = single.detach(),
                 single_inputs_repr = single_inputs.detach(),
                 pairwise_repr = pairwise.detach(),
@@ -3837,19 +3835,19 @@ def forward(
 
             if exists(pae_labels):
                 pae_labels = torch.where(pairwise_mask, pae_labels, ignore)
-                pae_loss = F.cross_entropy(logits.pae, pae_labels, ignore_index = ignore)
+                pae_loss = F.cross_entropy(ch_logits.pae, pae_labels, ignore_index = ignore)
 
             if exists(pde_labels):
                 pde_labels = torch.where(pairwise_mask, pde_labels, ignore)
-                pde_loss = F.cross_entropy(logits.pde, pde_labels, ignore_index = ignore)
+                pde_loss = F.cross_entropy(ch_logits.pde, pde_labels, ignore_index = ignore)
 
             if exists(plddt_labels):
                 plddt_labels = torch.where(mask, plddt_labels, ignore)
-                plddt_loss = F.cross_entropy(logits.plddt, plddt_labels, ignore_index = ignore)
+                plddt_loss = F.cross_entropy(ch_logits.plddt, plddt_labels, ignore_index = ignore)
 
             if exists(resolved_labels):
                 resolved_labels = torch.where(mask, resolved_labels, ignore)
-                resolved_loss = F.cross_entropy(logits.resolved, resolved_labels, ignore_index = ignore)
+                resolved_loss = F.cross_entropy(ch_logits.resolved, resolved_labels, ignore_index = ignore)
 
             confidence_loss = pae_loss + pde_loss + plddt_loss + resolved_loss