Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | R | S | T | U | V | W | X | Z _ __init__() (brainstate.transform.ProgressBar method) __module__ (brainstate.util.NestedDict attribute) _Annotation (class in brainstate.typing) _data (brainstate.util.FrozenDict attribute) _hash (brainstate.util.FrozenDict attribute) _JointGenericAlias (class in brainstate.mixin) _metric_names (brainstate.nn.MultiMetric attribute) _OneOfGenericAlias (class in brainstate.mixin) _T (class in brainstate.typing) A a (brainstate.nn.AffineT attribute) access() (brainstate.nn.Delay method) AccuracyMetric (class in brainstate.nn) AdaptiveAvgPool1d (class in brainstate.nn) AdaptiveAvgPool2d (class in brainstate.nn) AdaptiveAvgPool3d (class in brainstate.nn) AdaptiveMaxPool1d (class in brainstate.nn) AdaptiveMaxPool2d (class in brainstate.nn) AdaptiveMaxPool3d (class in brainstate.nn) add_after_update() (brainstate.nn.Dynamics method) add_before_update() (brainstate.nn.Dynamics method) add_context() (brainstate.nn.EnvironContext method) add_tag() (brainstate.State method) add_unique_key() (brainstate.util.DictManager method) add_unique_value() (brainstate.util.DictManager method) AffineT (class in brainstate.nn) after_updates (brainstate.nn.Dynamics attribute) (brainstate.nn.Dynamics property) algebraic_simplification() (in module brainstate.transform) all (class in brainstate.environ) AllToAll (class in brainstate.nn) AlphaDropout (class in brainstate.nn) append() (brainstate.nn.Sequential method) argname (brainstate.nn.AverageMetric attribute) (brainstate.nn.WelfordMetric attribute) args (brainstate.mixin.ParamDescriber attribute) Array (class in brainstate.typing) ArrayLike (in module brainstate.typing) assign() (brainstate.util.DictManager method) assign_dict_values() (brainstate.util.FlattedDict method) assign_state_vals() (brainstate.StateTraceStack method) assign_state_vals_v2() (brainstate.StateTraceStack method) assign_state_values (class in brainstate.nn) assign_values() (brainstate.StateDictManager method) at() (brainstate.nn.Delay method) (brainstate.nn.PrefetchDelay method) average (brainstate.nn.PrecisionMetric attribute) (brainstate.nn.RecallMetric attribute) AverageMetric (class in brainstate.nn) AvgPool1d (class in brainstate.nn) AvgPool2d (class in brainstate.nn) AvgPool3d (class in brainstate.nn) Axes (in module brainstate.typing) axis_env (brainstate.transform.StatefulFunction attribute) axis_name (brainstate.transform.StatefulMapping attribute) axis_size (brainstate.transform.StatefulMapping attribute) B b (brainstate.nn.AffineT attribute) base_module (brainstate.nn.LoRA attribute) BatchAxisError (class in brainstate) BatchNorm0d (class in brainstate.nn) BatchNorm1d (class in brainstate.nn) BatchNorm2d (class in brainstate.nn) BatchNorm3d (class in brainstate.nn) BatchState (class in brainstate) before_updates (brainstate.nn.Dynamics attribute) (brainstate.nn.Dynamics property) bernoulli (class in brainstate.random) bernoulli() (brainstate.random.RandomState method) beta (class in brainstate.random) beta() (brainstate.random.RandomState method) BetaReg (class in brainstate.nn) binomial (class in brainstate.random) binomial() (brainstate.random.RandomState method) bounded_while_loop() (in module brainstate.transform) brainstate.environ module brainstate.graph module brainstate.mixin module brainstate.nn module brainstate.random module brainstate.transform module brainstate.typing module brainstate.util module BrainStateError (class in brainstate) breakpoint_if() (in module brainstate.transform) C c (brainstate.nn.LSTMCell attribute) cache() (brainstate.nn.Param method) cache_stats (brainstate.nn.Param property) calculate_init_gain (class in brainstate.nn.init) call_all_fns (class in brainstate.nn) call_order (class in brainstate.nn) catch_new_states() (in module brainstate) categorical (class in brainstate.random) categorical() (brainstate.random.RandomState method) CauchyReg (class in brainstate.nn) CELU (class in brainstate.nn) celu (class in brainstate.nn) ChainedReg (class in brainstate.nn) ChainT (class in brainstate.nn) check_state_jax_tracer() (in module brainstate) check_state_value_tree() (in module brainstate) check_valid_context() (brainstate.graph.Node method) check_valid_trace() (brainstate.State method) checkpoint() (in module brainstate.transform) checkpointed_for_loop() (in module brainstate.transform) checkpointed_scan() (in module brainstate.transform) children (brainstate.nn.HiData attribute) children() (brainstate.nn.Module method) chisquare (class in brainstate.random) chisquare() (brainstate.random.RandomState method) choice (class in brainstate.random) choice() (brainstate.random.RandomState method) CircularPad1d (class in brainstate.nn) CircularPad2d (class in brainstate.nn) CircularPad3d (class in brainstate.nn) clear_buffer_memory (class in brainstate.util) clear_cache() (brainstate.nn.Param method) (brainstate.transform.StatefulFunction method) clear_hooks() (brainstate.State method) clip() (brainstate.nn.Param method) clip_grad_norm (class in brainstate.nn) ClipT (class in brainstate.nn) clone (class in brainstate.graph) clone() (brainstate.nn.HiData method) clone_rng (class in brainstate.random) cls (brainstate.mixin.ParamDescriber attribute) collect_values() (brainstate.StateDictManager method) common_subexpression_elimination() (in module brainstate.transform) compute() (brainstate.nn.AverageMetric method) (brainstate.nn.ConfusionMatrix method) (brainstate.nn.F1ScoreMetric method) (brainstate.nn.Metric method), [1] (brainstate.nn.MultiMetric method) (brainstate.nn.PrecisionMetric method) (brainstate.nn.RecallMetric method) (brainstate.nn.WelfordMetric method) cond() (in module brainstate.transform) ConfusionMatrix (class in brainstate.nn) Const (class in brainstate.nn) Constant (class in brainstate.nn.init) constant_fold() (in module brainstate.transform) ConstantPad1d (class in brainstate.nn) ConstantPad2d (class in brainstate.nn) ConstantPad3d (class in brainstate.nn) context (brainstate.nn.EnvironContext attribute) (class in brainstate.environ) contexts (brainstate.environ.EnvironmentState attribute) Conv1d (class in brainstate.nn) Conv2d (class in brainstate.nn) Conv3d (class in brainstate.nn) ConvTranspose1d (class in brainstate.nn) ConvTranspose2d (class in brainstate.nn) ConvTranspose3d (class in brainstate.nn) copy() (brainstate.State method) (brainstate.TreefyState method) (brainstate.util.DotDict method) (brainstate.util.FrozenDict method) copy_from() (brainstate.State method) copy_propagation() (in module brainstate.transform) count (brainstate.nn.AverageMetric attribute) (brainstate.nn.WelfordMetric attribute) count_parameters (class in brainstate.nn) current_inputs (brainstate.nn.Dynamics attribute) D dataclass (class in brainstate.util) dctype (class in brainstate.environ) dead_code_elimination() (in module brainstate.transform) debug_nan() (in module brainstate.transform) debug_nan_if() (in module brainstate.transform) decrease_stack_level() (brainstate.State method) deepcopy() (brainstate.util.DotDict method) DEFAULT_PRECISION (in module brainstate.environ) default_rng (class in brainstate.random) delay (brainstate.nn.Prefetch property) Delay (class in brainstate.nn) DelayAccess (class in brainstate.nn) DelayState (class in brainstate) delta_inputs (brainstate.nn.Dynamics attribute) DeltaOrthogonal (class in brainstate.nn.init) desc() (brainstate.mixin.ParamDesc class method) dftype (class in brainstate.environ) DictManager (class in brainstate.util) difference_by_keys() (brainstate.util.DictManager method) difference_by_values() (brainstate.util.DictManager method) dirichlet (class in brainstate.random) dirichlet() (brainstate.random.RandomState method) DirichletReg (class in brainstate.nn) ditype (class in brainstate.environ) DotDict (class in brainstate.util) draw() (in module brainstate.transform) draw_dot_graph() (in module brainstate.transform) Dropout (class in brainstate.nn) Dropout1d (class in brainstate.nn) Dropout2d (class in brainstate.nn) Dropout3d (class in brainstate.nn) DropoutFixed (class in brainstate.nn) DT (in module brainstate.environ) dtype (brainstate.nn.HiData property) (brainstate.typing.SupportsDType property) DType (in module brainstate.typing) DTypeLike (in module brainstate.typing) dutype (class in brainstate.environ) Dynamics (class in brainstate.nn) DynamicsGroup (in module brainstate.nn) E ElasticNetReg (class in brainstate.nn) elem_indent (brainstate.util.PrettyType attribute) ElementWiseBlock (class in brainstate.nn) ELU (class in brainstate.nn) elu (class in brainstate.nn) Embedding (class in brainstate.nn) embedding_size (brainstate.nn.Embedding attribute) empty_repr (brainstate.util.PrettyType attribute) end (brainstate.util.PrettyAttr attribute) (brainstate.util.PrettyType attribute) EntropyReg (class in brainstate.nn) EnvironContext (class in brainstate.nn) EnvironmentState (class in brainstate.environ) eps (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) (brainstate.nn.ScaledWSLinear attribute) eqns_to_closed_jaxpr() (in module brainstate.transform) eqns_to_jaxpr() (in module brainstate.transform) eval_shape() (in module brainstate.transform) EventFixedNumConn (class in brainstate.nn) EventFixedProb (in module brainstate.nn) EventLinear (class in brainstate.nn) exp_euler_step (class in brainstate.nn) exponential (class in brainstate.random) exponential() (brainstate.random.RandomState method) ExponentialReg (class in brainstate.nn) ExpT (class in brainstate.nn) extend() (brainstate.nn.Sequential method) F f (class in brainstate.random) f() (brainstate.random.RandomState method) F1ScoreMetric (class in brainstate.nn) FakeState (class in brainstate) false_negatives (brainstate.nn.RecallMetric attribute) false_positives (brainstate.nn.PrecisionMetric attribute) FeatureAlphaDropout (class in brainstate.nn) field (class in brainstate.util) Filter (in module brainstate.typing) filter() (brainstate.util.FlattedDict method) (brainstate.util.NestedDict method) filter_by_predicate() (brainstate.util.DictManager method) FilterLiteral (in module brainstate.typing) find() (brainstate.transform.StateFinder method) fit (brainstate.nn.Param attribute) FIT (in module brainstate.environ) fit_hyper (brainstate.nn.Regularization attribute) FixedNumConn (class in brainstate.nn) flat_mapping (class in brainstate.util) FlattedDict (class in brainstate.util) flatten (class in brainstate.graph) Flatten (class in brainstate.nn) fn_to_python_code() (in module brainstate.transform) for_loop() (in module brainstate.transform) forward() (brainstate.nn.AffineT method) (brainstate.nn.ChainT method) (brainstate.nn.ClipT method) (brainstate.nn.ExpT method) (brainstate.nn.IdentityT method) (brainstate.nn.LogT method) (brainstate.nn.MaskedT method) (brainstate.nn.NegativeT method) (brainstate.nn.NegSoftplusT method) (brainstate.nn.OrderedT method) (brainstate.nn.PositiveT method) (brainstate.nn.PowerT method) (brainstate.nn.ReluT method) (brainstate.nn.ScaledSigmoidT method) (brainstate.nn.SigmoidT method) (brainstate.nn.SimplexT method) (brainstate.nn.SoftplusT method) (brainstate.nn.SoftsignT method) (brainstate.nn.TanhT method) (brainstate.nn.Transform method), [1] (brainstate.nn.UnitVectorT method) freeze (brainstate.nn.Embedding attribute) from_dict() (brainstate.nn.HiData class method) (brainstate.util.DotDict class method) from_flat() (brainstate.util.NestedDict class method) from_nest() (brainstate.util.FlattedDict class method) from_pretrained() (brainstate.nn.Embedding class method) FrozenDict (class in brainstate.util) fun (brainstate.transform.StatefulFunction attribute) functions (brainstate.environ.EnvironmentState attribute) fwd_grad() (in module brainstate.transform) G gamma (class in brainstate.random) gamma() (brainstate.random.RandomState method) GammaReg (class in brainstate.nn) GaussianReg (class in brainstate.nn) GELU (class in brainstate.nn) gelu (class in brainstate.nn) geometric (class in brainstate.random) geometric() (brainstate.random.RandomState method) get (class in brainstate.environ) get() (brainstate.util.DotDict method) (brainstate.util.FrozenDict method) get_after_update() (brainstate.nn.Dynamics method) get_arg_cache_key() (brainstate.transform.StatefulFunction method) get_before_update() (brainstate.nn.Dynamics method) get_cache_stats() (brainstate.transform.StatefulFunction method) get_dt (class in brainstate.environ) get_host_device_count (class in brainstate.environ) get_item() (brainstate.nn.Prefetch method) get_item_value() (brainstate.nn.Prefetch method) get_jaxpr() (brainstate.transform.StatefulFunction method) get_jaxpr_by_cache() (brainstate.transform.StatefulFunction method) get_key (class in brainstate.random) get_metadata() (brainstate.TreefyState method) get_out_shapes() (brainstate.transform.StatefulFunction method) get_out_shapes_by_cache() (brainstate.transform.StatefulFunction method) get_out_treedef() (brainstate.transform.StatefulFunction method) get_out_treedef_by_cache() (brainstate.transform.StatefulFunction method) get_platform (class in brainstate.environ) get_precision (class in brainstate.environ) get_read_state_values() (brainstate.StateTraceStack method) get_read_states() (brainstate.StateTraceStack method) (brainstate.transform.StatefulFunction method) get_read_states_by_cache() (brainstate.transform.StatefulFunction method) get_state_trace() (brainstate.transform.StatefulFunction method) get_state_trace_by_cache() (brainstate.transform.StatefulFunction method) get_state_values() (brainstate.StateTraceStack method) get_states() (brainstate.transform.StatefulFunction method) get_states_by_cache() (brainstate.transform.StatefulFunction method) get_value() (brainstate.HiddenGroupState method) (brainstate.HiddenTreeState method) get_write_state_values() (brainstate.StateTraceStack method) get_write_states() (brainstate.StateTraceStack method) (brainstate.transform.StatefulFunction method) get_write_states_by_cache() (brainstate.transform.StatefulFunction method) GLU (class in brainstate.nn) glu (class in brainstate.nn) grad() (in module brainstate.transform) GradientTransform (class in brainstate.transform) graph_to_tree (class in brainstate.graph) GraphDef (class in brainstate.graph) graphdef (class in brainstate.graph) GroupLassoReg (class in brainstate.nn) GroupNorm (class in brainstate.nn) GRUCell (class in brainstate.nn) gumbel (class in brainstate.random) gumbel() (brainstate.random.RandomState method) H h (brainstate.nn.GRUCell attribute) (brainstate.nn.LSTMCell attribute) (brainstate.nn.MGUCell attribute) (brainstate.nn.ValinaRNNCell attribute) hard_shrink (class in brainstate.nn) hard_sigmoid (class in brainstate.nn) hard_silu (class in brainstate.nn) hard_swish (in module brainstate.nn) hard_tanh (class in brainstate.nn) Hardshrink (class in brainstate.nn) Hardsigmoid (class in brainstate.nn) Hardswish (class in brainstate.nn) Hardtanh (class in brainstate.nn) has_after_update() (brainstate.nn.Dynamics method) has_aux (brainstate.transform.GradientTransform attribute) has_before_update() (brainstate.nn.Dynamics method) has_hooks() (brainstate.State method) hashable (class in brainstate.mixin) HashableDict (class in brainstate.mixin) hessian() (in module brainstate.transform) HiData (class in brainstate.nn) HiddenGroupState (class in brainstate) HiddenState (class in brainstate) HiddenTreeState (class in brainstate) history (brainstate.nn.StateWithDelay attribute) hooks (brainstate.State property) HorseshoeReg (class in brainstate.nn) HOST_DEVICE_COUNT (in module brainstate.environ) HuberReg (class in brainstate.nn) hypergeometric (class in brainstate.random) hypergeometric() (brainstate.random.RandomState method) I I (in module brainstate.environ) identifier (brainstate.mixin.ParamDescriber attribute) (brainstate.mixin.ParamDescriber property) Identity (class in brainstate.nn) (class in brainstate.nn.init) IdentityT (class in brainstate.nn) ifelse() (in module brainstate.transform) in_axes (brainstate.transform.StatefulMapping attribute) in_channels (brainstate.nn.Conv1d attribute) (brainstate.nn.Conv2d attribute) (brainstate.nn.Conv3d attribute) (brainstate.nn.ConvTranspose1d attribute) (brainstate.nn.ConvTranspose2d attribute) (brainstate.nn.ConvTranspose3d attribute) (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) in_features (brainstate.nn.LoRA attribute) in_size (brainstate.nn.AllToAll attribute) (brainstate.nn.Conv1d attribute) (brainstate.nn.Conv2d attribute) (brainstate.nn.Conv3d attribute) (brainstate.nn.ConvTranspose1d attribute) (brainstate.nn.ConvTranspose2d attribute) (brainstate.nn.ConvTranspose3d attribute) (brainstate.nn.Dynamics attribute) (brainstate.nn.GRUCell attribute) (brainstate.nn.Linear attribute) (brainstate.nn.LoRA attribute) (brainstate.nn.LSTMCell attribute) (brainstate.nn.MGUCell attribute) (brainstate.nn.Module attribute) (brainstate.nn.OneToOne attribute) (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) (brainstate.nn.ScaledWSLinear attribute) (brainstate.nn.SignedWLinear attribute) (brainstate.nn.SparseLinear attribute) (brainstate.nn.ValinaRNNCell attribute) include_self (brainstate.nn.AllToAll attribute) increase_stack_level() (brainstate.State method) init() (brainstate.mixin.ParamDescriber method) (brainstate.nn.Param class method) init_all_states (class in brainstate.nn) init_state() (brainstate.nn.Delay method) (brainstate.nn.DropoutFixed method) (brainstate.nn.GRUCell method), [1] (brainstate.nn.LSTMCell method), [1] (brainstate.nn.MGUCell method), [1] (brainstate.nn.Module method) (brainstate.nn.RNNCell method) (brainstate.nn.StateWithDelay method) (brainstate.nn.URLSTMCell method), [1] (brainstate.nn.ValinaRNNCell method), [1] inline_jit() (in module brainstate.transform) insert() (brainstate.nn.Sequential method) intersection_by_keys() (brainstate.util.DictManager method) intersection_by_values() (brainstate.util.DictManager method) inverse() (brainstate.nn.AffineT method) (brainstate.nn.ChainT method) (brainstate.nn.ClipT method) (brainstate.nn.ExpT method) (brainstate.nn.IdentityT method) (brainstate.nn.LogT method) (brainstate.nn.MaskedT method) (brainstate.nn.NegativeT method) (brainstate.nn.NegSoftplusT method) (brainstate.nn.OrderedT method) (brainstate.nn.PositiveT method) (brainstate.nn.PowerT method) (brainstate.nn.ReluT method) (brainstate.nn.ScaledSigmoidT method) (brainstate.nn.SigmoidT method) (brainstate.nn.SimplexT method) (brainstate.nn.SoftplusT method) (brainstate.nn.SoftsignT method) (brainstate.nn.TanhT method) (brainstate.nn.Transform method), [1] (brainstate.nn.UnitVectorT method) InverseGammaReg (class in brainstate.nn) is_instance_eval (class in brainstate.util) items() (brainstate.nn.HiData method) (brainstate.util.FrozenDict method) iter_leaf (class in brainstate.graph) iter_node (class in brainstate.graph) J jacfwd() (in module brainstate.transform) jacobian() (in module brainstate.transform) jacrev() (in module brainstate.transform) jaxpr_call() (brainstate.transform.StatefulFunction method) jaxpr_call_auto() (brainstate.transform.StatefulFunction method) jaxpr_to_python_code() (in module brainstate.transform) jit() (in module brainstate.transform) JIT_ERROR_CHECK (in module brainstate.environ) jit_error_if() (in module brainstate.transform) jit_named_scope() (in module brainstate.transform) JointTypes (in module brainstate.mixin) K K (class in brainstate.typing) KaimingNormal (class in brainstate.nn.init) KaimingUniform (class in brainstate.nn.init) kernel_size (brainstate.nn.Conv1d attribute) (brainstate.nn.Conv2d attribute) (brainstate.nn.Conv3d attribute) (brainstate.nn.ConvTranspose1d attribute) (brainstate.nn.ConvTranspose2d attribute) (brainstate.nn.ConvTranspose3d attribute) (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) key (brainstate.util.PrettyAttr attribute) Key (class in brainstate.typing) keys() (brainstate.nn.HiData method) (brainstate.util.FrozenDict method) kwargs (brainstate.mixin.ParamDescriber attribute) L L1Reg (class in brainstate.nn) L2Reg (class in brainstate.nn) laplace (class in brainstate.random) laplace() (brainstate.random.RandomState method) layer (brainstate.nn.EnvironContext attribute) LayerNorm (class in brainstate.nn) leaky_relu (class in brainstate.nn) LeakyReLU (class in brainstate.nn) LecunNormal (class in brainstate.nn.init) LecunUniform (class in brainstate.nn.init) Linear (class in brainstate.nn) list_hooks() (brainstate.State method) list_registered_behaviors (class in brainstate.environ) locks (brainstate.environ.EnvironmentState attribute) log_abs_det_jacobian() (brainstate.nn.AffineT method) (brainstate.nn.ChainT method) (brainstate.nn.ExpT method) (brainstate.nn.IdentityT method) (brainstate.nn.LogT method) (brainstate.nn.PositiveT method) (brainstate.nn.SigmoidT method) (brainstate.nn.SoftplusT method) (brainstate.nn.Transform method), [1] log_sigmoid (class in brainstate.nn) log_softmax (class in brainstate.nn) loggamma (class in brainstate.random) loggamma() (brainstate.random.RandomState method) logistic (class in brainstate.random) logistic() (brainstate.random.RandomState method) lognormal (class in brainstate.random) lognormal() (brainstate.random.RandomState method) LogNormalReg (class in brainstate.nn) logseries (class in brainstate.random) logseries() (brainstate.random.RandomState method) LogSigmoid (class in brainstate.nn) LogSoftmax (class in brainstate.nn) logsumexp (class in brainstate.nn) LogT (class in brainstate.nn) LogUniformReg (class in brainstate.nn) LongTermState (class in brainstate) LoRA (class in brainstate.nn) loss() (brainstate.nn.BetaReg method) (brainstate.nn.CauchyReg method) (brainstate.nn.ChainedReg method) (brainstate.nn.DirichletReg method) (brainstate.nn.ElasticNetReg method) (brainstate.nn.EntropyReg method) (brainstate.nn.ExponentialReg method) (brainstate.nn.GammaReg method) (brainstate.nn.GaussianReg method) (brainstate.nn.GroupLassoReg method) (brainstate.nn.HorseshoeReg method) (brainstate.nn.HuberReg method) (brainstate.nn.InverseGammaReg method) (brainstate.nn.L1Reg method) (brainstate.nn.L2Reg method) (brainstate.nn.LogNormalReg method) (brainstate.nn.LogUniformReg method) (brainstate.nn.MaxNormReg method) (brainstate.nn.OrthogonalReg method) (brainstate.nn.Regularization method) (brainstate.nn.SpectralNormReg method) (brainstate.nn.SpikeAndSlabReg method) (brainstate.nn.StudentTReg method) (brainstate.nn.TotalVariationReg method) (brainstate.nn.UniformReg method) lower (brainstate.nn.ClipT attribute) (brainstate.nn.NegSoftplusT attribute) (brainstate.nn.SigmoidT attribute) (brainstate.nn.SoftplusT attribute) LPPool1d (class in brainstate.nn) LPPool2d (class in brainstate.nn) LPPool3d (class in brainstate.nn) LSTMCell (class in brainstate.nn) M m2 (brainstate.nn.WelfordMetric attribute) make_jaxpr() (brainstate.transform.StatefulFunction method) (in module brainstate.transform) map() (in module brainstate.transform) map_keys() (brainstate.util.DictManager method) map_values() (brainstate.util.DictManager method) mapping (brainstate.util.PrettyMapping attribute) mapping_fn (brainstate.transform.StatefulMapping attribute) MappingReprMixin (class in brainstate.util) mask (brainstate.nn.MaskedT attribute) MaskedT (class in brainstate.nn) matrix (brainstate.nn.ConfusionMatrix attribute) max_length (brainstate.nn.StateWithDelay attribute) max_norm (brainstate.nn.Embedding attribute) max_time (brainstate.nn.StateWithDelay attribute) MaxNormReg (class in brainstate.nn) MaxPool1d (class in brainstate.nn) MaxPool2d (class in brainstate.nn) MaxPool3d (class in brainstate.nn) MaxUnpool1d (class in brainstate.nn) MaxUnpool2d (class in brainstate.nn) MaxUnpool3d (class in brainstate.nn) maxwell (class in brainstate.random) maxwell() (brainstate.random.RandomState method) maybe_state() (in module brainstate) mean (brainstate.nn.GaussianReg attribute) (brainstate.nn.WelfordMetric attribute) merge() (brainstate.StateTraceStack method) (brainstate.util.FlattedDict static method) (brainstate.util.NestedDict static method) merge_context() (in module brainstate.graph) Metric (class in brainstate.nn) MetricState (class in brainstate.nn) MGUCell (class in brainstate.nn) Mish (class in brainstate.nn) mish (class in brainstate.nn) Missing (class in brainstate.typing) Mixin (class in brainstate.mixin) module brainstate.environ brainstate.graph brainstate.mixin brainstate.nn brainstate.random brainstate.transform brainstate.typing brainstate.util module (brainstate.nn.Vmap attribute) Module (class in brainstate.nn) ModuleMapper (in module brainstate.nn) modules() (brainstate.nn.Module method) MultiMetric (class in brainstate.nn) multinomial (class in brainstate.random) multinomial() (brainstate.random.RandomState method) multivariate_normal (class in brainstate.random) multivariate_normal() (brainstate.random.RandomState method) N name (brainstate.FakeState property) (brainstate.nn.Module attribute) (brainstate.nn.Module property) (brainstate.State attribute) (brainstate.State property) (brainstate.transform.StatefulFunction attribute) (brainstate.TreefyState property) named_children() (brainstate.nn.Module method) named_modules() (brainstate.nn.Module method) named_param_modules() (brainstate.nn.Module method) named_parameters() (brainstate.nn.Module method) negative_binomial (class in brainstate.random) negative_binomial() (brainstate.random.RandomState method) NegativeT (class in brainstate.nn) NegSoftplusT (class in brainstate.nn) nest_mapping (class in brainstate.util) NestedDict (class in brainstate.util) new_arg() (brainstate.StateTraceStack method) Node (class in brainstate.graph) NodeDef (class in brainstate.graph) NodeRef (class in brainstate.graph) nodes (class in brainstate.graph) nodes() (brainstate.nn.Module method) NodeStates (class in brainstate.graph) non_hashable_params (brainstate.mixin.ParamDesc attribute) noncentral_chisquare (class in brainstate.random) noncentral_chisquare() (brainstate.random.RandomState method) noncentral_f (class in brainstate.random) noncentral_f() (brainstate.random.RandomState method) norm_type (brainstate.nn.Embedding attribute) Normal (class in brainstate.nn.init) normal (class in brainstate.random) normal() (brainstate.random.RandomState method) NoSubclassMeta (class in brainstate.mixin) not_implemented (class in brainstate.mixin) not_instance_eval (class in brainstate.util) not_subset() (brainstate.util.DictManager method) num_classes (brainstate.nn.ConfusionMatrix attribute) (brainstate.nn.PrecisionMetric attribute) (brainstate.nn.RecallMetric attribute) num_embeddings (brainstate.nn.Embedding attribute) num_in (brainstate.nn.GRUCell attribute) (brainstate.nn.LSTMCell attribute) (brainstate.nn.MGUCell attribute) (brainstate.nn.ValinaRNNCell attribute) num_out (brainstate.nn.GRUCell attribute) (brainstate.nn.LSTMCell attribute) (brainstate.nn.MGUCell attribute) (brainstate.nn.ValinaRNNCell attribute) num_state (brainstate.HiddenGroupState property) (brainstate.HiddenState property) (brainstate.HiddenTreeState property) numel() (brainstate.State method) O one_hot (class in brainstate.nn) OneOfTypes (in module brainstate.mixin) OneToOne (class in brainstate.nn) optimize_jaxpr() (in module brainstate.transform) OrderedT (class in brainstate.nn) origin_fun (brainstate.transform.StatefulMapping attribute) original_state_values (brainstate.StateTraceStack property) Orthogonal (class in brainstate.nn.init) orthogonal (class in brainstate.random) orthogonal() (brainstate.random.RandomState method) OrthogonalReg (class in brainstate.nn) out_axes (brainstate.transform.StatefulMapping attribute) out_channels (brainstate.nn.Conv1d attribute) (brainstate.nn.Conv2d attribute) (brainstate.nn.Conv3d attribute) (brainstate.nn.ConvTranspose1d attribute) (brainstate.nn.ConvTranspose2d attribute) (brainstate.nn.ConvTranspose3d attribute) (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) out_features (brainstate.nn.LoRA attribute) out_size (brainstate.nn.AllToAll attribute) (brainstate.nn.Conv1d attribute) (brainstate.nn.Conv2d attribute) (brainstate.nn.Conv3d attribute) (brainstate.nn.ConvTranspose1d attribute) (brainstate.nn.ConvTranspose2d attribute) (brainstate.nn.ConvTranspose3d attribute) (brainstate.nn.Dynamics attribute) (brainstate.nn.Embedding attribute) (brainstate.nn.GRUCell attribute) (brainstate.nn.Linear attribute) (brainstate.nn.LoRA attribute) (brainstate.nn.LSTMCell attribute) (brainstate.nn.MGUCell attribute) (brainstate.nn.Module attribute) (brainstate.nn.OneToOne attribute) (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) (brainstate.nn.ScaledWSLinear attribute) (brainstate.nn.SignedWLinear attribute) (brainstate.nn.SparseLinear attribute) (brainstate.nn.ValinaRNNCell attribute) output_delay() (brainstate.nn.Dynamics method) OutputDelayAt (class in brainstate.nn) P padding_idx (brainstate.nn.Embedding attribute) Param (class in brainstate.nn) param (class in brainstate.nn.init) param_modules() (brainstate.nn.Module method) param_precompute() (brainstate.nn.Module method) ParamDesc (class in brainstate.mixin) ParamDescriber (class in brainstate.mixin) parameters() (brainstate.nn.Module method) ParamState (class in brainstate) pareto (class in brainstate.random) pareto() (brainstate.random.RandomState method) PathParts (in module brainstate.typing) permutation (class in brainstate.random) permutation() (brainstate.random.RandomState method) PLATFORM (in module brainstate.environ) pmap2() (in module brainstate.transform) pmap2_new_states() (in module brainstate.transform) poisson (class in brainstate.random) poisson() (brainstate.random.RandomState method) pop (class in brainstate.environ) pop() (brainstate.util.FrozenDict method) pop_by_keys() (brainstate.util.DictManager method) pop_by_values() (brainstate.util.DictManager method) pop_states (class in brainstate.graph) PositiveT (class in brainstate.nn) power (class in brainstate.random) power() (brainstate.random.RandomState method) PowerT (class in brainstate.nn) precision (brainstate.nn.GaussianReg attribute) PRECISION (in module brainstate.environ) precision_metric (brainstate.nn.F1ScoreMetric attribute) PrecisionMetric (class in brainstate.nn) precompute (brainstate.nn.Param attribute) Predicate (in module brainstate.typing) Prefetch (class in brainstate.nn) prefetch() (brainstate.nn.Dynamics method) prefetch_delay() (brainstate.nn.Dynamics method) PrefetchDelay (class in brainstate.nn) PrefetchDelayAt (class in brainstate.nn) PReLU (class in brainstate.nn) prelu (class in brainstate.nn) pretty_repr() (brainstate.util.FrozenDict method) PrettyAttr (class in brainstate.util) PrettyMapping (class in brainstate.util) PrettyRepr (class in brainstate.util) PrettyType (class in brainstate.util) ProgressBar (class in brainstate.transform) PyTree (class in brainstate.typing) PyTreeNode (class in brainstate.util) R raise_error_with_source_info() (brainstate.State method) rand (class in brainstate.random) rand() (brainstate.random.RandomState method) rand_like (class in brainstate.random) rand_like() (brainstate.random.RandomState method) randint (class in brainstate.random) randint() (brainstate.random.RandomState method) randint_like (class in brainstate.random) randint_like() (brainstate.random.RandomState method) randn (class in brainstate.random) randn() (brainstate.random.RandomState method) randn_like (class in brainstate.random) randn_like() (brainstate.random.RandomState method) random (class in brainstate.random) random() (brainstate.random.RandomState method) random_integers (class in brainstate.random) random_integers() (brainstate.random.RandomState method) random_sample (class in brainstate.random) random_sample() (brainstate.random.RandomState method) RandomState (class in brainstate.random) ranf (class in brainstate.random) ranf() (brainstate.random.RandomState method) raw_argnums (brainstate.transform.GradientTransform attribute) rayleigh (class in brainstate.random) rayleigh() (brainstate.random.RandomState method) read_its_value() (brainstate.StateTraceStack method) recall_metric (brainstate.nn.F1ScoreMetric attribute) RecallMetric (class in brainstate.nn) recovery_original_values() (brainstate.StateTraceStack method) ReflectionPad1d (class in brainstate.nn) ReflectionPad2d (class in brainstate.nn) ReflectionPad3d (class in brainstate.nn) RefMap (class in brainstate.graph) reg (brainstate.nn.Param attribute) reg_loss() (brainstate.nn.Module method) (brainstate.nn.Param method) register_default_behavior (class in brainstate.environ) register_delay() (brainstate.nn.Delay method) register_entry() (brainstate.nn.Delay method) register_graph_node_type (class in brainstate.graph) register_hook() (brainstate.State method) Regularization (class in brainstate.nn) regularizations (brainstate.nn.ChainedReg attribute) ReLU (class in brainstate.nn) relu (class in brainstate.nn) ReLU6 (class in brainstate.nn) relu6 (class in brainstate.nn) ReluT (class in brainstate.nn) remat() (in module brainstate.transform) replace() (brainstate.graph.NodeStates method) (brainstate.nn.HiData method) (brainstate.State method) (brainstate.TreefyState method) (brainstate.util.PyTreeNode method) replace_by_pure_dict() (brainstate.util.NestedDict method) ReplicationPad1d (class in brainstate.nn) ReplicationPad2d (class in brainstate.nn) ReplicationPad3d (class in brainstate.nn) reset (class in brainstate.environ) reset() (brainstate.nn.AverageMetric method) (brainstate.nn.ConfusionMatrix method) (brainstate.nn.F1ScoreMetric method) (brainstate.nn.Metric method), [1] (brainstate.nn.MultiMetric method) (brainstate.nn.PrecisionMetric method) (brainstate.nn.RecallMetric method) (brainstate.nn.WelfordMetric method) reset_all_states (class in brainstate.nn) reset_state() (brainstate.nn.Delay method) (brainstate.nn.GRUCell method), [1] (brainstate.nn.LSTMCell method), [1] (brainstate.nn.MGUCell method), [1] (brainstate.nn.Module method) (brainstate.nn.RNNCell method) (brainstate.nn.URLSTMCell method), [1] (brainstate.nn.ValinaRNNCell method), [1] reset_to_prior() (brainstate.nn.Param method) reset_value() (brainstate.nn.BetaReg method) (brainstate.nn.CauchyReg method) (brainstate.nn.ChainedReg method) (brainstate.nn.DirichletReg method) (brainstate.nn.ElasticNetReg method) (brainstate.nn.EntropyReg method) (brainstate.nn.ExponentialReg method) (brainstate.nn.GammaReg method) (brainstate.nn.GaussianReg method) (brainstate.nn.GroupLassoReg method) (brainstate.nn.HorseshoeReg method) (brainstate.nn.HuberReg method) (brainstate.nn.InverseGammaReg method) (brainstate.nn.L1Reg method) (brainstate.nn.L2Reg method) (brainstate.nn.LogNormalReg method) (brainstate.nn.LogUniformReg method) (brainstate.nn.MaxNormReg method) (brainstate.nn.OrthogonalReg method) (brainstate.nn.Regularization method) (brainstate.nn.SpectralNormReg method) (brainstate.nn.SpikeAndSlabReg method) (brainstate.nn.StudentTReg method) (brainstate.nn.TotalVariationReg method) (brainstate.nn.UniformReg method) restore_key (class in brainstate.random) restore_value() (brainstate.State method) retrieve_at_step() (brainstate.nn.Delay method) retrieve_at_time() (brainstate.nn.Delay method) return_only_write (brainstate.transform.StatefulFunction attribute) return_value (brainstate.transform.GradientTransform attribute) RMSNorm (class in brainstate.nn) RNNCell (class in brainstate.nn) RReLU (class in brainstate.nn) rrelu (class in brainstate.nn) S sample (class in brainstate.random) sample() (brainstate.random.RandomState method) sample_init() (brainstate.nn.BetaReg method) (brainstate.nn.CauchyReg method) (brainstate.nn.ChainedReg method) (brainstate.nn.DirichletReg method) (brainstate.nn.ElasticNetReg method) (brainstate.nn.EntropyReg method) (brainstate.nn.ExponentialReg method) (brainstate.nn.GammaReg method) (brainstate.nn.GaussianReg method) (brainstate.nn.GroupLassoReg method) (brainstate.nn.HorseshoeReg method) (brainstate.nn.HuberReg method) (brainstate.nn.InverseGammaReg method) (brainstate.nn.L1Reg method) (brainstate.nn.L2Reg method) (brainstate.nn.LogNormalReg method) (brainstate.nn.LogUniformReg method) (brainstate.nn.MaxNormReg method) (brainstate.nn.OrthogonalReg method) (brainstate.nn.Regularization method) (brainstate.nn.SpectralNormReg method) (brainstate.nn.SpikeAndSlabReg method) (brainstate.nn.StudentTReg method) (brainstate.nn.TotalVariationReg method) (brainstate.nn.UniformReg method) scale_grad_by_freq (brainstate.nn.Embedding attribute) ScaledSigmoidT (class in brainstate.nn) ScaledWSConv1d (class in brainstate.nn) ScaledWSConv2d (class in brainstate.nn) ScaledWSConv3d (class in brainstate.nn) ScaledWSLinear (class in brainstate.nn) scan() (in module brainstate.transform) seed (class in brainstate.random) seed() (brainstate.random.RandomState method) seed_context (class in brainstate.random) SeedOrKey (in module brainstate.typing) self_assign_multi_keys (class in brainstate.random) self_assign_multi_keys() (brainstate.random.RandomState method) SELU (class in brainstate.nn) selu (class in brainstate.nn) Sequential (class in brainstate.nn) set (class in brainstate.environ) set_host_device_count (class in brainstate.environ) set_key (class in brainstate.random) set_platform (class in brainstate.environ) set_precision (class in brainstate.environ) set_value() (brainstate.HiddenGroupState method) (brainstate.HiddenTreeState method) (brainstate.nn.Param method) setdefault() (brainstate.util.DotDict method) settings (brainstate.environ.EnvironmentState attribute) Shape (in module brainstate.typing) ShortTermState (class in brainstate) shuffle (class in brainstate.random) shuffle() (brainstate.random.RandomState method) Sigmoid (class in brainstate.nn) sigmoid (class in brainstate.nn) SigmoidT (class in brainstate.nn) SignedWLinear (class in brainstate.nn) SiLU (class in brainstate.nn) silu (class in brainstate.nn) SimplexT (class in brainstate.nn) Size (in module brainstate.typing) sofo_grad() (in module brainstate.transform) soft_shrink (class in brainstate.nn) soft_sign (class in brainstate.nn) Softmax (class in brainstate.nn) softmax (class in brainstate.nn) Softmax2d (class in brainstate.nn) Softmin (class in brainstate.nn) softmin (class in brainstate.nn) Softplus (class in brainstate.nn) softplus (class in brainstate.nn) SoftplusT (class in brainstate.nn) Softshrink (class in brainstate.nn) Softsign (class in brainstate.nn) SoftsignT (class in brainstate.nn) source_info (brainstate.State property) spar_mat (brainstate.nn.SparseLinear attribute) sparse_plus (class in brainstate.nn) sparse_sigmoid (class in brainstate.nn) SparseLinear (class in brainstate.nn) SpectralNormReg (class in brainstate.nn) SpikeAndSlabReg (class in brainstate.nn) SpikeBitwise (class in brainstate.nn) split() (brainstate.StateDictManager method) (brainstate.util.DictManager method) (brainstate.util.FlattedDict method) (brainstate.util.NestedDict method) split_context() (in module brainstate.graph) split_key (class in brainstate.random) split_key() (brainstate.random.RandomState method) split_keys (class in brainstate.random) split_values() (brainstate.StateDictManager method) squareplus (class in brainstate.nn) stack_level (brainstate.State property) standard_cauchy (class in brainstate.random) standard_cauchy() (brainstate.random.RandomState method) standard_exponential (class in brainstate.random) standard_exponential() (brainstate.random.RandomState method) standard_gamma (class in brainstate.random) standard_gamma() (brainstate.random.RandomState method) standard_normal (class in brainstate.random) standard_normal() (brainstate.random.RandomState method) standard_t (class in brainstate.random) standard_t() (brainstate.random.RandomState method) standardize (class in brainstate.nn) start (brainstate.util.PrettyAttr attribute) (brainstate.util.PrettyType attribute) state (brainstate.nn.StateWithDelay attribute) State (class in brainstate) state_in_axes (brainstate.transform.StatefulMapping attribute) state_out_axes (brainstate.transform.StatefulMapping attribute) state_size (brainstate.nn.HiData property) state_subset() (brainstate.StateTraceStack method) state_trees() (brainstate.nn.Module method) StateDictManager (class in brainstate) StateFinder (class in brainstate.transform) stateful_target (brainstate.transform.GradientTransform attribute) StatefulFunction (class in brainstate.transform) StatefulMapping (class in brainstate.transform) states (class in brainstate.graph) states() (brainstate.nn.Module method) StateTraceStack (class in brainstate) StateWithDelay (class in brainstate.nn) static_argnames (brainstate.transform.StatefulFunction attribute) static_argnums (brainstate.transform.StatefulFunction attribute) StudentTReg (class in brainstate.nn) subset() (brainstate.util.DictManager method) SUPPORTED_PLATFORMS (in module brainstate.environ) SUPPORTED_PRECISIONS (in module brainstate.environ) SupportsDType (class in brainstate.typing) swish (in module brainstate.nn) switch() (in module brainstate.transform) T t (brainstate.nn.Param attribute) (class in brainstate.random) T (in module brainstate.environ) t() (brainstate.random.RandomState method) tag (brainstate.State attribute) Tanh (class in brainstate.nn) tanh (class in brainstate.nn) tanh_shrink (class in brainstate.nn) Tanhshrink (class in brainstate.nn) TanhT (class in brainstate.nn) target (brainstate.transform.GradientTransform attribute) temporary_hook() (brainstate.State method) Threshold (class in brainstate.nn) to_dict() (brainstate.nn.HiData method) (brainstate.util.DictManager method) (brainstate.util.DotDict method) to_dict_values() (brainstate.StateDictManager method) (brainstate.util.FlattedDict method) to_flat() (brainstate.util.NestedDict method) to_nest() (brainstate.util.FlattedDict method) to_predicate (class in brainstate.util) to_pure_dict() (brainstate.util.NestedDict method) to_state() (brainstate.TreefyState method) tolerance (class in brainstate.environ) total (brainstate.nn.AverageMetric attribute) TotalVariationReg (class in brainstate.nn) transform (brainstate.nn.MaskedT attribute) Transform (class in brainstate.nn) transforms (brainstate.nn.ChainT attribute) tree_flatten() (brainstate.util.DictManager method) (brainstate.util.DotDict method) tree_flatten_with_keys() (brainstate.util.FrozenDict method) tree_to_graph (class in brainstate.graph) tree_unflatten() (brainstate.util.DictManager class method) (brainstate.util.DotDict class method) (brainstate.util.FrozenDict class method) treefy_merge (class in brainstate.graph) treefy_split (class in brainstate.graph) treefy_states (class in brainstate.graph) TreefyState (class in brainstate) triangular (class in brainstate.random) triangular() (brainstate.random.RandomState method) true_argnums (brainstate.transform.GradientTransform attribute) true_positives (brainstate.nn.PrecisionMetric attribute) (brainstate.nn.RecallMetric attribute) truncated_normal (class in brainstate.random) truncated_normal() (brainstate.random.RandomState method) TruncatedNormal (class in brainstate.nn.init) type (brainstate.util.PrettyType attribute) type_name (brainstate.util.PrettyMapping attribute) U unflatten (class in brainstate.graph) Unflatten (class in brainstate.nn) unfreeze() (brainstate.util.FrozenDict method) Uniform (class in brainstate.nn.init) uniform (class in brainstate.random) uniform() (brainstate.random.RandomState method) UniformReg (class in brainstate.nn) unique() (brainstate.util.DictManager method) unique_() (brainstate.util.DictManager method) unit (brainstate.nn.NegSoftplusT attribute) (brainstate.nn.SigmoidT attribute) (brainstate.nn.SoftplusT attribute) UnitVectorT (class in brainstate.nn) unregister_default_behavior (class in brainstate.environ) unregister_hook() (brainstate.State method) unvmap() (in module brainstate.transform) update() (brainstate.nn.AccuracyMetric method) (brainstate.nn.AverageMetric method) (brainstate.nn.ConfusionMatrix method) (brainstate.nn.Delay method) (brainstate.nn.Embedding method) (brainstate.nn.EnvironContext method) (brainstate.nn.F1ScoreMetric method) (brainstate.nn.GroupNorm method) (brainstate.nn.GRUCell method) (brainstate.nn.LayerNorm method) (brainstate.nn.LSTMCell method) (brainstate.nn.Metric method), [1] (brainstate.nn.MGUCell method) (brainstate.nn.MultiMetric method) (brainstate.nn.PrecisionMetric method) (brainstate.nn.RecallMetric method) (brainstate.nn.RMSNorm method) (brainstate.nn.RNNCell method) (brainstate.nn.Sequential method) (brainstate.nn.StateWithDelay method) (brainstate.nn.URLSTMCell method), [1] (brainstate.nn.ValinaRNNCell method) (brainstate.nn.Vmap method) (brainstate.nn.WelfordMetric method) (brainstate.util.DotDict method) update_from_ref() (brainstate.State method) update_states (class in brainstate.graph) upper (brainstate.nn.ClipT attribute) URLSTMCell (class in brainstate.nn) V val (brainstate.nn.Param attribute) validate_all_states() (brainstate.transform.StatefulFunction method) validate_states() (brainstate.transform.StatefulFunction method) ValinaRNNCell (class in brainstate.nn) value (brainstate.FakeState property) (brainstate.State attribute) (brainstate.State property) (brainstate.util.PrettyAttr attribute) value() (brainstate.nn.Param method) value_call() (brainstate.State method) value_sep (brainstate.util.PrettyType attribute) values() (brainstate.nn.HiData method) (brainstate.util.FrozenDict method) VarianceScaling (class in brainstate.nn.init) varshape (brainstate.HiddenGroupState property) (brainstate.HiddenState property) (brainstate.HiddenTreeState property) (brainstate.nn.Dynamics property) vector_grad() (in module brainstate.transform) view_pydot() (in module brainstate.transform) Vmap (class in brainstate.nn) vmap() (in module brainstate.transform) vmap2() (in module brainstate.transform) vmap2_new_states() (in module brainstate.transform) vmap_call_all_fns (class in brainstate.nn) vmap_init_all_states (class in brainstate.nn) vmap_new_states() (in module brainstate.transform) vmap_reset_all_states (class in brainstate.nn) vmapped_fn (brainstate.nn.Vmap attribute) vonmises (class in brainstate.random) vonmises() (brainstate.random.RandomState method) W w_mask (brainstate.nn.Linear attribute) (brainstate.nn.ScaledWSLinear attribute) w_sign (brainstate.nn.SignedWLinear attribute) wald (class in brainstate.random) wald() (brainstate.random.RandomState method) weibull (class in brainstate.random) weibull() (brainstate.random.RandomState method) weibull_min (class in brainstate.random) weibull_min() (brainstate.random.RandomState method) weight (brainstate.nn.AllToAll attribute) (brainstate.nn.ChainedReg attribute) (brainstate.nn.Conv1d attribute) (brainstate.nn.Conv2d attribute) (brainstate.nn.Conv3d attribute) (brainstate.nn.ConvTranspose1d attribute) (brainstate.nn.ConvTranspose2d attribute) (brainstate.nn.ConvTranspose3d attribute) (brainstate.nn.Embedding attribute) (brainstate.nn.GaussianReg attribute) (brainstate.nn.L1Reg attribute) (brainstate.nn.L2Reg attribute) (brainstate.nn.Linear attribute) (brainstate.nn.LoRA attribute) (brainstate.nn.OneToOne attribute) (brainstate.nn.PReLU attribute) (brainstate.nn.ScaledWSConv1d attribute) (brainstate.nn.ScaledWSConv2d attribute) (brainstate.nn.ScaledWSConv3d attribute) (brainstate.nn.ScaledWSLinear attribute) (brainstate.nn.SignedWLinear attribute) (brainstate.nn.SparseLinear attribute) weight_standardization (class in brainstate.nn) WelfordMetric (class in brainstate.nn) while_loop() (in module brainstate.transform) width (brainstate.nn.SigmoidT attribute) write_its_value() (brainstate.StateTraceStack method) X XavierNormal (class in brainstate.nn.init) XavierUniform (class in brainstate.nn.init) Z ZeroInit (class in brainstate.nn.init) ZeroPad1d (class in brainstate.nn) ZeroPad2d (class in brainstate.nn) ZeroPad3d (class in brainstate.nn) zipf (class in brainstate.random) zipf() (brainstate.random.RandomState method)
