Fix several latent bugs in MSA, ranking, and structure-context code

chai_lab/data/dataset/msas/colabfold.py

@@ -61,8 +61,8 @@ def submit(seqs, mode, N=101):
             query += f">{n}\n{seq}\n"
             n += 1
 
+        error_count = 0
         while True:
-            error_count = 0
             try:
                 # https://requests.readthedocs.io/en/latest/user/advanced/#advanced
                 # "good practice to set connect timeouts to slightly larger than a multiple of 3"
@@ -95,8 +95,8 @@ def submit(seqs, mode, N=101):
         return out
 
     def status(ID):
+        error_count = 0
         while True:
-            error_count = 0
             try:
                 res = requests.get(
                     f"{host_url}/ticket/{ID}", timeout=6.02, headers=headers

chai_lab/data/dataset/structure/all_atom_structure_context.py

@@ -382,7 +382,7 @@ def merge(
         token_backbone_frame_index = torch.cat(
             [
                 x.token_backbone_frame_index + count
-                for x, count in zip(contexts, token_offsets)
+                for x, count in zip(contexts, atom_offsets)
             ]
         )
 

chai_lab/data/features/generators/msa.py

@@ -240,7 +240,9 @@ def _generate(
         query = msa_dataset_source_to_int[MSADataSource.QUERY]
         none = msa_dataset_source_to_int[MSADataSource.NONE]
         # chai-1 specific: replace QUERY with NONE
-        msa_sequence_source[msa_sequence_source.eq(query)] = none
+        msa_sequence_source = msa_sequence_source.masked_fill(
+            msa_sequence_source.eq(query), none
+        )
         # use none for masking.
         msa_sequence_source = msa_sequence_source.masked_fill(~msa_mask, none)
         return self.make_feature(data=msa_sequence_source.unsqueeze(-1))

chai_lab/data/features/generators/token_dist_restraint.py

@@ -120,11 +120,11 @@ def __init__(
             )
         ).long()
         self.key_entity_types = torch.tensor(
-            [
+            (
                 [e.value for e in key_entity_types]
                 if key_entity_types is not None
                 else [e.value for e in EntityType]
-            ]
+            )
         ).long()
 
     def get_num_restraints(self, batch_size) -> list[int]:

chai_lab/ranking/clashes.py

@@ -79,7 +79,8 @@ def has_inter_chain_clashes(
     ).ge(max_clash_ratio)
 
     has_clashes |= (
-        per_chain_pair_clashes / rearrange(atoms_per_chain, "b c -> b 1 c").clamp(min=1)
+        per_chain_pair_clashes
+        / rearrange(atoms_per_chain, "... c -> ... 1 c").clamp(min=1)
     ).ge(max_clash_ratio)
 
     # only consider clashes between pairs of polymer chains