⬆️ bump brainstate from 0.3.0 to 0.4.0

[dependabot]

Bumps brainstate from 0.3.0 to 0.4.0.

Release notes

Sourced from brainstate's releases.

Version 0.4.0

This release modernizes brainstate.random with JAX typed PRNG keys and comprehensive physical-unit support, ships inline type information (PEP 561) gated by a mypy CI ratchet, adds a new brainstate.interop module for converting models to/from Flax NNX, Flax Linen, and Equinox, and expands brainstate.transform with several new state-aware transformations (vjp/jvp, shard_map, named_call, and the checkify runtime-check family).

Breaking Changes

• Renamed jit_named_scope to named_scope: The brainstate.transform.jit_named_scope decorator is now exported as brainstate.transform.named_scope. Update any usage accordingly.
• Removed brainstate.transform.sofo_grad: the second-order forward-mode (SOFO) gradient helper has moved to braintools. Replace brainstate.transform.sofo_grad(fn, ...) with the braintools.optim.SOFO optimizer (see examples/009_sofo_mnist.py for the updated usage).
• Removed brainstate.graph.NodeDef and brainstate.graph.NodeRef: the graph representation was reworked. A flattened graph is now described by brainstate.graph.NodeSpec together with the new edge types (NodeEdge, StateEdge, StateLeafEdge, PytreeEdge, StaticEdge, Static). Code that referenced NodeDef/NodeRef directly must migrate to these types; users of the high-level graph.flatten / graph.treefy_split / graph.treefy_merge API are unaffected.

Typed PRNG Keys in brainstate.random

brainstate.random now uses JAX's modern typed PRNG keys (jax.random.key, dtype key<fry>, scalar shape ()) everywhere a key is produced, replacing the legacy raw uint32[2] representation.

• get_key(), split_key(), split_keys(), self_assign_multi_keys(), and RandomState.value now return typed keys. A single key has shape () (was (2,)); a batch of n keys has shape (n,) (was (n, 2)). Code that asserted key.shape == (2,) or key.dtype == uint32, or that indexed the raw words of a key, must be updated.
• Key inputs accept three forms: an integer seed, a typed JAX key, or a legacy uint32[2] array (the last is auto-wrapped via jax.random.wrap_key_data). Passing an integer seed array of size 1 is also accepted. Invalid inputs now raise TypeError (previously ValueError in some paths).
• RandomState remains transform-compatible: typed keys vmap/jit/grad cleanly over their leading axis, and state-aware transformations that special-case RandomState continue to work unchanged.
• The module-level DEFAULT generator still constructs without triggering JAX backend initialization at import time: it holds a lazy uint32[2] placeholder that is materialized into a typed key (via wrap_key_data, preserving the exact seed) on first use.

Migration: to recover the raw uint32[2] words from a typed key, use the new brainstate.random.get_key_data() or jax.random.key_data(key).

New Features

Inline Type Information (PEP 561)

• py.typed marker added: brainstate now ships inline type information, so downstream projects' type checkers (mypy, pyright, etc.) pick up brainstate's annotations automatically.
• Typing correctness gate: a mypy configuration with a per-module "ratchet" enforces type correctness in CI, starting with brainstate.typing. Coverage expands module-by-module over time.
• All annotations are evaluated lazily (from __future__ import annotations), so they impose no import-time or runtime cost.

Physical Unit Support in brainstate.random

Random distributions are now comprehensively and strictly compatible with brainunit physical units, with a consistent location–scale convention.

• Location/scale parameters carry the output unit: normal, laplace, logistic, gumbel, wald, and truncated_normal propagate the unit of their loc/scale (or mean/bounds) into the samples. When only one of loc/scale carries a unit, the plain value is interpreted in that same unit; a compatible-but-different unit (e.g. volt against mV) is converted, while an incompatible one raises UnitMismatchError.
• Scale-only distributions carry the scale unit: exponential, gamma, rayleigh, and weibull_min propagate the unit of their scale parameter.
• multivariate_normal carries the unit of mean (with cov required to be mean-unit squared).
• Shape / rate / count / probability parameters are strictly dimensionless: parameters such as df, a/b, lam, n, p, alpha, logits, kappa, concentration, and friends reject a dimensional Quantity with a clear ValueError. A genuinely dimensionless Quantity (e.g. 3.0 * u.UNITLESS) is accepted.
• No units → plain arrays: every distribution returns a plain array when given plain inputs, so existing unitless code is unaffected.

Raw Key Interop Helper

• brainstate.random.get_key_data() returns the current global key as a raw uint32[2] array (via jax.random.key_data), for interfacing with code that still expects the legacy representation.

Framework Interoperability (brainstate.interop)

A new brainstate.interop module converts modules to and from other JAX frameworks, with an extensible layer registry:

... (truncated)

Changelog

Sourced from brainstate's changelog.

Version 0.4.0 (2026-06-01)

Breaking Changes

