Update fork to AlphaFold 3 v3.0.4 by DimaMolod · Pull Request #3 · KosinskiLab/alphafold3

DimaMolod · 2026-10-06T18:28:51Z

Summary

Update the fork to AlphaFold 3 v3.0.4 (google-deepmind/alphafold3 85c4d20). This is a regular merge of the upstream tag into fork main, with no conflicts. The fork's own changes (gapped and duplicate residue IDs in folding_input.py, parsing.py and structure_tables.py) apply unchanged on top. The result differs from upstream v3.0.4 only in those three files.

What upstream v3.0.4 brings:

Dependencies: jax 0.10.2 and tokamax 0.0.12; uv.lock updated.
Fixes: lower-memory OuterProductMean; a numerical-stability fix; loading of empty template maps; early checks that seeds are uint32; chain IDs in the summary confidence JSON.
New: CPU and Mac backends, and example input files.
Plus one later upstream commit, cherry-picked with -x: 43525e6 "Make the data test robust to RDKit/CCD updates" (test file and golden only). Without it, CI fails run_alphafold_data_test.py::test_featurisation, because the CCD it downloads has drifted. Unchanged fork main fails the same way (run 37510781084), and so did v3.0.4 alone.

Validation

git merge v3.0.4 into fork main: clean. Against upstream v3.0.4, only the fork's three files differ.
The fused-triangle work replays cleanly on top: 90/90 tests on jax 0.10.2 / tokamax 0.0.12, and the golden model_config.json matches. That PR follows.
CI on this PR.
GPU validation on a jax 0.10.2 AlphaPulldown image after the merge, before any AlphaPulldown release.

🤖 Generated with Claude Code

Inspired by google-deepmind/alphafold3#685 PiperOrigin-RevId: 936470606 Change-Id: I415b47de84c6337e4dff87ea2996dc7020974031

PiperOrigin-RevId: 937865794 Change-Id: I69046a04cfe1007a297e2d6fae2e072fd4984319

…g `etils.epath`. This change migrates path operations in AlphaFold3 from Python's standard `pathlib.Path` and `os.path` to the unified filesystem wrapper `etils.epath.Path`. This enables several core flags (`--json_path`, `--input_dir`, `--output_dir`, `--model_dir`, and `--pdb_database_path`) to natively support GCS cloud-hosted paths (`gs://...`). Support for Google Cloud Storage paths is treated as an optional dependency. Open source users can opt into GCS path support by installing the `gcsfs` add-in, e.g., via: `uv sync --extra` , or adding `--build-arg UV_EXTRAS="--extra gcsfs"` to their docker build command. PiperOrigin-RevId: 938058741 Change-Id: I3463fb159a082deca4e6752fceec7dae3d75c1d3

@mooreneural

* Uses significantly less RAM (256 -> 72 MB for 1024-token input). Inspired by google-deepmind/alphafold3#674, thank you very much for the suggestion @mooreneural. PiperOrigin-RevId: 941052434 Change-Id: I0f1cfa6a16f7f2da4371a0c0c53642ef26b1328c

