fix to tests

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3071,9 +3071,7 @@ def __init__(
         # residue or nucleotide modifications
 
         num_molecule_mods = default(num_molecule_mods, 0)
-
         has_molecule_mod_embeds = num_molecule_mods > 0
-        self.num_molecule_mods = num_molecule_mods
 
         if has_molecule_mod_embeds:
             self.molecule_mod_embeds = nn.Embedding(num_molecule_mods, dim_single)
@@ -3233,6 +3231,11 @@ def __init__(
 
         self.register_buffer('zero', torch.tensor(0.), persistent = False)
 
+        # some shorthand for jaxtyping
+
+        self.dapi = self.dim_atompair_inputs
+        self.dai = self.dim_atom_inputs
+
     @property
     def device(self):
         return self.zero.device
@@ -3309,16 +3312,16 @@ def init_and_load(
     def forward(
         self,
         *,
-        atom_inputs: Float['b m dai'],
-        atompair_inputs: Float['b m m dapi'] | Float['b nw w1 w2 dapi'],
+        atom_inputs: Float['b m {self.dai}'],
+        atompair_inputs: Float['b m m {self.dapi}'] | Float['b nw w1 w2 {self.dapi}'],
         additional_molecule_feats: Int[f'b n {ADDITIONAL_MOLECULE_FEATS}'],
         is_molecule_types: Bool[f'b n {IS_MOLECULE_TYPES}'],
         molecule_atom_lens: Int['b n'],
         molecule_ids: Int['b n'],
         additional_token_feats: Float['b n {self.dim_additional_token_feats}'] | None = None,
         atom_ids: Int['b m'] | None = None,
         atompair_ids: Int['b m m'] | Int['b nw w1 w2'] | None = None,
-        is_molecule_mod: Bool['b n {self.num_molecule_mods}'] | None = None,
+        is_molecule_mod: Bool['b n num_mods'] | None = None,
         atom_mask: Bool['b m'] | None = None,
         atom_parent_ids: Int['b m'] | None = None,
         token_bonds: Bool['b n n'] | None = None,

pyproject.toml

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

tests/test_af3.py

@@ -473,6 +473,7 @@ def test_alphafold3(
         atoms_per_window = atoms_per_window,
         dim_template_feats = 44,
         num_dist_bins = 38,
+        num_molecule_mods = num_molecule_mods,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
         ),