What problem does it solve? Building neural models in JAX requires threading mutable state through functional transforms, which is error-prone. This Skill guides the construction of BrainState models where mutable State, Module graphs, environment-scoped simulations, and state-aware jit/grad/vmap transforms work together correctly. ## Core Features & Use Cases - Stateful Module Construction: Create Modules with State roles (HiddenState, ShortTermState, ParamState), registered child graphs, and size-aware Sequential composition with .desc() inference. - State-Aware Transforms: Apply brainstate.transform jit, grad, and vmap that track State reads and writes, enabling compiled training steps and batched ensembles. - Simulation Environments: Run environment-scoped simulations with dt, time, and fit-mode contexts, collective state initialization/reset, and reproducible randomness. - Use Case: Train a custom RNN on temporal integration tasks by defining a Module with trainable ParamState weights, wrapping the loss in brainstate.transform.grad, and compiling the training step with jit. ## Quick Start Use the brainstate skill to build a stateful RNN module with trainable parameters and a jit-compiled training step.
