fix: jax 0.11.1 compatibility, release v0.5.4

[chaoming0625]

Summary

JAX 0.11.1 removed core.CallPrimitive and completed the Jaxpr/ClosedJaxpr merge begun in 0.11.0. Together these broke import brainstate outright and silently disabled a validation check in the IR tooling. This PR restores compatibility and cuts release v0.5.4.

Fixes

_compatible_import — import brainstate failed on 0.11.1

CallPrimitive is gone from both jax.extend.core and jax._src.core, replaced by create_call_primitive() / register_call_primitive_rules(). It was imported unconditionally on jax>=0.10, so every 0.11.1 install failed at import time. The import is now attempted first, falling back to a thin Primitive subclass that applies register_call_primitive_rules() at construction — preserving the public CallPrimitive name and its constructor behavior. JAX 0.11.0 still ships the original class and continues to use it.

transform/_ir_processing — silent validation hole

0.11.1 dropped Jaxpr._num_consts, so an input counts as a constvar only when a constant value is attached. eqns_to_closed_jaxpr() validated len(consts) against jaxpr.constvars read back off an intermediate open jaxpr — now always empty. That both rejected valid consts and, more seriously, silently accepted mismatched ones, defeating the check entirely.

Binder resolution is factored into _resolve_jaxpr_vars() so validation runs against the resolved constvars; the returned ClosedJaxpr(jaxpr, consts) restores the constvars/invars split identically on all supported versions.

The boundary is 0.11.1, not 0.11.0 — 0.11.0 still tracked the constvar count independently of attached values. Confirmed against the 0.11.0 sdist rather than inferred.

interop — opaque failure for installed-but-broken optional deps

lazy_import() caught only ImportError, so a framework that is present yet fails to import surfaced as a traceback through its own internals. It now raises InteropError naming the package, the underlying error, and the installed JAX version. MissingDependencyError still signals a genuinely absent package.

Known issue: flax

flax is not yet compatible with JAX 0.11.1. The latest release (0.12.8) fails to import under it (jax.experimental.hijax.MutableHiType no longer exists). This is upstream with no brainstate-side fix; flax interop is unavailable on 0.11.1 until they ship a compatible release. equinox interop is unaffected. The affected tests now skip cleanly instead of producing 38 collection errors. Pin jax==0.11.0 if flax interop is required.

CI

Added a pinned jax==0.11.0 matrix entry (one per supported minor, plus unpinned latest, which now exercises 0.11.1). The jax>=0.8.0 floor is unchanged — no packaging constraint needed changing.

Verification

jax	result
0.11.1	5303 passed, 60 skipped, 0 failed
0.11.0	5342 passed, 21 skipped, 0 failed
0.10.0	IR + interop: 211 passed, 0 failed

The 39-test delta between 0.11.1 and 0.11.0 is exactly the flax interop suite — which also confirms the new skip helper does not over-skip when flax works.

Summary by Sourcery

Restore compatibility with JAX 0.11.1 while hardening IR validation and optional framework interoperability.

Bug Fixes:

• Restore brainstate imports and compatibility with JAX 0.11.1 after removal of CallPrimitive.
• Fix IR constant-variable validation and ClosedJaxpr reconstruction across supported JAX versions.
• Report actionable errors when installed optional frameworks fail to import due to version incompatibilities.

Enhancements:

• Preserve optional-framework test behavior by skipping unusable integrations cleanly, including flax under JAX 0.11.1.

CI:

• Add JAX 0.11.0 to the CI compatibility matrix while continuing to test the latest release.

Documentation:

• Document the JAX 0.11.1 compatibility fixes, flax limitation, supported-version verification, and v0.5.4 release.

Tests:

• Update IR tests for JAX-specific constvar behavior and optional interop tests for installed-but-unusable dependencies.

Chores:

• Bump the project version to 0.5.4.

[sourcery-ai]

Reviewer's Guide

Restores compatibility with JAX 0.11.1, fixes a const/var validation hole in IR processing, hardens optional-framework import error handling, adjusts CI to exercise 0.11.1, and cuts release v0.5.4.

Sequence diagram for lazy_import error handling in interop