…nto`. PiperOrigin-RevId: 941158582 Change-Id: I2d4c92690440b2f99f11407a72476c910d481937

PiperOrigin-RevId: 942076492 Change-Id: I93c1ff09ba89919f0089ab90215353ea4ce1c368

…rectory PiperOrigin-RevId: 943802483 Change-Id: Ice5dcd44bd8a7b84c83ccb976f300d14c050bed3

PiperOrigin-RevId: 944697091 Change-Id: I8eb3b1ef0b8f91412c920ef9844513126422058f

From google-deepmind/alphafold3#692. PiperOrigin-RevId: 945022442 Change-Id: Idbebdf27f9eb16a8231b6fa128717a92a0d09964

@ntnn19

Adapted from google-deepmind/alphafold3#678, thanks @ntnn19. PiperOrigin-RevId: 945057068 Change-Id: Ic46e7fc9e5a00bd1f6261b1b28c23065d9d85f92

@Fred-Jing

…ymer residues Inspired by google-deepmind/alphafold3#662, thanks @Fred-Jing. PiperOrigin-RevId: 945075811 Change-Id: I9479f460108e71883a0e7a73054a34349081e746

PiperOrigin-RevId: 945749261 Change-Id: Ic6b33b6e15f92ad1a2970625a229d1d5544da95c

PiperOrigin-RevId: 946977350 Change-Id: I033f1e72d88d0a54713838baaecb1ac41e398031

@wuooo339

Discovered by @wuooo339 in google-deepmind/alphafold3#273 (comment) PiperOrigin-RevId: 948202523 Change-Id: I7566974d2dbd04abdabfcaa471a01b8ca3e12f33

…_getaffinity (e.g. MacOS) PiperOrigin-RevId: 948990027 Change-Id: I46bba29bbd0a9a160063dc885ac57675f56c9790

PiperOrigin-RevId: 949000038 Change-Id: I14871ff35a313a8e97b2d24bd92e8b2931b8e97a

PiperOrigin-RevId: 949536584 Change-Id: I01433bcd34f3326617767aacb9601ac12fdb1735

…well PiperOrigin-RevId: 951547559 Change-Id: I6a212de48617a7d0ff12b283711c6fd813772240

PiperOrigin-RevId: 952763752 Change-Id: I4d0dc0bebf06720b8969d49315c48e184a96bbc1

PiperOrigin-RevId: 952769192 Change-Id: I8e596d4d0768c820191bc1c58708744d0e3d4e76

PiperOrigin-RevId: 952803413 Change-Id: Ia5ae477f5169b07727fddde418acd15958afa67b

PiperOrigin-RevId: 954188170 Change-Id: I047edf9a7dcdd8ac064db9197a6ac635c9bb42ba

PiperOrigin-RevId: 954517731 Change-Id: I2e3b4ff9e0083de182eaf1d68d5621126d5cf3a7

@SuhasSrinivasan

* Heavily inspired by google-deepmind/alphafold3#699, thanks @SuhasSrinivasan! * Roughly 3x faster inference time than CPU on a MacBook Pro M3. PiperOrigin-RevId: 955210867 Change-Id: Ie9064488ce69dfd0442f263b13886d489ee39604

Merge upstream v3.0.4 (google-deepmind/alphafold3, 85c4d20) into the fork. No conflicts with the fork's changes (gapped and duplicate residue IDs). Upstream brings jax 0.10.2 and tokamax 0.0.12, a lower-memory OuterProductMean, CPU/Mac backends and input-validation fixes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

PiperOrigin-RevId: 964146796 Change-Id: I40a0faba4da75a071adad96484e700d3958a2071 (cherry picked from commit 43525e6d629cd1db78f2f5d6fa3ce9b68e81dcb8)

, #4) Fork main 7809148: upstream v3.0.4 (#3) plus the fused triangle kernels (#4). Same tree as the tested ab98e22 apart from upstream's data-test fix (43525e6). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Augustin-Zidek and others added 26 commits June 23, 2026 00:07

Add early validation that RNG seeds are uint32s

a84d263

Inspired by google-deepmind/alphafold3#685 PiperOrigin-RevId: 936470606 Change-Id: I415b47de84c6337e4dff87ea2996dc7020974031

Add missing type annotations

b2f3d45

PiperOrigin-RevId: 937865794 Change-Id: I69046a04cfe1007a297e2d6fae2e072fd4984319

Use less memory in the OuterProductMean

7b4c2a5

* Uses significantly less RAM (256 -> 72 MB for 1024-token input). Inspired by google-deepmind/alphafold3#674, thank you very much for the suggestion @mooreneural. PiperOrigin-RevId: 941052434 Change-Id: I0f1cfa6a16f7f2da4371a0c0c53642ef26b1328c

Fix model loading for etils file providers that do not support `readi…

b1f7802

