small cleanup

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -124,6 +124,9 @@ def pack_one(t, pattern):
 def unpack_one(t, ps, pattern):
     return unpack(t, ps, pattern)[0]
 
+def exclusive_cumsum(t, dim = -1):
+    return t.cumsum(dim = dim) - t
+
 # decorators
 
 def maybe(fn):
@@ -262,8 +265,7 @@ def repeat_consecutive_with_lens(
     window_size = mask.shape[-1]
     arange = torch.arange(window_size, device = device)
 
-    cumsum_len = lens.cumsum(dim = -1)
-    offsets = F.pad(cumsum_len, (1, -1), value = 0)
+    offsets = exclusive_cumsum(lens)
     indices = einx.add('w, b n -> b n w', arange, offsets)
 
     # create output tensor + a sink position on the very right (index max_len)
@@ -3270,7 +3272,7 @@ def forward(
         # handle offsets for molecule atom indices
 
         if exists(molecule_atom_indices):
-            molecule_atom_indices = molecule_atom_indices + F.pad(molecule_atom_lens.cumsum(dim = -1), (1, -1), value = 0)
+            molecule_atom_indices = molecule_atom_indices + exclusive_cumsum(molecule_atom_lens)
 
         # get atom sequence length and molecule sequence length depending on whether using packed atomic seq
 

pyproject.toml

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

tests/test_trainer.py

@@ -174,11 +174,11 @@ def test_trainer():
     # saving and loading from trainer
 
     trainer.save('./some/nested/folder2/training.pt', overwrite = True)
-    trainer.load('./some/nested/folder2/training.pt')
+    trainer.load('./some/nested/folder2/training.pt', strict = False)
 
     # allow for only loading model, needed for fine-tuning logic
 
-    trainer.load('./some/nested/folder2/training.pt', only_model = True)
+    trainer.load('./some/nested/folder2/training.pt', only_model = True, strict = False)
 
     # also allow for loading Alphafold3 directly from training ckpt