chaobrain/brainx: BrainX 2026.8.18
Abstract
This release brings the bundle to JAX 0.11. JAX 0.11.0 and 0.11.1 changed the tracing internals the whole stack is built on — the scan primitive dropped its num_consts / num_carry parameters, Jaxpr and ClosedJaxpr were merged into a single class, core.CallPrimitive was deleted, and is_constant_dim and concrete_or_error moved namespaces — and every affected component has now adapted. Five packages advance: BrainUnit 0.5.2, BrainEvent 0.2.1, BrainState 0.5.4, BrainTrace 0.2.6, and BrainPy 2.8.2. The supported JAX range widens from =0.2.8 Highlights BrainEvent 0.2.1 (spans 0.2.0) — just-in-time connectivity now draws one matrix everywhere: the numba CPU kernels were rebuilt on the CUDA light-RNG walk and the mm path folded onto the 32-lane mv walk, so jits / jitn / jitu and every derived entry point materialize the same matrix on CPU and GPU and across the matrix-vector and matrix-matrix paths. Adds a first-class Dense data representation, int64 indptr for CSR / CSC past the int32 nnz range, a tunable CSR hybrid CUDA scheduler (HybridConfig, init_csr_config), a 'cublas' dense GPU backend, and numba CPU kernels for the previously CUDA-only JITC CSR and dt2t paths. An operator-registration audit closes 19 defects (stale backend dispatch, dropped JVP rules, incomplete cache keys, incorrect vmap of numba.cuda kernels), and a CUDA sweep fixes an out-of-bounds shared-memory abort on float64 binary_csrmm, widens 54 index expressions that wrapped past INT32_MAX, and restores the warp-per-row CSRMV tier (up to 3.4×). Breaking: JITC values recorded against 0.1.x will not reproduce and seeds are not portable — re-record any golden outputs; CSR / CSC now validate at construction and raise immediately; pararnn is removed. 0.2.0's matrix_mode keyword and mat.mv / mat.mm views were reverted in 0.2.1, so upgrading straight from 0.1.2 sees no signature change. BrainState 0.5.4 (spans 0.5.3) — pure JAX-compatibility patches with no public API change. 0.5.3 repairs brainstate.transform's IR tooling for JAX 0.11.0 via a version-agnostic scan_num_consts_carry() helper (18 failures → a clean 1282-test pass); 0.5.4 restores import brainstate itself on 0.11.1 after CallPrimitive was deleted, fixes an eqns_to_closed_jaxpr() length check that had silently stopped validating, and turns an installed-but-unimportable optional interop framework into an actionable InteropError. Known issue: flax==0.12.8 cannot import under JAX 0.11.1, so brainstate.interop's Flax conversions are unavailable there (Equinox unaffected) — pin jax==0.11.0 if you need them. BrainTrace 0.2.6 (spans 0.2.5) — adds the sequence-driver API (etrace_grad / etrace_evolve), which removes the hand-written scan-and-accumulate loop from every call site, and five new learning rules — SnAp, UORO, ThreeFactor, DNI and the public RandomProjectionVjpAlgorithm — expressed as coordinates in the new ETraceConfig six-axis space rather than as bespoke implementations. Control-flow canonicalization fixpoints are now bounded (a non-converging rewrite raises instead of hanging the compiler) and the hidden↔gradient correspondence is checked with explicit raises rather than asserts that vanished under python -O. 0.2.6 then fixes the merged-Jaxpr representation, under which transition jaxprs began reporting constvars == [] and took down essentially every gradient path: 690 of 2902 tests failed on JAX 0.11.1, all passing again. Breaking: OTTT, OSTTP, OTPE and PresynapticTrace are removed — migrate OTTT → pp_prop, OTPE → D_RTRL or pp_prop, OSTTP → EProp(feedback='random') — and IODimVjpAlgorithm.decay is now a read-only property. BrainPy 2.8.2 — imports concrete_or_error from jax.extend.core on JAX >= 0.11, where it was removed from jax.interpreters.partial_eval; completes the master → main branch rename begun in 2.8.1; and repoints the brainpy.state documentation links at the canonical brainx.chaobrain.com/brainpy-state URLs. BrainUnit 0.5.2 — a drop-in upgrade that restores brainunit.math.einsum on JAX 0.11, where is_constant_dim was removed from the public jax.core namespace with no jax.extend.core replacement; the symbol now resolves through the version-aware _compatible_import shim across JAX 0.6–0.11. JAX range widened to
Publication details
- Journal / source
- Zenodo (CERN European Organization for Nuclear Research)
- Publication date
- 18 August 2026