…nto`. PiperOrigin-RevId: 941158582 Change-Id: I2d4c92690440b2f99f11407a72476c910d481937

No public description

4958a9b

PiperOrigin-RevId: 942076492 Change-Id: I93c1ff09ba89919f0089ab90215353ea4ce1c368

Add a flag to control whether Terms of Use are saved in the output di…

7c1f195

…rectory PiperOrigin-RevId: 943802483 Change-Id: Ice5dcd44bd8a7b84c83ccb976f300d14c050bed3

Fix np.savez_compressed failure writing to unseekable epath backends

ded9d57

PiperOrigin-RevId: 944697091 Change-Id: I8eb3b1ef0b8f91412c920ef9844513126422058f

Do not serialize empty templateIndices as None

62136ec

From google-deepmind/alphafold3#692. PiperOrigin-RevId: 945022442 Change-Id: Idbebdf27f9eb16a8231b6fa128717a92a0d09964

Add chain IDs in the summary confidence JSON for easier use

e68269f

Adapted from google-deepmind/alphafold3#678, thanks @ntnn19. PiperOrigin-RevId: 945057068 Change-Id: Ic46e7fc9e5a00bd1f6261b1b28c23065d9d85f92

Add a note that user-provided CCD can be used for custom modified pol…

cbdd420

…ymer residues Inspired by google-deepmind/alphafold3#662, thanks @Fred-Jing. PiperOrigin-RevId: 945075811 Change-Id: I9479f460108e71883a0e7a73054a34349081e746

Add a directory with example input JSON files

4d610bc

PiperOrigin-RevId: 945749261 Change-Id: Ic6b33b6e15f92ad1a2970625a229d1d5544da95c

Delete the ref_pos entry only after the pickle is written

0d3facb

PiperOrigin-RevId: 946977350 Change-Id: I033f1e72d88d0a54713838baaecb1ac41e398031

Fix numerical instability that prevents AF3 from running without a GPU

fd39d2c

Discovered by @wuooo339 in google-deepmind/alphafold3#273 (comment) PiperOrigin-RevId: 948202523 Change-Id: I7566974d2dbd04abdabfcaa471a01b8ca3e12f33

Make the num_cpus function compatible with platforms without os.sched…

27f4d6d

…_getaffinity (e.g. MacOS) PiperOrigin-RevId: 948990027 Change-Id: I46bba29bbd0a9a160063dc885ac57675f56c9790

Add a flag to run on CPU only

7ecd64f

PiperOrigin-RevId: 949000038 Change-Id: I14871ff35a313a8e97b2d24bd92e8b2931b8e97a

Make run_alphafold runnable without a GPU and on Mac OS

f43a5c4

PiperOrigin-RevId: 949536584 Change-Id: I01433bcd34f3326617767aacb9601ac12fdb1735

Bump JAX and Tokamax to newer versions to fix unified memory on Black…

a3b8355

…well PiperOrigin-RevId: 951547559 Change-Id: I6a212de48617a7d0ff12b283711c6fd813772240

Update the instructions on how to obtain AlphaFold 3 weights

dd1a7ba

PiperOrigin-RevId: 952763752 Change-Id: I4d0dc0bebf06720b8969d49315c48e184a96bbc1

Fix broken formatting in the README

4e62292

PiperOrigin-RevId: 952769192 Change-Id: I8e596d4d0768c820191bc1c58708744d0e3d4e76

Replace --use_cpu_only with --jax_backend for future expandability

98a15dc

PiperOrigin-RevId: 952803413 Change-Id: Ia5ae477f5169b07727fddde418acd15958afa67b

Suppressions for pyrefly

01bd8d1

PiperOrigin-RevId: 954188170 Change-Id: I047edf9a7dcdd8ac064db9197a6ac635c9bb42ba

Add a missing backslash and remove a trailing one in a multiline command

8a84535

PiperOrigin-RevId: 954517731 Change-Id: I2e3b4ff9e0083de182eaf1d68d5621126d5cf3a7

Add support for using Apple Silicon GPU

85c4d20

* Heavily inspired by google-deepmind/alphafold3#699, thanks @SuhasSrinivasan! * Roughly 3x faster inference time than CPU on a MacBook Pro M3. PiperOrigin-RevId: 955210867 Change-Id: Ie9064488ce69dfd0442f263b13886d489ee39604

Make the data test robust to RDKit/CCD updates

ba044f4

PiperOrigin-RevId: 964146796 Change-Id: I40a0faba4da75a071adad96484e700d3958a2071 (cherry picked from commit 43525e6d629cd1db78f2f5d6fa3ce9b68e81dcb8)

DimaMolod merged commit 758d767 into main Oct 7, 2026
1 check passed

This was referenced Oct 7, 2026

Opt-in fused triangle kernels (Pallas) for inference #4

Merged

Opt-in fused triangle kernels (Pallas) for inference #2

Closed

DimaMolod mentioned this pull request Oct 7, 2026

AF3: --fast_kernels with fused triangle kernels; AlphaFold 3 v3.0.4 KosinskiLab/AlphaPulldown#646

Merged

3 of 4 tasks

Navigation Menu

Repository navigation

Update fork to AlphaFold 3 v3.0.4 - #3

Update fork to AlphaFold 3 v3.0.4#3
DimaMolod merged 26 commits into
mainfrom
exp/af3-v3.0.4

DimaMolod commented Oct 6, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

DimaMolod commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

DimaMolod commented Oct 6, 2026 •

edited

Loading