just default to serial true and save a bunch of people headache

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -1768,7 +1768,7 @@ def __init__(
         attn_num_memory_kv = False,
         trans_expansion_factor = 2,
         num_register_tokens = 0,
-        serial = False,
+        serial = True,
         add_residual = True,
         use_linear_attn = False,
         checkpoint = False,
@@ -2112,7 +2112,7 @@ def __init__(
         token_transformer_heads = 16,
         atom_decoder_depth = 3,
         atom_decoder_heads = 4,
-        serial = False,
+        serial = True,
         atom_encoder_kwargs: dict = dict(),
         atom_decoder_kwargs: dict = dict(),
         token_transformer_kwargs: dict = dict(),
@@ -4276,7 +4276,6 @@ def __init__(
             token_transformer_heads = 16,
             atom_decoder_depth = 3,
             atom_decoder_heads = 4,
-            serial = True # believe they have an error on Algorithm 23. lacking a residual - default to serial architecture until further news
         ),
         edm_kwargs: dict = dict(
             sigma_min = 0.002,

pyproject.toml

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

tests/test_af3.py

@@ -683,12 +683,22 @@ def test_alphafold3_without_msa_and_templates():
             depth = 1
         ),
         pairformer_stack = dict(
+            checkpoint = True,
             depth = 2
         ),
         diffusion_module_kwargs = dict(
-            atom_encoder_depth = 1,
-            token_transformer_depth = 1,
-            atom_decoder_depth = 1,
+            atom_encoder_depth = 2,
+            atom_encoder_kwargs = dict(
+                checkpoint = True,
+            ),
+            token_transformer_depth = 2,
+            token_transformer_kwargs = dict(
+                checkpoint = True,
+            ),
+            atom_decoder_depth = 2,
+            atom_decoder_kwargs = dict(
+                checkpoint = True,
+            ),
         ),
     )