einx indexing has performance issues

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -104,7 +104,8 @@
     ExpressCoordinatesInFrame,
     RigidFrom3Points,
     RigidFromReference3Points,
-    calculate_weighted_rigid_align_weights
+    calculate_weighted_rigid_align_weights,
+    pack_one
 )
 
 from alphafold3_pytorch.utils.model_utils import distance_to_dgram
@@ -250,15 +251,6 @@ def max_neg_value(t: Tensor):
 def dict_to_device(d, device):
     return tree_map(lambda t: t.to(device) if is_tensor(t) else t, d)
 
-def pack_one(t, pattern):
-    packed, ps = pack([t], pattern)
-
-    def unpack_one(to_unpack, unpack_pattern = None):
-        unpacked, = unpack(to_unpack, ps, default(unpack_pattern, pattern))
-        return unpacked
-
-    return packed, unpack_one
-
 def exclusive_cumsum(t, dim = -1):
     return t.cumsum(dim = dim) - t
 

alphafold3_pytorch/inputs.py

@@ -6,6 +6,7 @@
 import os
 import random
 import statistics
+import traceback
 from collections import defaultdict
 from collections.abc import Iterable
 from contextlib import redirect_stderr
@@ -51,17 +52,20 @@
     rna_constants,
     ligand_constants
 )
+
 from alphafold3_pytorch.common.biomolecule import (
     Biomolecule,
     _from_mmcif_object,
     get_residue_constants,
 )
+
 from alphafold3_pytorch.data import (
     mmcif_parsing,
     msa_pairing,
     msa_parsing,
     template_parsing,
 )
+
 from alphafold3_pytorch.data.data_pipeline import (
     FeatureDict,
     get_assembly,
@@ -70,7 +74,9 @@
     make_template_features,
     merge_chain_features,
 )
+
 from alphafold3_pytorch.data.weighted_pdb_sampler import WeightedPDBSampler
+
 from alphafold3_pytorch.life import (
     ATOM_BONDS,
     ATOMS,
@@ -82,6 +88,7 @@
     reverse_complement,
     reverse_complement_tensor,
 )
+
 from alphafold3_pytorch.utils.data_utils import (
     PDB_INPUT_RESIDUE_MOLECULE_TYPE,
     extract_mmcif_metadata_field,
@@ -91,6 +98,7 @@
     is_polymer,
     make_one_hot,
 )
+
 from alphafold3_pytorch.utils.model_utils import (
     distance_to_dgram,
     exclusive_cumsum,
@@ -99,8 +107,11 @@
     offset_only_positive,
     remove_consecutive_duplicate,
     to_pairwise_mask,
+    pack_one
 )
+
 from alphafold3_pytorch.tensor_typing import Bool, Float, Int, typecheck
