rename fn (#140)

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -294,7 +294,7 @@ def mean_pool_with_lens(
     return avg
 
 @typecheck
-def repeat_consecutive_with_lens(
+def batch_repeat_interleave(
     feats: Float['b n ...'] | Bool['b n ...'] | Bool['b n'] | Int['b n'],
     lens: Int['b n'],
     mask_value: float | int | bool | None = None,
@@ -2280,7 +2280,7 @@ def forward(
 
         single_repr_cond = self.single_repr_to_atom_feat_cond(conditioned_single_repr)
 
-        single_repr_cond = repeat_consecutive_with_lens(single_repr_cond, molecule_atom_lens)
+        single_repr_cond = batch_repeat_interleave(single_repr_cond, molecule_atom_lens)
         single_repr_cond = pad_or_slice_to(single_repr_cond, length = atom_feats_cond.shape[1], dim = 1)
 
         atom_feats_cond = single_repr_cond + atom_feats_cond
@@ -2299,7 +2299,7 @@ def forward(
         indices = torch.arange(seq_len, device = device)
         indices = repeat(indices, 'n -> b n', b = batch_size)
 
-        indices = repeat_consecutive_with_lens(indices, molecule_atom_lens)
+        indices = batch_repeat_interleave(indices, molecule_atom_lens)
         indices = pad_or_slice_to(indices, atom_seq_len, dim = -1)
         indices = pad_and_window(indices, w)
 
@@ -2392,7 +2392,7 @@ def forward(
 
         atom_decoder_input = self.tokens_to_atom_decoder_input_cond(tokens)
 
-        atom_decoder_input = repeat_consecutive_with_lens(atom_decoder_input, molecule_atom_lens)
+        atom_decoder_input = batch_repeat_interleave(atom_decoder_input, molecule_atom_lens)
         atom_decoder_input = pad_or_slice_to(atom_decoder_input, length = atom_feats_skip.shape[1], dim = 1)
 
         atom_decoder_input = atom_decoder_input + atom_feats_skip
@@ -2707,7 +2707,7 @@ def forward(
         if exists(is_molecule_types):
             is_nucleotide_or_ligand_fields = is_molecule_types.unbind(dim = -1)
 
-            is_nucleotide_or_ligand_fields = tuple(repeat_consecutive_with_lens(t, molecule_atom_lens) for t in is_nucleotide_or_ligand_fields)
+            is_nucleotide_or_ligand_fields = tuple(batch_repeat_interleave(t, molecule_atom_lens) for t in is_nucleotide_or_ligand_fields)
             is_nucleotide_or_ligand_fields = tuple(pad_or_slice_to(t, length = align_weights.shape[-1], dim = -1) for t in is_nucleotide_or_ligand_fields)
 
             _, atom_is_dna, atom_is_rna, atom_is_ligand, _ = is_nucleotide_or_ligand_fields
@@ -3429,13 +3429,13 @@ def forward(
         # handle maybe atom level resolution
 
         if self.atom_resolution:
-            single_repr = repeat_consecutive_with_lens(single_repr, molecule_atom_lens)
+            single_repr = batch_repeat_interleave(single_repr, molecule_atom_lens)
 
-            pairwise_repr = repeat_consecutive_with_lens(pairwise_repr, molecule_atom_lens)
+            pairwise_repr = batch_repeat_interleave(pairwise_repr, molecule_atom_lens)
 
             molecule_atom_lens = repeat(molecule_atom_lens, 'b ... -> (b r) ...', r = pairwise_repr.shape[1])
             pairwise_repr, unpack_one = pack_one(pairwise_repr, '* n d')
-            pairwise_repr = repeat_consecutive_with_lens(pairwise_repr, molecule_atom_lens)
+            pairwise_repr = batch_repeat_interleave(pairwise_repr, molecule_atom_lens)
             pairwise_repr = unpack_one(pairwise_repr)
 
             interatomic_dist = torch.cdist(pred_atom_pos, pred_atom_pos, p = 2)
@@ -3744,8 +3744,8 @@ def forward(
         valid_indices = torch.ones_like(indices).bool()
 
         # valid_indices at padding position has value False
-        indices = repeat_consecutive_with_lens(indices, molecule_atom_lens)
-        valid_indices = repeat_consecutive_with_lens(valid_indices, molecule_atom_lens)
+        indices = batch_repeat_interleave(indices, molecule_atom_lens)
+        valid_indices = batch_repeat_interleave(valid_indices, molecule_atom_lens)
 
         if exists(atom_mask):
             valid_indices = valid_indices * atom_mask
@@ -3811,8 +3811,8 @@ def compute_full_complex_metric(
         valid_indices = torch.ones_like(indices).bool()
 
         # valid_indices at padding position has value False
-        indices = repeat_consecutive_with_lens(indices, molecule_atom_lens)
-        valid_indices = repeat_consecutive_with_lens(valid_indices, molecule_atom_lens)
+        indices = batch_repeat_interleave(indices, molecule_atom_lens)
+        valid_indices = batch_repeat_interleave(valid_indices, molecule_atom_lens)
 
         # broadcast is_molecule_types to atom
 
@@ -4265,8 +4265,8 @@ def compute_weighted_lddt(
         batch_size = pred_coords.shape[0]
 
         # broadcast asym_id and is_molecule_types to atom level
-        atom_asym_id = repeat_consecutive_with_lens(asym_id, molecule_atom_lens, mask_value=-1)
-        atom_is_molecule_types = repeat_consecutive_with_lens(is_molecule_types, molecule_atom_lens)
+        atom_asym_id = batch_repeat_interleave(asym_id, molecule_atom_lens, mask_value=-1)
+        atom_is_molecule_types = batch_repeat_interleave(is_molecule_types, molecule_atom_lens)
 
         weighted_lddt = torch.zeros(batch_size, device=device)
 

alphafold3_pytorch/utils/model_utils.py

@@ -374,7 +374,7 @@ def mean_pool_with_lens(
 
 
 @typecheck
-def repeat_consecutive_with_lens(
+def batch_repeat_interleave(
     feats: Float["b n ..."] | Bool["b n"] | Int["b n"],  # type: ignore
     lens: Int["b n"],  # type: ignore
 ) -> Float["b m ..."] | Bool["b m"] | Int["b m"]:  # type: ignore

pyproject.toml

@@ -1,9 +1,10 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.101"
+version = "0.2.102"
 description = "Alphafold 3 - Pytorch"
 authors = [
-    { name = "Phil Wang", email = "[email protected]" }
+    { name = "Phil Wang", email = "[email protected]" },
+    { name = "Alex Morehead", email = "[email protected]"}
 ]
 readme = "README.md"
 requires-python = ">= 3.8"

tests/test_af3.py

@@ -43,7 +43,7 @@
 
 from alphafold3_pytorch.alphafold3 import (
     mean_pool_with_lens,
-    repeat_consecutive_with_lens,
+    batch_repeat_interleave,
     full_pairwise_repr_to_windowed,
     get_cid_molecule_type,
 )
@@ -75,10 +75,10 @@ def test_mean_pool_with_lens():
 
     assert torch.allclose(pooled, torch.tensor([[[1.], [2.], [1.]]]))
 
-def test_repeat_consecutive_with_lens():
+def test_batch_repeat_interleave():
     seq = torch.tensor([[[1.], [2.], [4.]], [[1.], [2.], [4.]]])
     lens = torch.tensor([[3, 4, 2], [2, 5, 1]]).long()
-    repeated = repeat_consecutive_with_lens(seq, lens)
+    repeated = batch_repeat_interleave(seq, lens)
     assert torch.allclose(repeated, torch.tensor([[[1.], [1.], [1.], [2.], [2.], [2.], [2.], [4.], [4.]], [[1.], [1.], [2.], [2.], [2.], [2.], [2.], [4.], [0.]]]))
 
 def test_smooth_lddt_loss():