torch compile flag in trainer

Author: lucidrains

alphafold3_pytorch/tensor_typing.py

@@ -69,5 +69,6 @@ def __getitem__(self, shapes: str):
     Int,
     Bool,
     typecheck,
+    should_typecheck,
     beartype_isinstance
 ]

alphafold3_pytorch/trainer.py

@@ -10,6 +10,7 @@
 from typing import TypedDict, List, Callable
 
 from alphafold3_pytorch.tensor_typing import (
+    should_typecheck,
     typecheck,
     Int, Bool, Float
 )
@@ -278,7 +279,8 @@ def __init__(
         ema_kwargs: dict = dict(
             use_foreach = True
         ),
-        use_adam_atan2: bool = False
+        use_adam_atan2: bool = False,
+        use_torch_compile: bool = False
     ):
         super().__init__()
 
@@ -288,10 +290,6 @@ def __init__(
         self.fabric = fabric
         fabric.launch()
 
-        # model
-
-        self.model = model
-
         # exponential moving average
 
         self.ema_model = None
@@ -305,6 +303,16 @@ def __init__(
                 **ema_kwargs
             )
 
+        # maybe torch compile
+
+        if use_torch_compile:
+            assert not should_typecheck, f'does not work well with jaxtyping, please invoke your training script with the environment flag `TYPECHECK=False` - ex. `TYPECHECK=False python train_af3.py`'
+            model = torch.compile(model)
+
+        # model
+
+        self.model = model
+
         # optimizer
 
         if not exists(optimizer):

pyproject.toml

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