Skip to main content

Ctrl+K

Tutorial Series

Part 1: Basics
Part 2: Neural Network
Part 3: Transformations
Part 4: Utilities
Part 5: Examples
Part 6: Migration Guide
- Migrating Concepts from PyTorch to BrainState

API Reference

Release Notes
brainstate module
brainstate.graph moduel
brainstate.transform module
brainstate.nn module
- Module
- ElementWiseBlock
- Sequential
- Param
- Const
- Transform
- IdentityT
- ClipT
- AffineT
- SigmoidT
- TanhT
- SoftsignT
- ScaledSigmoidT
- SoftplusT
- NegSoftplusT
- LogT
- ExpT
- ReluT
- PositiveT
- NegativeT
- PowerT
- OrderedT
- SimplexT
- UnitVectorT
- ChainT
- MaskedT
- Regularization
- L1Reg
- L2Reg
- ElasticNetReg
- HuberReg
- GroupLassoReg
- TotalVariationReg
- MaxNormReg
- EntropyReg
- OrthogonalReg
- SpectralNormReg
- ChainedReg
- GaussianReg
- StudentTReg
- CauchyReg
- UniformReg
- BetaReg
- LogNormalReg
- ExponentialReg
- GammaReg
- InverseGammaReg
- LogUniformReg
- HorseshoeReg
- SpikeAndSlabReg
- DirichletReg
- EnvironContext
- Vmap
- ModuleMapper
- Linear
- ScaledWSLinear
- SignedWLinear
- SparseLinear
- LoRA
- AllToAll
- OneToOne
- Conv1d
- Conv2d
- Conv3d
- ScaledWSConv1d
- ScaledWSConv2d
- ScaledWSConv3d
- ConvTranspose1d
- ConvTranspose2d
- ConvTranspose3d
- Flatten
- Unflatten
- AvgPool1d
- AvgPool2d
- AvgPool3d
- MaxPool1d
- MaxPool2d
- MaxPool3d
- MaxUnpool1d
- MaxUnpool2d
- MaxUnpool3d
- LPPool1d
- LPPool2d
- LPPool3d
- AdaptiveAvgPool1d
- AdaptiveAvgPool2d
- AdaptiveAvgPool3d
- AdaptiveMaxPool1d
- AdaptiveMaxPool2d
- AdaptiveMaxPool3d
- ReflectionPad1d
- ReflectionPad2d
- ReflectionPad3d
- ReplicationPad1d
- ReplicationPad2d
- ReplicationPad3d
- ZeroPad1d
- ZeroPad2d
- ZeroPad3d
- ConstantPad1d
- ConstantPad2d
- ConstantPad3d
- CircularPad1d
- CircularPad2d
- CircularPad3d
- BatchNorm0d
- BatchNorm1d
- BatchNorm2d
- BatchNorm3d
- LayerNorm
- RMSNorm
- GroupNorm
- weight_standardization
- Dropout
- Dropout1d
- Dropout2d
- Dropout3d
- AlphaDropout
- FeatureAlphaDropout
- DropoutFixed
- Embedding
- Threshold
- ReLU
- RReLU
- Hardtanh
- ReLU6
- Sigmoid
- Hardsigmoid
- Tanh
- SiLU
- Mish
- Hardswish
- ELU
- CELU
- SELU
- GLU
- GELU
- Hardshrink
- LeakyReLU
- LogSigmoid
- Softplus
- Softshrink
- PReLU
- Softsign
- Tanhshrink
- Softmin
- Softmax
- Softmax2d
- LogSoftmax
- Identity
- SpikeBitwise
- tanh
- relu
- squareplus
- softplus
- soft_sign
- sigmoid
- silu
- swish
- log_sigmoid
- elu
- leaky_relu
- hard_tanh
- celu
- selu
- gelu
- glu
- logsumexp
- log_softmax
- softmax
- standardize
- one_hot
- relu6
- hard_sigmoid
- hard_silu
- hard_swish
- hard_shrink
- rrelu
- mish
- soft_shrink
- prelu
- tanh_shrink
- softmin
- sparse_plus
- sparse_sigmoid
- FixedNumConn
- EventFixedNumConn
- EventFixedProb
- EventLinear
- RNNCell
- ValinaRNNCell
- GRUCell
- MGUCell
- LSTMCell
- URLSTMCell
- DynamicsGroup
- Dynamics
- Prefetch
- PrefetchDelay
- PrefetchDelayAt
- OutputDelayAt
- Delay
- DelayAccess
- StateWithDelay
- call_order
- call_all_fns
- vmap_call_all_fns
- init_all_states
- vmap_init_all_states
- reset_all_states
- vmap_reset_all_states
- assign_state_values
- exp_euler_step
- MetricState
- Metric
- AverageMetric
- WelfordMetric
- AccuracyMetric
- MultiMetric
- PrecisionMetric
- RecallMetric
- F1ScoreMetric
- ConfusionMatrix
- HiData
- count_parameters
- clip_grad_norm
- param
- calculate_init_gain
- ZeroInit
- Constant
- Identity
- Normal
- TruncatedNormal
- Uniform
- VarianceScaling
- KaimingUniform
- KaimingNormal
- XavierUniform
- XavierNormal
- LecunUniform
- LecunNormal
- Orthogonal
- DeltaOrthogonal
brainstate.random module
- RandomState
- seed
- seed_context
- default_rng
- clone_rng
- set_key
- get_key
- restore_key
- split_key
- split_keys
- self_assign_multi_keys
- rand
- randn
- random
- random_sample
- ranf
- sample
- randint
- random_integers
- rand_like
- randint_like
- randn_like
- choice
- permutation
- shuffle
- beta
- exponential
- gamma
- gumbel
- laplace
- logistic
- normal
- pareto
- standard_cauchy
- standard_exponential
- standard_gamma
- standard_normal
- standard_t
- uniform
- truncated_normal
- lognormal
- power
- rayleigh
- triangular
- vonmises
- wald
- weibull
- weibull_min
- maxwell
- t
- loggamma
- bernoulli
- binomial
- categorical
- geometric
- hypergeometric
- logseries
- multinomial
- negative_binomial
- poisson
- zipf
- chisquare
- dirichlet
- f
- multivariate_normal
- noncentral_chisquare
- noncentral_f
- orthogonal
brainstate.util module
brainstate.typing module
- Array
- ArrayLike
- Shape
- Size
- Axes
- DType
- DTypeLike
- SupportsDType
- PyTree
- Key
- PathParts
- Predicate
- Filter
- FilterLiteral
- SeedOrKey
- K
- _T
- _Annotation
- Missing
brainstate.mixin module
brainstate.environ module
- reset
- context
- get
- set
- pop
- all
- EnvironmentState
- get_platform
- set_platform
- get_host_device_count
- set_host_device_count
- get_precision
- set_precision
- dftype
- ditype
- dutype
- dctype
- tolerance
- get_dt
- I
- T
- DT
- PRECISION
- PLATFORM
- HOST_DEVICE_COUNT
- JIT_ERROR_CHECK
- FIT
- DEFAULT_PRECISION
- SUPPORTED_PLATFORMS
- SUPPORTED_PRECISIONS
- register_default_behavior
- unregister_default_behavior
- list_registered_behaviors

.rst

ModuleMapper

ModuleMapper#

brainstate.nn.ModuleMapper#: alias of Map

Contents

ModuleMapper

By BrainState Developers

Last updated on May 22, 2026.