• Renamed jit_named_scope to named_scope: The brainstate.transform.jit_named_scope decorator is now exported as brainstate.transform.named_scope. Update any usage accordingly.
• Removed brainstate.transform.sofo_grad: the second-order forward-mode (SOFO) gradient helper has moved to braintools. Replace brainstate.transform.sofo_grad(fn, ...) with the braintools.optim.SOFO optimizer (see examples/009_sofo_mnist.py for the updated usage).
• Removed brainstate.graph.NodeDef and brainstate.graph.NodeRef: the graph representation was reworked. A flattened graph is now described by brainstate.graph.NodeSpec together with the new edge types (NodeEdge, StateEdge, StateLeafEdge, PytreeEdge, StaticEdge, Static). Code that referenced NodeDef/NodeRef directly must migrate to these types; users of the high-level graph.flatten / graph.treefy_split / graph.treefy_merge API are unaffected.

Typed PRNG Keys in brainstate.random

brainstate.random now uses JAX's modern typed PRNG keys (jax.random.key, dtype key<fry>, scalar shape ()) everywhere a key is produced, replacing the legacy raw uint32[2] representation.

• get_key(), split_key(), split_keys(), self_assign_multi_keys(), and RandomState.value now return typed keys. A single key has shape () (was (2,)); a batch of n keys has shape (n,) (was (n, 2)). Code that asserted key.shape == (2,) or key.dtype == uint32, or that indexed the raw words of a key, must be updated.
• Key inputs accept three forms: an integer seed, a typed JAX key, or a legacy uint32[2] array (the last is auto-wrapped via jax.random.wrap_key_data). Passing an integer seed array of size 1 is also accepted. Invalid inputs now raise TypeError (previously ValueError in some paths).
• RandomState remains transform-compatible: typed keys vmap/jit/grad cleanly over their leading axis, and state-aware transformations that special-case RandomState continue to work unchanged.
• The module-level DEFAULT generator still constructs without triggering JAX backend initialization at import time: it holds a lazy uint32[2] placeholder that is materialized into a typed key (via wrap_key_data, preserving the exact seed) on first use.

Migration: to recover the raw uint32[2] words from a typed key, use the new brainstate.random.get_key_data() or jax.random.key_data(key).

New Features

Inline Type Information (PEP 561)

• py.typed marker added: brainstate now ships inline type information, so downstream projects' type checkers (mypy, pyright, etc.) pick up brainstate's annotations automatically.
• Typing correctness gate: a mypy configuration with a per-module "ratchet" enforces type correctness in CI, starting with brainstate.typing. Coverage expands module-by-module over time.
• All annotations are evaluated lazily (from __future__ import annotations), so they impose no import-time or runtime cost.

Physical Unit Support in brainstate.random

Random distributions are now comprehensively and strictly compatible with brainunit physical units, with a consistent location–scale convention.

• Location/scale parameters carry the output unit: normal, laplace, logistic, gumbel, wald, and truncated_normal propagate the unit of their loc/scale (or mean/bounds) into the samples. When only one of loc/scale carries a unit, the plain value is interpreted in that same unit; a compatible-but-different unit (e.g. volt against mV) is converted, while an incompatible one raises UnitMismatchError.
• Scale-only distributions carry the scale unit: exponential, gamma, rayleigh, and weibull_min propagate the unit of their scale parameter.
• multivariate_normal carries the unit of mean (with cov required to be mean-unit squared).
• Shape / rate / count / probability parameters are strictly dimensionless: parameters such as df, a/b, lam, n, p, alpha, logits, kappa, concentration, and friends reject a dimensional Quantity with a clear ValueError. A genuinely dimensionless Quantity (e.g. 3.0 * u.UNITLESS) is accepted.
• No units → plain arrays: every distribution returns a plain array when given plain inputs, so existing unitless code is unaffected.

Raw Key Interop Helper

• brainstate.random.get_key_data() returns the current global key as a raw uint32[2] array (via jax.random.key_data), for interfacing with code that still expects the legacy representation.

Framework Interoperability (brainstate.interop)

A new brainstate.interop module converts modules to and from other JAX frameworks, with an extensible layer registry:

... (truncated)

Commits

• 4ecbb68 docs: use braintools.init instead of brainstate.nn.init; drop init API page (...
• 5c7c7ff release: prep 0.4.0 — fix publish version check, finalize changelog, pin deps...
• 74d70bb fix: stop emitting jax.core.DropVar deprecation warning (#196)
• ff5f0f7 docs: update navigation structure and add new how-to guides
• 33fe069 worktree docs reorganization (#195)
• 0a852a5 docs(api): reorganize API reference; split nn/transform, add interop & hooks ...
• 8492403 docs: reformat Google/Sphinx docstrings to NumPy style (#193)
• 9498284 feat(transform): add tqdm-free fallback for ProgressBar, make tqdm optional (...
• 93aeb75 docs(api): add docstrings and type annotations across public APIs (#191)
• 390cbcd feat(interop): convert models between brainstate.nn and flax.nnx / flax.linen...
• Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)