able to override default atom field padding values during collation, set molecule_atom_lens to 0

Author: lucidrains

alphafold3_pytorch/inputs.py

@@ -159,7 +159,15 @@ def inner(x, *args, **kwargs):
 # atom level, what Alphafold3 accepts
 
 UNCOLLATABLE_ATOM_INPUT_FIELDS = {'filepath'}
-ATOM_INPUT_EXCLUDE_MODEL_FIELDS = {'filepath', 'chains'}
+
+ATOM_INPUT_EXCLUDE_MODEL_FIELDS = {
+    'filepath',
+    'chains'
+}
+
+ATOM_DEFAULT_PAD_VALUES = dict(
+    molecule_atom_lens = 0
+)
 
 @typecheck
 @dataclass

alphafold3_pytorch/trainer.py

@@ -27,6 +27,7 @@
     PDBInput,
     maybe_transform_to_atom_inputs,
     UNCOLLATABLE_ATOM_INPUT_FIELDS,
+    ATOM_DEFAULT_PAD_VALUES,
 )
 
 from alphafold3_pytorch.data import (
@@ -171,7 +172,9 @@ def collate_inputs_to_batched_atom_input(
 
         # use -1 for padding int values, for assuming int are labels - if not, handle within alphafold3
 
-        if dtype in (torch.int, torch.long):
+        if key in ATOM_DEFAULT_PAD_VALUES:
+            pad_value = ATOM_DEFAULT_PAD_VALUES[key]
+        elif dtype in (torch.int, torch.long):
             pad_value = int_pad_value
         elif dtype == torch.bool:
             pad_value = False

pyproject.toml

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