release: v0.6.0 by chaoming0625 · Pull Request #247 · chaobrain/brainstate

chaoming0625 · 2026-10-20T11:36:46Z

Release brainstate 0.6.0.

Changes in this PR: the version goes to 0.6.0, the minimum Python version goes to 3.12, the Python 3.15 classifier is added, and the 0.6.0 changelog is finalized.

After merging, create the v0.6.0 GitHub release. That triggers Publish.yml, which uploads to PyPI.

brainstate 0.6.0 replaces the pretty representation of modules and states, fixes about thirty correctness bugs found in an audit of the codebase, and supports JAX 0.9 through 0.12. The minimum supported versions are now Python 3.12 and JAX 0.9.0. A few of the fixes change results that existing code may depend on; read Upgrade notes before upgrading.

Highlights

Compact, width-aware repr for Module, Node, State and the brainstate.util containers. It summarises arrays by shape and dtype and adds a parameter-count footer, and can be configured with the new brainstate.util.repr_options() (feat: redesign pretty repr for Node, Module, State and friends #246).
Correctness fixes in State hooks, environ, random, nn and transform (fix: High/Medium audit bugs across state, environ, random, nn, transform #244).
import brainstate respects jax_enable_x64: importing the package no longer turns float64 off (Respect jax_enable_x64 when brainstate is imported #242).
JAX 0.9.0–0.12.0 supported and tested in CI (fix: jax 0.11.2 compatibility, raise floor to jax>=0.9.0 #240, ci: pin jax 0.12.0 in test matrix; add jax 0.12 compatibility audit #243).

Upgrade notes

Most code needs no changes. Check the following if they apply to you:

Area	What changed	What to do
Python	The minimum Python version is now 3.12 (was 3.11).	Upgrade Python, or stay on brainstate 0.5.x.
Dependencies	The minimum JAX version is now `jax>=0.9.0` (was `0.8.0`).	Upgrade JAX, or stay on brainstate 0.5.x.
`brainstate.random.pareto()`	It now samples Pareto II (Lomax), with support `[0, inf)`, the same as `numpy.random.pareto`. Earlier samples were 1 too large.	To reproduce the old values, add 1 to the result.
`nn.MaxPool*(return_indices=True)`	Indices now use a layout that does not depend on the input size.	Indices are unchanged for unpadded inputs whose size divides evenly. Otherwise, don't use them to index the input directly (`x.ravel()[indices]`). Pass them to `MaxUnpool*`.
`brainstate.environ`	Settings are shared by all threads in the process. Only `context()` overrides are still per thread.	If code relied on each thread having its own settings, use `environ.context(...)` inside each thread.
Float64 at import	brainstate no longer forces `jax_enable_x64=False` when imported.	Nothing to do with JAX's default. If you turned x64 on before importing brainstate, it now stays on.
`mixin.Mode` equality	`Batching` instances are equal only if `batch_size` and `batch_axis` match. `JointMode` instances compare their component modes.	When a `Mode` is a static `jit` argument, a different batch size now triggers a retrace.
`repr()` / `print()` output	The output format is different (see below).	Update any doctests or string comparisons that check `repr` output.

New Features

Pretty representation redesign (#246)

Node, nn.Module, State, TreefyState and the brainstate.util containers now share one width-aware layout engine. A group goes on a single line when it fits within the line width. If it doesn't, each item goes on its own line and is followed by a comma. Arrays are shown as dtype and shape. Trailing comments give per-child and total parameter counts:

>>> import brainstate
>>> net = brainstate.nn.Sequential(
...     brainstate.nn.Linear(3, 4), brainstate.nn.ReLU(), brainstate.nn.Linear(4, 2))
>>> net
Sequential(
  (3,) -> (2,),
  layers=[
    [0] Linear((3,) -> (4,), weight=ParamState{bias: f32[4], weight: f32[3,4]}),  # 16 params
    [1] ReLU(),
    [2] Linear((4,) -> (2,), weight=ParamState{bias: f32[2], weight: f32[4,2]}),  # 10 params
  ],
)  # 26 params (104 B)

Formatting:
- Arrays look like f32[3,4], scalars like f32(0.5), and brainunit quantities like f32[3] mV.
- Modules show their input and output sizes as (3,) -> (4,).
- A state's value is shown positionally, and a dict value uses braces: ParamState{...}.
- Sequences of objects get [i] index labels.
- An object that has already appeared is shown as <shared: path>, and a reference back to an enclosing object as <cycle: path>.
Summary: the parameter and state-value counts and byte sizes are computed from unique states. Computing them does not trigger state read hooks.
New options API in brainstate.util:
- ReprOptions has the fields width, max_depth, max_items, summary, color and short_dtypes.
- set_repr_options() and get_repr_options() set and read the process-wide defaults.
- The repr_options() context manager applies a nestable override for the current thread only.
```
>>> with brainstate.util.repr_options(max_depth=0):
...     print(repr(net))
Sequential((3,) -> (2,), layers=[Linear(...), ReLU(...), Linear(...)])  # 26 params (104 B)
```
brainstate.util.pretty_format(value) formats any value with the same engine, including lists, dicts, dataclasses and arrays.
Optional ANSI color, applied only by str() and print(), never by repr(). Turn it on with set_repr_options(color=True) or BRAINSTATE_REPR_COLOR=1. The NO_COLOR environment variable turns it off.
The custom __pretty_repr__ protocol is backward compatible.

The redesign also fixes these representations:

TreefyState printed full arrays and its internal hook manager.
nn.Map and nn.Param printed attributes that were None.
mixin.Batching labelled batch_size as in_size.
interop.PermuteAxes printed as Transpose.
HiData subclasses printed as HiData.

Bug Fixes

`brainstate.State` and hooks (#244)

Copying a state no longer shares its hooks or tags:
- State.update_from_ref() used to replace the target's hook manager with the source's. It now keeps the target's own.
- State.replace() and TreefyState.to_state() now get a fresh hook manager.
- copy(), replace(), copy_from() and to_state() now copy the tag set instead of sharing it.
Raising HookCancellationError in a hook cancels the write. It used to be wrapped in HookExecutionError, so the write went ahead anyway. Logged hook errors now report the original exception.
Hook registration and execution are thread-safe. The hook cache is rebuilt under a lock, and callbacks run outside the lock, which prevents deadlocks.
StateTraceStack no longer hides exceptions. If an exception was raised inside the block, leaving the block now lets it propagate instead of replacing it with an "out of order" RuntimeError.

`brainstate.mixin` (#244)

Mode.__eq__ and Mode.__hash__ take parameters into account, so Batching(2) != Batching(5). Before, jit could reuse code compiled for a different batch size. Comparing a mode with a non-Mode object returns NotImplemented.

`brainstate.environ` (#242, #244)

import brainstate no longer disables float64. The global precision now starts from JAX's setting: 64 if x64 is on, otherwise 32. Importing the package doesn't change the JAX config. Explicit set(precision=...), set_precision() and context(precision=...) still control jax_enable_x64.
Settings are process-global. Worker threads used to start with empty settings (get_dt() raised KeyError) while jax_enable_x64 was set for the whole process. If context(precision=...) blocks overlap in different threads, one thread no longer turns x64 off while another still needs it.

`brainstate.random` (#244)

pareto() samples Pareto II (Lomax), matching its documentation and NumPy (see Upgrade notes).
multinomial() no longer drops draws when pvals does not sum to 1. As in NumPy, the last probability is taken to be 1 - sum(pvals[:-1]), so the counts in every row sum to n.
zipf() accepts size as a list.

`brainstate.nn` (#244)

BatchNorm* with track_running_stats=False now normalizes using batch statistics. Before, it returned its input unchanged.
BatchNorm* running statistics are no longer corrupted when feature_axis is not the last axis.
MaxUnpool* restores odd-sized and padded inputs correctly without output_size. Before, values were misplaced, sometimes into other batch elements. output_size is still honored (see Upgrade notes for the change to the index format).
MaxPool*(return_indices=True, in_size=...) no longer fails in the constructor.
Delay computes the correct number of steps for float delays. If delay / dt is within 1e-6 of an integer, it is now rounded instead of rounded up, so 0.3 / 0.1 gives 3 steps instead of 4. update_every gets the same fix.
init_maybe_prefetch() accepts PrefetchDelay targets instead of raising TypeError.

`brainstate.transform` (#244)

IR optimizations keep side effects.
- Dead-code elimination no longer removes equations with effects, such as jax.debug.print, host callbacks and jit_error_if.
- Constant folding no longer runs effectful equations, including those inside a nested jit, while the optimization pass is running.
The x - x → 0 and 0 * x → 0 simplifications only apply to 0-d integer outputs. They used to produce outputs of the wrong shape and to turn inf - inf into 0 instead of NaN.
vmap_new_states, vmap2_new_states and pmap2_new_states detect mapped tracers leaking into existing states. When the function writes a mapped value into a state that existed before the call, that state is rolled back and a clear error is raised. States excluded with state_to_exclude are now cleaned up properly instead of keeping invalid tracers.
Batched and unbatched leaves can be mixed in one state under the mapping transforms. Unbatched leaves are broadcast.
jit works out static_argnums from static_argnames (and the reverse) using the function signature, like jax.jit. The same applies to donate_argnums and donate_argnames. A static argument passed by position is now treated as static.
StatefulFunction and make_jaxpr accept negative static_argnums.
Python code generated for scan with constants runs correctly (fn_to_python_code(), jaxpr_to_python_code()). This also covers bodies containing operations that only use constants, and nested jaxprs with unused inputs.
Jaxprs built by brainstate always have a DebugInfo attached. JAX has deprecated jaxprs without one and plans to make them an error (jax-ml/jax#26480). This affects eqns_to_jaxpr(), eqns_to_closed_jaxpr(), optimize_jaxpr(), fn_to_python_code() and jaxpr_to_python_code() (fix: jax 0.11.2 compatibility, raise floor to jax>=0.9.0 #240).

Build & CI

The minimum Python version is raised to 3.12 (requires-python = '>=3.12'). The Python 3.11 classifier is removed and a Python 3.15 classifier is added.
The minimum JAX version is raised to jax>=0.9.0 in pyproject.toml and requirements.txt, and the code paths for older JAX versions are removed (fix: jax 0.11.2 compatibility, raise floor to jax>=0.9.0 #240).
The CI matrix tests JAX 0.9.0, 0.10.0, 0.11.0 and 0.12.0, plus the latest release. The jax==0.8.0 entry is dropped (fix: jax 0.11.2 compatibility, raise floor to jax>=0.9.0 #240, ci: pin jax 0.12.0 in test matrix; add jax 0.12 compatibility audit #243).
Tests now fail if any code builds a jaxpr without DebugInfo (fix: jax 0.11.2 compatibility, raise floor to jax>=0.9.0 #240).
Interop tests are skipped, rather than failing, when an installed optional framework (for example Flax) requires a newer JAX than the one under test (test: skip interop tests when a framework's declared jax requirement is unmet #241).
Fixed file ownership in the documentation deployment workflow (deploy: chown release to root so files are not owned by the runner uid 1001 #238).

Quality

Full test suite, including slow tests, on Python 3.13:
- JAX 0.12.0: 5529 passed, 41 skipped, 0 failed (420 subtests).
- JAX 0.9.0: 5461 passed, 109 skipped, 0 failed (420 subtests). Most of the extra skips are Flax interop tests, because the installed Flax requires jax>=0.11.1, and visualization tests that need pydot.
mypy reports no issues in the 108 checked source files.

Summary by Sourcery

Release brainstate 0.6.0 with redesigned representations, broad correctness fixes, and support for current Python and JAX versions.

New Features:

Introduce compact, width-aware representations for modules, nodes, states, and utility containers with configurable formatting and summaries.

Bug Fixes:

Correct state hook and tag copying, hook cancellation and concurrency, environment precision and thread behavior, random sampling, neural-network operations, and transformation correctness across the library.

Enhancements:

Preserve JAX side effects during IR optimization and attach debug information to generated jaxprs.
Respect JAX's existing float64 configuration when importing brainstate.

Build:

Raise the minimum supported versions to Python 3.12 and JAX 0.9.0, remove older-JAX compatibility paths, and publish the 0.6.0 package with Python 3.12.

CI:

Expand CI coverage across JAX 0.9 through 0.12, strengthen jaxpr validation, and skip incompatible optional-framework interop tests.

Deployment:

Update the publishing workflow to use Python 3.12.

Documentation:

Finalize the 0.6.0 release notes with upgrade guidance, feature highlights, and compatibility changes.

Tests:

Expand the validated release test coverage and report clean full-suite results across supported JAX versions.

Chores:

Update the project version to 0.6.0 and align mypy with Python 3.12.

- Bump version to 0.6.0. - Raise the minimum Python version to 3.12; add the Python 3.15 classifier. - Finalize the 0.6.0 changelog: highlights, upgrade notes, pretty repr redesign (#246), audit fixes (#244), JAX 0.9-0.12 support.

sourcery-ai · 2026-10-20T11:36:53Z

Reviewer's Guide

This release-only PR publishes brainstate 0.6.0 by updating version and packaging metadata, raising the supported Python floor to 3.12, aligning the PyPI workflow and type-checking configuration, and finalizing comprehensive release notes for the representation redesign, correctness fixes, JAX compatibility, upgrade impacts, and validation results.

File-Level Changes

Change	Details	Files
Finalized the 0.6.0 release metadata and documentation, including compatibility requirements, upgrade guidance, feature highlights, bug fixes, CI coverage, and validation results.	Bumped the package version to 0.6.0 and finalized the changelog. Raised the minimum Python version to 3.12, removed the 3.11 classifier, and added the 3.15 classifier. Updated mypy to target Python 3.12. Documented the new representation system, correctness fixes, behavioral changes, and upgrade notes. Recorded JAX 0.9–0.12 support, CI changes, and test-quality results.	`brainstate/_version.py` `changelog.md` `pyproject.toml`
Aligned the publishing workflow with the new supported Python runtime.	Changed the release publishing job to use Python 3.12.	`.github/workflows/Publish.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 found 1 issue

Prompt for AI Agents

Please address the comments from this code review:

## Individual Comments

### Comment 1
<location path="pyproject.toml" line_range="28" />
<code_context>
 description = "A State-based Transformation System for Brain Modeling."
 readme = 'README.md'
-requires-python = '>=3.11'
+requires-python = '>=3.12'
 authors = [{ name = 'BrainState Developers', email = 'chao.brain@qq.com' }]
 classifiers = [
</code_context>
<issue_to_address>
**Documented Python installs fail**

When a user follows the installation guide on Python 3.10 or 3.11, `pip` enforces `requires-python = '>=3.12'` and rejects installation, so users following the installation guide on Python 3.10 or 3.11 cannot install this release.

Update the installation guide to state that Python 3.12 or newer is required.
</issue_to_address>

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

sourcery-ai · 2026-10-20T11:39:14Z

 description = "A State-based Transformation System for Brain Modeling."
 readme = 'README.md'
-requires-python = '>=3.11'
+requires-python = '>=3.12'


🟡 Medium · Documented Python installs fail

When a user follows the installation guide on Python 3.10 or 3.11, pip enforces requires-python = '>=3.12' and rejects installation, so users following the installation guide on Python 3.10 or 3.11 cannot install this release.

Update the installation guide to state that Python 3.12 or newer is required.

Prompt for AI agents

In `pyproject.toml` at line 28: **Documented Python installs fail** When a user follows the installation guide on Python 3.10 or 3.11, `pip` enforces `requires-python = '>=3.12'` and rejects installation, so users following the installation guide on Python 3.10 or 3.11 cannot install this release. Update the installation guide to state that Python 3.12 or newer is required.

codecov · 2026-10-20T11:49:56Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

release: v0.6.0

3c58b45

- Bump version to 0.6.0. - Raise the minimum Python version to 3.12; add the Python 3.15 classifier. - Finalize the 0.6.0 changelog: highlights, upgrade notes, pretty repr redesign (#246), audit fixes (#244), JAX 0.9-0.12 support.

sourcery-ai Bot reviewed Oct 20, 2026

View reviewed changes

chaoming0625 merged commit 521ef49 into main Oct 20, 2026
8 checks passed

chaoming0625 deleted the release/v0.6.0 branch October 20, 2026 13:08

Navigation Menu

Repository navigation

release: v0.6.0 - #247

release: v0.6.0#247
chaoming0625 merged 1 commit into
mainfrom
release/v0.6.0

chaoming0625 commented Oct 20, 2026 •

edited by sourcery-ai Bot

Loading

Uh oh!

sourcery-ai Bot commented Oct 20, 2026 •

edited

Loading

Interacting with Sourcery

Customizing Your Experience

Getting Help

Uh oh!

sourcery-ai Bot left a comment

Uh oh!

sourcery-ai Bot Oct 20, 2026

Uh oh!

codecov Bot commented Oct 20, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

chaoming0625 commented Oct 20, 2026 • edited by sourcery-ai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Highlights

Upgrade notes

New Features

Pretty representation redesign (#246)

Bug Fixes

brainstate.State and hooks (#244)

brainstate.mixin (#244)

brainstate.environ (#242, #244)

brainstate.random (#244)

brainstate.nn (#244)

brainstate.transform (#244)

Build & CI

Quality

Summary by Sourcery

Uh oh!

sourcery-ai Bot commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviewer's Guide

File-Level Changes

Interacting with Sourcery

Customizing Your Experience

Getting Help

Uh oh!

sourcery-ai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

sourcery-ai Bot Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Oct 20, 2026

Codecov Report

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

chaoming0625 commented Oct 20, 2026 •

edited by sourcery-ai Bot

Loading

`brainstate.State` and hooks (#244)

`brainstate.mixin` (#244)

`brainstate.environ` (#242, #244)

`brainstate.random` (#244)

`brainstate.nn` (#244)

`brainstate.transform` (#244)

sourcery-ai Bot commented Oct 20, 2026 •

edited

Loading