Replace flax usage with equinox

[jd61349]

PR Type

• Refactoring (no functional changes)

Description

Replace all usage of Flax with Equinox. Fixes #582.

How Has This Been Tested?

Related tests have been changed to reflect updated code, however what they test has remained the same. All tests pass.

Does this PR introduce a breaking change?

No.

Checklist before requesting a review

• I have made sure that my PR is not a duplicate.
• My code follows the style guidelines of this project.
• I have ensured my code is easy to understand, including docstrings and comments where necessary.
• I have performed a self-review of my code.
• I have made corresponding changes to the documentation.
• My changes generate no new warnings.
• New and existing unit tests pass locally with my changes.
• Any dependent changes have been merged and published in downstream modules.
• I have updated CHANGELOG.md, if appropriate.

[CLAassistant]

All committers have signed the CLA.

[sylvesterkaczmarek]

MLP.__init__ now passes mode="batch" to eqx.nn.BatchNorm, but pyproject.toml still supports equinox>=0.11.11. In Equinox 0.11.11 BatchNorm.__init__ has no mode argument, so a lowest-resolution install fails as soon as this benchmark constructs the model. Could we either raise the lower bound to the version that introduced this API or avoid the newer argument, and add a lowest-dependency benchmark smoke test?

[jd61349]

MLP.__init__ now passes mode="batch" to eqx.nn.BatchNorm, but pyproject.toml still supports equinox>=0.11.11. In Equinox 0.11.11 BatchNorm.__init__ has no mode argument, so a lowest-resolution install fails as soon as this benchmark constructs the model. Could we either raise the lower bound to the version that introduced this API or avoid the newer argument, and add a lowest-dependency benchmark smoke test?

Good spot! Thanks

[github-actions]

Performance review

Commit d6460a3 - Merge 45e5a4b into 9a804d3

Statistically significant changes

• basic_rpc:
  • OLD: compilation 0.7062 units ± 0.01754 units; execution 0.02056 units ± 0.0001399 units
  • NEW: compilation 0.6986 units ± 0.03884 units; execution 0.02276 units ± 0.0002806 units
  • Significant increase in execution time (10.70%, p=7.695e-12)

---

Normalisation values for new data:
Compilation: 1 unit = 646.35 ms
Execution: 1 unit = 697.58 ms

Last updated: Wed, 16 Sep 2026 11:46:16 GMT

[sylvesterkaczmarek]

Re-reviewed current 15a48daf. The compatibility issue I raised is resolved: the Equinox lower bound is now >=0.13.0, matching the BatchNorm(..., mode="batch") API used by the benchmark, and the lowest/highest dependency test matrix is green across supported Python versions. No remaining blocker from me.

[tm91236]

This looks good, but it does introduce some breaking changes without any deprecation warnings. I think we should keep the old network.py module, but mark the ScoreNetwork as deprecated for removal in v1.2.0.

[sylvesterkaczmarek]

Re-reviewed current 45e5a4b5. Restoring coreax/networks.py for the requested deprecation path introduces a new compatibility problem: that module imports flax, but this PR still removes flax from the package dependencies. A normal installation can therefore no longer import the deprecated module that is being kept specifically to preserve compatibility. The current pre-commit run confirms this with unresolved imports for flax, flax.linen, and flax.training in coreax/networks.py. Please retain the flax dependency for the deprecation window (or otherwise make the compatibility module self-contained) and add a simple normal-install import regression for coreax.networks/ScoreNetwork.