Version 0.2.7
BrainState 0.2.7 modernizes the experimental compilation stack, deepens the transformation APIs, and tightens runtime infrastructure across the project.
Experimental Compiler and Visualization
- Introduced the experimental
neuroircompiler built on dataclass-based graph IR elements and an explicitCompilationContext, improving dependency tracking, hidden-state mapping, and ClosedJaxpr fidelity even for self-connections and delay buffers. - Added GraphDisplayer and TextDisplayer backends with hierarchical and force-directed layouts, plus richer diagnostics and tests that cover large sample networks and neuro-graph visualizations.
Transformations and Autodiff
- Added the
jit_named_scopedecorator and supporting utilities so nested transformations emit meaningful names inside traced functions, together with_make_jaxprrefinements that separate dynamic/static arguments and improve caching semantics forStatefulFunction. - Expanded the gradient toolkit by exporting the new Jacobian (forward and reverse), Hessian, and SOFO transforms, unifying gradient handling for classes, auxiliary returns, and state-aware updates through the transform module.
State and Runtime Enhancements
- Replaced the experimental
ArrayParamwith a dedicatedDelayState, propagating the new state through the compiler, delay modules, and neuro-IR so historical buffers participate in tracing and optimization just like other states. - Environment helpers can now run against injected
EnvironmentStateinstances, enabling sandboxed or per-thread configurations while DelayState-aware unit tests extend coverage of the updated modules.
Experimental and Infrastructure Updates
- Completed the neuron IR → neuroir rename, aligned the GDiist BPU codebase with the new terminology, and added new sample networks plus placeholder skips to keep the growing compiler/displayer test surface manageable.
- Added
braincellto the development requirements, refreshed documentation wording, and kept CI dependencies current for the GitHub Actions runners.
Bug Fixes
- Hardened caching, randomness, and initialization logic by fixing
get_arg_cache_key, removing stale decorator parameters, validating truncated normal draws, and correcting the exported version metadata. - Declared Python 3.14 support and cleaned up compiler import ordering to keep linting noise low.
What's Changed
- Add Jaxpr optimization passes and integrate constant folding by @chaoming0625 in #114
- Enhance compilation with BPU parser, IR optimizations, and device wrappers by @chaoming0625 in #115
- deps(deps): bump actions/checkout from 5 to 6 by @dependabot[bot] in #116
- update document by @xinzhu-L in #117
- Support custom EnvironmentState instances in environ helpers by @chaoming0625 in #118
- Add unified gradient transforms with Jacobian, Hessian, SOFO by @chaoming0625 in #119
- Refactor(state): remove ArrayParam, add DelayState by @chaoming0625 in #122
- feat(neuroir): add NeuroIR core module and API by @chaoming0625 in #123
Full Changelog: v0.2.4...v0.2.7