switch to pytest-shard (#137)

switch to pytest-shard

Author: lucidrains

.github/workflows/test.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        group: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        group: [0, 1, 2, 3, 4]
 
     steps:
     - uses: actions/checkout@v4
@@ -27,4 +27,4 @@ jobs:
         python -m pip install -e .[test]
     - name: Test with pytest
       run: |
-        python -m pytest --test-group-count 10 --test-group-random-seed 42 --test-group ${{ matrix.group }} tests/
+        python -m pytest --num-shards 5 --shard-id ${{ matrix.group }} tests/

README.md

@@ -77,6 +77,7 @@ atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
 additional_molecule_feats = torch.randint(0, 2, (2, seq_len, 5))
 additional_token_feats = torch.randn(2, seq_len, 2)
 is_molecule_types = torch.randint(0, 2, (2, seq_len, 5)).bool()
+is_molecule_mod = torch.randint(0, 2, (2, seq_len, 4)).bool()
 molecule_ids = torch.randint(0, 32, (2, seq_len))
 
 template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
@@ -107,6 +108,7 @@ loss = alphafold3(
     additional_molecule_feats = additional_molecule_feats,
     additional_token_feats = additional_token_feats,
     is_molecule_types = is_molecule_types,
+    is_molecule_mod = is_molecule_mod,
     msa = msa,
     msa_mask = msa_mask,
     templates = template_feats,
@@ -134,6 +136,7 @@ sampled_atom_pos = alphafold3(
     additional_molecule_feats = additional_molecule_feats,
     additional_token_feats = additional_token_feats,
     is_molecule_types = is_molecule_types,
+    is_molecule_mod = is_molecule_mod,
     msa = msa,
     msa_mask = msa_mask,
     templates = template_feats,
@@ -180,6 +183,7 @@ alphafold3 = Alphafold3(
     atoms_per_window = 27,
     dim_template_feats = 44,
     num_dist_bins = 38,
+    num_molecule_mods = 0,
     confidence_head_kwargs = dict(
         pairformer_depth = 1
     ),

alphafold3_pytorch/alphafold3.py

@@ -4431,7 +4431,7 @@ def __init__(
         num_molecule_types: int = NUM_MOLECULE_IDS,     # restype in additional residue information, apparently 32. will do 33 to account for metal ions
         num_atom_embeds: int | None = None,
         num_atompair_embeds: int | None = None,
-        num_molecule_mods: int | None = None,
+        num_molecule_mods: int | None = DEFAULT_NUM_MOLECULE_MODS,
         distance_bins: List[float] = torch.linspace(3, 20, 38).float().tolist(),
         ignore_index = -1,
         num_dist_bins: int | None = None,

alphafold3_pytorch/inputs.py

@@ -7,7 +7,7 @@
 from functools import partial, wraps
 from itertools import groupby
 from collections import defaultdict
-from collections.abc import Iterableassignment
+from collections.abc import Iterable
 from dataclasses import asdict, dataclass, field
 from typing import Any, Callable, Dict, List, Literal, Set, Tuple, Type
 

alphafold3_pytorch/mocks.py

@@ -18,7 +18,7 @@ def __init__(
         data_length,
         max_seq_len = 16,
         atoms_per_window = 4,
-        has_molecule_mods = False
+        has_molecule_mods = True
     ):
         self.data_length = data_length
         self.max_seq_len = max_seq_len

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.96"
+version = "0.2.97"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -62,7 +62,7 @@ Repository = "https://github.com/lucidrains/alphafold3-pytorch"
 examples = []
 test = [
     "pytest",
-    "pytest-split-tests",
+    "pytest-shard",
 ]
 
 [build-system]

tests/configs/trainer_with_atom_dataset_created_from_pdb.yaml

@@ -17,7 +17,6 @@ model:
     num_plddt_bins: 50
     num_pde_bins: 64
     num_pae_bins: 64
-    num_molecule_mods: 1
     sigma_data: 16
     diffusion_num_augmentations: 4
     loss_confidence_weight: 0.0001

tests/configs/trainer_with_pdb_dataset.yaml

@@ -17,7 +17,6 @@ model:
     num_plddt_bins: 50
     num_pde_bins: 64
     num_pae_bins: 64
-    num_molecule_mods: 1
     sigma_data: 16
     diffusion_num_augmentations: 4
     loss_confidence_weight: 0.0001

tests/configs/trainer_with_pdb_dataset_and_weighted_sampling.yaml

@@ -17,7 +17,6 @@ model:
     num_plddt_bins: 50
     num_pde_bins: 64
     num_pae_bins: 64
-    num_molecule_mods: 1
     sigma_data: 16
     diffusion_num_augmentations: 4
     loss_confidence_weight: 0.0001

tests/test_af3.py

@@ -689,6 +689,7 @@ def test_alphafold3_without_msa_and_templates():
         dim_atom_inputs = 77,
         dim_template_feats = 44,
         num_dist_bins = 38,
+        num_molecule_mods = 0,
         checkpoint_trunk_pairformer = True,
         checkpoint_diffusion_token_transformer = True,
         confidence_head_kwargs = dict(
@@ -767,6 +768,7 @@ def test_alphafold3_force_return_loss():
         dim_atom_inputs = 77,
         dim_template_feats = 44,
         num_dist_bins = 38,
+        num_molecule_mods = 0,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
         ),
@@ -851,6 +853,7 @@ def test_alphafold3_force_return_loss_with_confidence_logits():
         dim_atom_inputs = 77,
         dim_template_feats = 44,
         num_dist_bins = 38,
+        num_molecule_mods = 0,
         confidence_head_kwargs = dict(
             pairformer_depth = 1
         ),
@@ -913,6 +916,7 @@ def test_alphafold3_with_atom_and_bond_embeddings():
     alphafold3 = Alphafold3(
         num_atom_embeds = 7,
         num_atompair_embeds = 3,
+        num_molecule_mods = 0,
         dim_atom_inputs = 77,
         dim_template_feats = 44
     )