Align the jax stack, name the CUDA extra for CUDA, and drop Python 3.11

[mahdi-shafiei]

Brings DiffBio onto the stack the sibling repositories now share: jax 0.11.1, jaxlib 0.11.1, flax 0.12.9, optax 0.2.8, grain 0.2.18, on Python 3.12.

gpu → cuda12

JAX names its own extras for the CUDA major version — cuda12, cuda12-local, cuda13 — and publishes no gpu extra. This project also ships metal, which is a GPU, so gpu could not distinguish the two. setup.sh already spoke of a cuda12 backend while syncing --extra gpu.

The extra also drops its jaxlib line: every jax[cuda12*] extra already requires jaxlib at its own version, so restating it only creates a second place for the two to disagree.

Floors that state requirements

A lock refresh does not move a dependency that still satisfies its floor, so each of these is declared rather than left to the resolver:

declaration	reason
flax>=0.12.1	first release with nnx.Variable.set_value, which the pipeline calls
optax>=0.2.8	below it, optax sets jax_pmap_shmap_merge, removed in jax 0.10 — importing flax raises AttributeError and collection dies for the entire suite
dm-haiku>=0.0.17 (constraint)	arrives via jax-md; up to 0.0.16 it binds jax.core.DropVar at import, removed in jax 0.11.0

Python 3.11

Follows jax 0.11.0 dropping it. This also collapses the resolution to a single unforked lockfile, which is most of the lock diff. Raising ruff's target to py312 enables the PEP 695 rules; the four aliases converted here have no runtime uses anywhere, and the module-level TypeVar its generic replaced is removed with it.

Two defects found while verifying, each in its own commit

top_k_mask returned values outside its documented [0, 1] — one float32 ULP over, from the isotonic solver's pooled averaging. The same input gives the same value on the previous jax release, so this predates the version move; the property test only reaches it on some seeds. The forward value is saturated onto the bound while the backward pass stays on the projection's own gradient — a plain clip would zero the gradient of exactly the entries the mask selects. Over 540 combinations the gradients are unchanged and only the four out-of-bound masks move.

A benchmark bypassed the optimizer ownership boundary, failing a contract test on main. It went unnoticed because no CI job runs tests/benchmarks/ — the core-and-misc shard ignores it, no other shard names it, and the e2e and performance jobs select by marker and collect nothing there. Routing it through create_benchmark_optimizer required exposing weight_decay, because OptimizerConfig defaults it to 0.0 while optax.adamw defaults to 1e-4; passing 1e-4 keeps the optimizer bit-identical, verified against the previous call with a deliberately wrong value as a control.

Verification

All five unit shards on Python 3.12 with the extras CI installs: 3126 passed, 48 skipped, 0 failed. ruff check, ruff format --check (606 files) and lint-imports (6 contracts) all clean.

tests/benchmarks/ remains outside every CI job's selection. That gap is left as-is here rather than widened into a CI change alongside a dependency move.