brainstate.transform module#
The brainstate.transform module provides powerful transformations for
neural computation and scientific computing. It extends JAX’s transformation capabilities
with stateful computation support, enabling efficient compilation, automatic differentiation,
parallelization, and control flow for brain simulation and machine learning workloads.
Browse the reference by category below.