Skip to content

ci: pin jax 0.12.0 in test matrix; add jax 0.12 compatibility audit #243

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chaoming0625 merged 1 commit into from
Oct 20, 2026
Merged

ci: pin jax 0.12.0 in test matrix; add jax 0.12 compatibility audit #243

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/CI.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ jobs:
fail-fast: false
matrix:
python-version: ["3.13"]
jax-version: ["0.9.0", "0.10.0", "0.11.0", ""]
jax-version: ["0.9.0", "0.10.0", "0.11.0", "0.12.0", ""]

steps:
- uses: actions/checkout@v6
Expand Down
47 changes: 47 additions & 0 deletions doc/specs/jax-0.12-compat.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
# JAX 0.12.0 compatibility audit

## Problem

JAX 0.12.0 (2026-10-19) and 0.11.2 (2026-09-17) shipped several breaking changes. The
`""` (latest) leg of the CI matrix now installs 0.12.0, so we need to know whether brainstate
is affected before that leg goes red.

## Method

1. Cross-check the 0.11.2 and 0.12.0 changelog entries against brainstate's source.
2. Run the full suite in an isolated venv with `jax[cpu]==0.12.0` and the CI dependency
set (`requirements-dev.txt`). Repeat with `--run-slow` and
`XLA_FLAGS=--xla_force_host_platform_device_count=8` to cover the multi-device tests.
3. Probe the changes the suite might not cover, and compare against a jax 0.11.1 venv.

## Findings

**No incompatibilities.** On jax 0.12.0, `import brainstate` is clean under
`-W error::DeprecationWarning`, and the suite passes: 5314 passed / 75 skipped by default,

@sourcery-ai sourcery-ai Bot Oct 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚪ Low · Compatibility test totals conflict

When readers compare the audit with the PR description, the audit reports 5,314 passed / 75 skipped by default and 5,333 / 56 with slow tests and 8 devices, while the PR description reports 5,323 passed with CI dependencies. The reports do not say which run the PR total represents, so readers cannot verify the compatibility claim.

Label the exact test configuration behind each total in the audit and PR description, and reconcile the totals for matching runs.

Prompt for AI agents 
In `doc/specs/jax-0.12-compat.md` at line 20:

**Compatibility test totals conflict**

When readers compare the audit with the PR description, the audit reports 5,314 passed / 75 skipped by default and 5,333 / 56 with slow tests and 8 devices, while the PR description reports 5,323 passed with CI dependencies. The reports do not say which run the PR total represents, so readers cannot verify the compatibility claim.

Label the exact test configuration behind each total in the audit and PR description, and reconcile the totals for matching runs.

5333 passed / 56 skipped with slow + 8 devices. No `brainstate/` change is needed.

| JAX change | brainstate exposure | Result |
|---|---|---|
| `jax.custom_remat` removed | not used | n/a |
| `jnp.atleast_*d` return tuples | not used | n/a |
| `jax.random.multinomial` `shape` is now the batch shape | `random/_impl.py` has its own implementation | shapes match NumPy for `size=None/int/tuple` |
| `custom_vjp` under `vmap` errors when the cotangent of an unbatched primal would be counted per example | `nn/_embedding.py` `_lookup` | `vmap(grad)` with weight unbatched/batched and `grad(vmap)` agree with per-example grads, for all three lookup variants |
| Duplicate / closed-over-and-passed `Ref` errors in `checkpoint`/`custom_vjp`/`shard_map` | brainstate does not pass jax `Ref`s; `State` values are threaded as arrays | `transform.remat` that reads the same State twice is fine; shard_map tests pass with 8 devices |
| `jit(inline=True)` now means `Inline.JAX_LATE` (0.11.2) | `transform/_jit.py` passes `inline` through | nested inlined jit is correct; the docstring doesn't promise early/late semantics |
| `jnp.array` `copy/order/ndmin` keyword-only; 2D `jnp.cross` removed; `shard_alike`, `psend/precv`, `ProfileOptions` | not used | n/a |
| NumPy >= 2.2 required | jax's own requirement takes care of it | n/a |
| Private `jax._src` imports (batching tracers, `trace_ctx`, `UNSPECIFIED`, `flatten_axes`, `TypedNdArray`, ...) | `_compatible_import.py` and others | all import and work |
| Optional interop frameworks | flax 0.12.10, equinox 0.13.8, optax 0.2.8 | interop tests run (not skipped) and pass |

## Follow-ups (not 0.12 regressions)

- `transform/_debug_test.py` triggers jax's `Python iterable type 'generator' is treated as a
leaf in PyTree` DeprecationWarning (~30 times). It is present on 0.11.1 too. JAX says this
will become an error in a future release. The warning comes from jaxlib's C++ flatten
inside a callback path, and turning it into an error does not fail any test, so its exact
source is still unidentified.

## CI

Added `"0.12.0"` to the `jax-version` matrix in `.github/workflows/CI.yml`, so 0.12 stays
tested after `""` moves on to 0.13.
Loading