test on the atom inputs collation function individually

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -42,6 +42,7 @@
 from alphafold3_pytorch.trainer import (
     Trainer,
     DataLoader,
+    collate_af3_inputs
 )
 
 from alphafold3_pytorch.configs import (

tests/test_trainer.py

@@ -14,6 +14,7 @@
     DataLoader,
     Trainer,
     ConductorConfig,
+    collate_af3_inputs,
     create_trainer_from_yaml,
     create_trainer_from_conductor_yaml,
     create_alphafold3_from_yaml
@@ -186,6 +187,38 @@ def test_trainer():
 
     alphafold3 = Alphafold3.init_and_load('./some/nested/folder2/training.pt')
 
+# test use of collation fn outside of trainer
+
+def test_collate_fn():
+    alphafold3 = Alphafold3(
+        dim_atom_inputs = 77,
+        dim_template_feats = 44,
+        num_dist_bins = 38,
+        confidence_head_kwargs = dict(
+            pairformer_depth = 1
+        ),
+        template_embedder_kwargs = dict(
+            pairformer_stack_depth = 1
+        ),
+        msa_module_kwargs = dict(
+            depth = 1
+        ),
+        pairformer_stack = dict(
+            depth = 1
+        ),
+        diffusion_module_kwargs = dict(
+            atom_encoder_depth = 1,
+            token_transformer_depth = 1,
+            atom_decoder_depth = 1,
+        ),
+    )
+
+    dataset = MockAtomDataset(1)
+
+    batched_atom_inputs = collate_af3_inputs([dataset[0]])
+
+    _, breakdown = alphafold3(**batched_atom_inputs, return_loss_breakdown = True)
+
 # test creating trainer + alphafold3 from config
 
 def test_trainer_config():