+
 from alphafold3_pytorch.utils.utils import default, exists, first, not_exists
 
 from alphafold3_pytorch.attention import (
@@ -759,7 +770,10 @@ def molecule_to_atom_input(mol_input: MoleculeInput) -> AtomInput:
             num_atoms = mol.GetNumAtoms()
             mol_atompair_ids = torch.zeros(num_atoms, num_atoms).long()
 
-            for bond in mol.GetBonds():
+            bonds = mol.GetBonds()
+            num_bonds = len(bonds)
+
+            for bond in has_bonds:
                 atom_start_index = bond.GetBeginAtomIdx()
                 atom_end_index = bond.GetEndAtomIdx()
 
@@ -785,12 +799,21 @@ def molecule_to_atom_input(mol_input: MoleculeInput) -> AtomInput:
 
                 updates.extend([bond_to, bond_from])
 
-            coordinates = tensor(coordinates).long()
-            updates = tensor(updates).long()
+            if num_bonds > 0:
+                coordinates = tensor(coordinates).long()
+                updates = tensor(updates).long()
 
-            mol_atompair_ids = einx.set_at(
-                "[h w], c [2], c -> [h w]", mol_atompair_ids, coordinates, updates
-            )
+                # mol_atompair_ids = einx.set_at("[h w], c [2], c -> [h w]", mol_atompair_ids, coordinates, updates)
+
+                molpair_strides = tensor(mol_atompair_ids.stride())
+                flattened_coordinates = (coordinates * molpair_strides).sum(dim = -1)
+
+                packed_atompair_ids, unpack_one = pack_one(mol_atompair_ids, '*')
+                packed_atompair_ids[flattened_coordinates] = updates
+
+                mol_atompair_ids = unpack_one(packed_atompair_ids)
+
+            # /einx.set_at
 
             row_col_slice = slice(offset, offset + num_atoms)
             atompair_ids[row_col_slice, row_col_slice] = mol_atompair_ids
@@ -1110,11 +1133,12 @@ def molecule_lengthed_molecule_input_to_atom_input(
 
         if mol_is_one_token_per_atom:
             coordinates = []
-            updates = []
 
             has_bond = torch.zeros(num_atoms, num_atoms).bool()
+            bonds = mol.GetBonds()
+            num_bonds = len(bonds)
 
-            for bond in mol.GetBonds():
+            for bond in bonds:
                 atom_start_index = bond.GetBeginAtomIdx()
                 atom_end_index = bond.GetEndAtomIdx()
 
@@ -1125,12 +1149,19 @@ def molecule_lengthed_molecule_input_to_atom_input(
                     ]
                 )
 
-                updates.extend([True, True])
+            if num_bonds > 0:
+                coordinates = tensor(coordinates).long()
 
-            coordinates = tensor(coordinates).long()
-            updates = tensor(updates).bool()
+                # has_bond = einx.set_at("[h w], c [2], c -> [h w]", has_bond, coordinates, updates)
 
-            has_bond = einx.set_at("[h w], c [2], c -> [h w]", has_bond, coordinates, updates)
+                has_bond_stride = tensor(has_bond.stride())
+                flattened_coordinates = (coordinates * has_bond_stride).sum(dim = -1)
+                packed_has_bond, unpack_has_bond = pack_one(has_bond, '*')
+
+                packed_has_bond[flattened_coordinates] = True
+                has_bond = unpack_has_bond(packed_has_bond, '*')
+
+                # / ein.set_at
 
             row_col_slice = slice(offset, offset + num_atoms)
             token_bonds[row_col_slice, row_col_slice] = has_bond
@@ -1279,9 +1310,7 @@ def molecule_lengthed_molecule_input_to_atom_input(
             coordinates = tensor(coordinates).long()
             updates = tensor(updates).long()
 
-            mol_atompair_ids = einx.set_at(
-                "[h w], c [2], c -> [h w]", mol_atompair_ids, coordinates, updates
-            )
+            # mol_atompair_ids = einx.set_at("[h w], c [2], c -> [h w]", mol_atompair_ids, coordinates, updates)
 
             row_col_slice = slice(offset, offset + num_atoms)
             atompair_ids[row_col_slice, row_col_slice] = mol_atompair_ids
@@ -4638,7 +4667,7 @@ def maybe_transform_to_atom_input(i: Any, raise_exception: bool = False) -> Atom
     try:
         return maybe_to_atom_fn(i)
     except Exception as e:
-        logger.error(f"Failed to convert input {i} to AtomInput due to: {e}")
+        logger.error(f"Failed to convert input {i} to AtomInput due to: {e}, {traceback.format_exc()}")
         if raise_exception:
             raise e
         return None

alphafold3_pytorch/utils/model_utils.py

@@ -4,7 +4,7 @@
 import einx
 import torch
 import torch.nn.functional as F
-from einops import einsum, pack, rearrange, reduce, repeat, unpack
+from einops import einsum, pack, unpack, rearrange, reduce, repeat
 from torch import Tensor
 from torch.nn import Module
 
@@ -17,6 +17,7 @@
 
 # helper functions
 
+
 # default scheduler used in paper w/ warmup
 
 

pyproject.toml

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