Base Module Classes

Core building blocks for neural network construction. Module is the base class for all components in BrainState, providing utilities for parameter management, state traversal, and hierarchical composition. Sequential enables easy chaining of modules for feedforward architectures.

Module	Base class for neural network modules in BrainState.
ElementWiseBlock	Marker base class for element-wise neural network modules.
Sequential	A sequential input-output module.

Common Wrappers

Utility wrappers for context management and vectorization. EnvironContext manages environment-specific configurations, while Vmap and Map enable efficient batching and vectorization of module operations across multiple inputs.

EnvironContext	Wrap a module so it executes inside a brainstate environment context.
Vmap	Vectorize a module with brainstate.transform.vmap.
Map	Vectorize or parallelize a module using brainstate.transform.vmap2 or pmap2.