jaxtyping unable to help here, so just do asserts on shape manually

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -138,6 +138,9 @@ def log(t, eps = 1e-20):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def compact(*args):
+    return tuple(filter(exists, args))
+
 # tensor helpers
 
 def max_neg_value(t: Tensor):
@@ -3905,6 +3908,9 @@ def forward(
 
             # cross entropy losses
 
+            assert all([t.shape[-1] for t in compact(pde_labels, plddt_labels, resolved_labels)])
+            assert pde_labels.shape[-1] == ch_logits.pde.shape[-1]
+
             if exists(pae_labels):
                 pae_labels = torch.where(label_pairwise_mask, pae_labels, ignore)
                 pae_loss = F.cross_entropy(ch_logits.pae, pae_labels, ignore_index = ignore)

pyproject.toml

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