sequenceDiagram
    participant User
    participant lazy_import
    participant importlib
    participant MissingDependencyError
    participant InteropError

    User->>lazy_import: lazy_import(module_name)
    lazy_import->>importlib: import_module(module_name)
    alt import raises ImportError
        importlib-->>lazy_import: ImportError
        lazy_import->>MissingDependencyError: raise MissingDependencyError(top, hint)
    else import raises other Exception
        importlib-->>lazy_import: Exception
        lazy_import->>InteropError: raise InteropError(message)
    else import succeeds
        importlib-->>lazy_import: module
        lazy_import-->>User: module
    end

Loading

File-Level Changes

Change	Details	Files
Refactor Jaxpr binder resolution so constvars/invars are resolved explicitly and consts-length validation works across JAX versions, including 0.11.1 where open Jaxprs no longer expose constvars.	Introduce _resolve_jaxpr_vars(eqns, invars, outvars, constvars) to compute and validate binder lists Update eqns_to_jaxpr to use _resolve_jaxpr_vars and document constvars/invars behavior across JAX versions Change eqns_to_closed_jaxpr to call _resolve_jaxpr_vars directly, validate len(consts) against resolved constvars, and construct Jaxpr manually Extend _ir_processing_test to assert version-conditional constvars placement and that ClosedJaxpr reconstruction restores the constvars/invars split	brainstate/transform/_ir_processing.py brainstate/transform/_ir_processing_test.py
Introduce a compatibility shim for the removed jax.extend.core.CallPrimitive so brainstate imports successfully on JAX 0.11.1 while preserving the CallPrimitive public API and constructor semantics.	Stop importing CallPrimitive directly from jax.extend.core in the main import block Add a guarded import that uses jax.extend.core.CallPrimitive on jax<0.11.1 and defines a CallPrimitive(Primitive) subclass on jax>=0.11.1 In the shim subclass, set call_primitive=True and register standard call-primitive rules via register_call_primitive_rules in init	brainstate/_compatible_import.py
Strengthen optional-framework import handling so installed-but-broken packages raise a clear InteropError instead of surfacing opaque tracebacks.	Change lazy_import to import InteropError and update its docstring to describe new failure modes Add a broad Exception handler in lazy_import that raises InteropError with a message describing the failing package, underlying exception, and installed JAX version	brainstate/interop/_common.py
Make interop conversion tests treat frameworks that fail to import for non-ImportError reasons as unavailable, skipping their tests cleanly (e.g., flax under JAX 0.11.1).	Add an _importorskip helper using importlib.import_module and pytest.skip on any exception to mark a framework as unusable Replace pytest.importorskip(..., exc_type=ImportError) calls in conversion tests with _importorskip for flax.nnx, flax.linen, and equinox	brainstate/interop/_conversion_test.py
Cut release v0.5.4 with updated version metadata, changelog, and CI matrix that explicitly pins JAX 0.11.0 while the latest entry exercises 0.11.1.	Update version to 0.5.4 in the version module Add a 0.5.4 section to changelog.md describing JAX 0.11.1 compatibility, IR and interop fixes, CI updates, and verification results Extend the GitHub Actions CI matrix to include a pinned jax==0.11.0 entry in addition to existing minors and the latest slot	brainstate/_version.py changelog.md .github/workflows/CI.yml

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've left some high level feedback:

• The new _resolve_jaxpr_vars helper currently returns a bare tuple; consider tightening the return type annotation to Tuple[List[Var], List[Var], List[Var]] (and importing Tuple) to make downstream usage clearer and more type-checker-friendly.
• In _conversion_test._importorskip, catching all Exception and immediately skipping can make debugging unexpected errors harder; you might want to special-case known interop issues (e.g., AttributeError from flax) or at least log the underlying exception somewhere before skipping.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- The new `_resolve_jaxpr_vars` helper currently returns a bare `tuple`; consider tightening the return type annotation to `Tuple[List[Var], List[Var], List[Var]]` (and importing `Tuple`) to make downstream usage clearer and more type-checker-friendly.
- In `_conversion_test._importorskip`, catching all `Exception` and immediately skipping can make debugging unexpected errors harder; you might want to special-case known interop issues (e.g., `AttributeError` from flax) or at least log the underlying exception somewhere before skipping.

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[codecov]

Codecov Report

❌ Patch coverage is 80.95238% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
brainstate/_compatible_import.py	77.77%	2 Missing ⚠️
brainstate/interop/_common.py	33.33%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!