Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | R | S | T | U | V | W | Y | Z _ __init__() (brainstate.transform.ProgressBar method) __module__ (brainstate.util.NestedDict attribute) (brainstate.util.PrettyDict attribute) (brainstate.util.PrettyList 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) All (class in brainstate.util) all_checks (in module brainstate.transform) AllToAll (class in brainstate.nn) AlphaDropout (class in brainstate.nn) Any (class in brainstate.util) 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() (in module brainstate.nn) assign_values() (brainstate.StateDictManager method) at() (brainstate.nn.Delay method) (brainstate.nn.PrefetchDelay method) automatic_checks (in module brainstate.transform) 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) batch_axis (brainstate.mixin.Batching attribute) batch_size (brainstate.mixin.Batching attribute) BatchAxisError (class in brainstate) Batching (class in brainstate.mixin) 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) BoundedCache (class in brainstate.util) brainstate.environ module brainstate.graph module brainstate.interop 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) call_all_fns() (in module brainstate.nn) call_order() (in module brainstate.nn) callback (brainstate.Hook attribute) cancel (brainstate.MutableWriteHookContext attribute) cancel_reason (brainstate.MutableWriteHookContext attribute) 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() (in module brainstate.nn) ChainedReg (class in brainstate.nn) ChainT (class in brainstate.nn) check() (in module brainstate.transform) check_error() (in module brainstate.transform) 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) checkify() (in module brainstate.transform) 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() (brainstate.util.BoundedCache method) clear_buffer_memory (class in brainstate.util) clear_cache() (brainstate.nn.Param method) (brainstate.transform.StatefulFunction method) clear_hooks() (brainstate.HookManager method) (brainstate.State method) clear_state_hooks() (in module brainstate) clip() (brainstate.nn.Param method) clip_grad_norm() (in module 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_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) ConversionError (class in brainstate.interop) 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) (in module brainstate.util) copy_from() (brainstate.State method) copy_propagation() (in module brainstate.transform) count (brainstate.nn.AverageMetric attribute) (brainstate.nn.WelfordMetric attribute) count_parameters() (in module 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 (in module brainstate.random) 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) desc() (brainstate.mixin.ParamDesc class method) dftype (class in brainstate.environ) dict_vmap_states (brainstate.nn.Map attribute) 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) disable() (brainstate.HookHandle method) disable_on_error (brainstate.HookConfig attribute) ditype (class in brainstate.environ) div_checks (in module brainstate.transform) 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) E ElasticNetReg (class in brainstate.nn) elem_indent (brainstate.util.PrettyType attribute) ElementWiseBlock (class in brainstate.nn) ELU (class in brainstate.nn) elu() (in module brainstate.nn) Embedding (class in brainstate.nn) embedding_size (brainstate.nn.Embedding attribute) empty_repr (brainstate.util.PrettyType attribute) enable() (brainstate.HookHandle method) enabled (brainstate.Hook 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) error_logger (brainstate.HookConfig attribute) eval_shape() (in module brainstate.transform) EventFixedNumConn (class in brainstate.nn) EventFixedProb (in module brainstate.nn) EventLinear (class in brainstate.nn) Everything (class in brainstate.util) execute() (brainstate.Hook method) execute_init_hooks() (brainstate.HookManager method) execute_read_hooks() (brainstate.HookManager method) execute_restore_hooks() (brainstate.HookManager method) execute_write_after_hooks() (brainstate.HookManager method) execute_write_before_hooks() (brainstate.HookManager method) exp_euler_step() (in module 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) (brainstate.util.PrettyDict 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) flatten_dict() (in module brainstate.util) float_checks (in module brainstate.transform) 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) freeze() (in module brainstate.util) from_dict() (brainstate.nn.HiData class method) (brainstate.util.DotDict class method) from_equinox() (in module brainstate.interop) from_flat() (brainstate.util.NestedDict class method) from_linen() (in module brainstate.interop) from_nest() (brainstate.util.FlattedDict class method) from_nnx() (in module brainstate.interop) 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() (in module brainstate.nn) geometric (class in brainstate.random) geometric() (brainstate.random.RandomState method) get (class in brainstate.environ) get() (brainstate.nn.InterpolationRegistry class method) (brainstate.util.BoundedCache method) (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_hooks() (brainstate.HookManager method) 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_key_data (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_stats() (brainstate.util.BoundedCache method) get_unique_name() (in module brainstate.util) 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) GlobalHookRegistry (class in brainstate) GLU (class in brainstate.nn) glu() (in module 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() (in module brainstate.nn) hard_sigmoid() (in module brainstate.nn) hard_silu() (in module brainstate.nn) hard_swish() (in module brainstate.nn) hard_tanh() (in module brainstate.nn) Hardshrink (class in brainstate.nn) Hardsigmoid (class in brainstate.nn) Hardswish (class in brainstate.nn) Hardtanh (class in brainstate.nn) has() (brainstate.mixin.JointMode method) (brainstate.mixin.Mode method) has_after_update() (brainstate.nn.Dynamics method) has_aux (brainstate.transform.GradientTransform attribute) has_before_update() (brainstate.nn.Dynamics method) has_hooks() (brainstate.HookManager method) (brainstate.State method) has_state_hooks() (in module brainstate) 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) Hook (class in brainstate) hook_id (brainstate.Hook attribute) HookCancellationError (class in brainstate) HookConfig (class in brainstate) HookContext (class in brainstate) HookError (class in brainstate) HookExecutionError (class in brainstate) HookHandle (class in brainstate) HookManager (class in brainstate) HookRegistrationError (class in brainstate) hooks (brainstate.State property) HookWarning (class in brainstate) 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) IdentityT (class in brainstate.nn) ifelse() (in module brainstate.transform) 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) index (brainstate.graph.GraphDef property) index_checks (in module brainstate.transform) init() (brainstate.mixin.ParamDescriber method) (brainstate.nn.Param class method) init_all_states() (brainstate.nn.Map method) (in module brainstate.nn) init_map_size (brainstate.nn.Map attribute) init_maybe_prefetch() (in module 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] InitHookContext (class in brainstate) inline_jit() (in module brainstate.transform) insert() (brainstate.nn.Sequential method) instance() (brainstate.GlobalHookRegistry class method) InteropError (class in brainstate.interop) InterpolationRegistry (class in brainstate.nn) 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_a() (brainstate.mixin.JointMode method) (brainstate.mixin.Mode method) is_dataclass() (in module brainstate.util) is_enabled() (brainstate.HookHandle method) is_instance_eval (class in brainstate.util) is_removed() (brainstate.HookHandle method) 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) JointMode (class in brainstate.mixin) JointTypes (in module brainstate.mixin) jvp() (in module brainstate.transform) K K (class in brainstate.typing) 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.PathContains attribute) (brainstate.util.PrettyAttr attribute) Key (class in brainstate.typing) keys() (brainstate.nn.HiData method) (brainstate.util.BoundedCache 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) LayerMapping (class in brainstate.interop) LayerNorm (class in brainstate.nn) leaky_relu() (in module brainstate.nn) LeakyReLU (class in brainstate.nn) Linear (class in brainstate.nn) list_hooks() (brainstate.State method) list_methods() (brainstate.nn.InterpolationRegistry class method) list_registered_behaviors (class in brainstate.environ) list_state_hooks() (in module brainstate) 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() (in module brainstate.nn) log_softmax() (in module 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() (in module 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 (class in brainstate.nn) map() (brainstate.nn.Map method) (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_errors_per_hook (brainstate.HookConfig 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) maxsize (brainstate.util.BoundedCache property) 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) (brainstate.util.PrettyDict method) merge_context() (in module brainstate.graph) merge_dicts() (in module brainstate.util) metadata (brainstate.HookContext attribute) (brainstate.InitHookContext attribute) Metric (class in brainstate.nn) MetricState (class in brainstate.nn) MGUCell (class in brainstate.nn) Mish (class in brainstate.nn) mish() (in module brainstate.nn) Missing (class in brainstate.typing) MissingDependencyError (class in brainstate.interop) MissingShapeError (class in brainstate.interop) Mixin (class in brainstate.mixin) Mode (class in brainstate.mixin) modes (brainstate.mixin.JointMode attribute) module brainstate.environ brainstate.graph brainstate.interop brainstate.mixin brainstate.nn brainstate.random brainstate.transform brainstate.typing brainstate.util module (brainstate.nn.Map attribute) (brainstate.nn.Vmap attribute) Module (class in 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) MutableWriteHookContext (class in brainstate) N name (brainstate.FakeState property) (brainstate.Hook attribute) (brainstate.HookHandle property) (brainstate.nn.Module attribute) (brainstate.nn.Module property) (brainstate.State attribute) (brainstate.State property) (brainstate.transform.StatefulFunction attribute) (brainstate.TreefyState property) named_call() (in module brainstate.transform) 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) named_scope() (in module brainstate.transform) nan_checks (in module brainstate.transform) 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) NodeEdge (class in brainstate.graph) nodes (class in brainstate.graph) nodes() (brainstate.nn.Module method) NodeSpec (class in brainstate.graph) 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.random) normal() (brainstate.random.RandomState method) NoSubclassMeta (class in brainstate.mixin) Not (class in brainstate.util) not_implemented (class in brainstate.mixin) not_instance_eval (class in brainstate.util) not_receive_update_input() (in module brainstate.nn) not_receive_update_output() (in module brainstate.nn) not_subset() (brainstate.util.DictManager method) Nothing (class in brainstate.util) 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 OfType (class in brainstate.util) old_value (brainstate.MutableWriteHookContext attribute) (brainstate.RestoreHookContext attribute) (brainstate.WriteHookContext attribute) on_error (brainstate.HookConfig attribute) one_hot() (in module brainstate.nn) OneOfTypes (in module brainstate.mixin) OneToOne (class in brainstate.nn) operation (brainstate.HookContext attribute) 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.random) orthogonal() (brainstate.random.RandomState method) OrthogonalReg (class in brainstate.nn) 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_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) PathContains (class in brainstate.util) 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.BoundedCache method) (brainstate.util.FrozenDict method) (in module brainstate.util) 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 (brainstate.util.Not attribute) Predicate (in module brainstate.typing) predicates (brainstate.util.All attribute) (brainstate.util.Any attribute) 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() (in module brainstate.nn) pretty_repr() (brainstate.util.FrozenDict method) (in module brainstate.util) PrettyAttr (class in brainstate.util) PrettyDict (class in brainstate.util) PrettyList (class in brainstate.util) PrettyMapping (class in brainstate.util) PrettyObject (class in brainstate.util) PrettyRepr (class in brainstate.util) PrettyType (class in brainstate.util) priority (brainstate.Hook attribute) (brainstate.HookHandle property) ProgressBar (class in brainstate.transform) PyTree (class in brainstate.typing) PytreeEdge (class in brainstate.graph) 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) ReadHookContext (class in brainstate) recall_metric (brainstate.nn.F1ScoreMetric attribute) RecallMetric (class in brainstate.nn) receive_update_input() (in module brainstate.nn) receive_update_output() (in module 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() (brainstate.nn.InterpolationRegistry class 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.HookManager method) (brainstate.State method) register_layer_mapping() (in module brainstate.interop) register_prim_handler() (in module brainstate.transform) register_state_hook() (in module brainstate) Regularization (class in brainstate.nn) regularizations (brainstate.nn.ChainedReg attribute) release_original_values() (brainstate.StateTraceStack method) ReLU (class in brainstate.nn) relu() (in module brainstate.nn) ReLU6 (class in brainstate.nn) relu6() (in module brainstate.nn) ReluT (class in brainstate.nn) remat() (in module brainstate.transform) remove() (brainstate.HookHandle method) replace() (brainstate.graph.NodeStates method) (brainstate.nn.HiData method) (brainstate.State method) (brainstate.TreefyState method) (brainstate.util.BoundedCache 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.GlobalHookRegistry class method) (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() (in module 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) RestoreHookContext (class in brainstate) 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() (in module 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() (in module brainstate.nn) Sequential (class in brainstate.nn) set (class in brainstate.environ) set() (brainstate.util.BoundedCache method) 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) shard_map() (in module brainstate.transform) ShortTermState (class in brainstate) shuffle (class in brainstate.random) shuffle() (brainstate.random.RandomState method) Sigmoid (class in brainstate.nn) sigmoid() (in module brainstate.nn) SigmoidT (class in brainstate.nn) SignedWLinear (class in brainstate.nn) SiLU (class in brainstate.nn) silu() (in module brainstate.nn) SimplexT (class in brainstate.nn) Size (in module brainstate.typing) soft_shrink() (in module brainstate.nn) soft_sign() (in module brainstate.nn) Softmax (class in brainstate.nn) softmax() (in module brainstate.nn) Softmax2d (class in brainstate.nn) Softmin (class in brainstate.nn) softmin() (in module brainstate.nn) Softplus (class in brainstate.nn) softplus() (in module 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() (in module brainstate.nn) sparse_sigmoid() (in module 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) (brainstate.util.PrettyDict 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_total() (in module brainstate.util) split_values() (brainstate.StateDictManager method) squareplus() (in module 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() (in module brainstate.nn) start (brainstate.util.PrettyAttr attribute) (brainstate.util.PrettyType attribute) state (brainstate.HookContext property) (brainstate.nn.StateWithDelay attribute) State (class in brainstate) state_name (brainstate.HookContext property) state_ref (brainstate.HookContext attribute) state_size (brainstate.nn.HiData property) state_subset() (brainstate.StateTraceStack method) state_trees() (brainstate.nn.Module method) StateDictManager (class in brainstate) StateEdge (class in brainstate.graph) StateFinder (class in brainstate.transform) stateful_target (brainstate.transform.GradientTransform attribute) StatefulFunction (class in brainstate.transform) StatefulMapping (class in brainstate.transform) StateJaxTracer (class in brainstate.util) StateLeafEdge (class in brainstate.graph) states (class in brainstate.graph) states() (brainstate.nn.Module method) StateTraceStack (class in brainstate) StateWithDelay (class in brainstate.nn) Static (class in brainstate.graph) static_argnames (brainstate.transform.StatefulFunction attribute) static_argnums (brainstate.transform.StatefulFunction attribute) StaticEdge (class in brainstate.graph) StudentTReg (class in brainstate.nn) subset() (brainstate.util.DictManager method) (brainstate.util.PrettyDict method) supported_layers() (in module brainstate.interop) 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) (brainstate.util.WithTag attribute) Tanh (class in brainstate.nn) tanh() (in module brainstate.nn) tanh_shrink() (in module 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) timestamp (brainstate.HookContext attribute) to_dict() (brainstate.nn.HiData method) (brainstate.util.DictManager method) (brainstate.util.DotDict method) (brainstate.util.PrettyDict method) to_dict_values() (brainstate.StateDictManager method) (brainstate.util.FlattedDict method) to_equinox() (in module brainstate.interop) to_flat() (brainstate.util.NestedDict method) to_linen() (in module brainstate.interop) to_nest() (brainstate.util.FlattedDict method) to_nnx() (in module brainstate.interop) to_predicate() (in module 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) TraceContextError (class in brainstate) Training (class in brainstate.mixin) transform (brainstate.nn.MaskedT attribute) Transform (class in brainstate.nn) transformed_value (brainstate.MutableWriteHookContext attribute) transforms (brainstate.nn.ChainT attribute) tree_flatten() (brainstate.util.DictManager method) (brainstate.util.DotDict method) (brainstate.util.PrettyList 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) (brainstate.util.PrettyList class method) treefy_merge (class in brainstate.graph) treefy_split (class in brainstate.graph) treefy_state() (brainstate.util.PrettyDict method) 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) type (brainstate.graph.GraphDef property) (brainstate.util.OfType attribute) (brainstate.util.PrettyType attribute) type_name (brainstate.util.PrettyMapping attribute) types (brainstate.mixin.JointMode attribute) U unflatten (class in brainstate.graph) Unflatten (class in brainstate.nn) unflatten_dict() (in module brainstate.util) unfreeze() (brainstate.util.FrozenDict method) (in module brainstate.util) 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) UnmappedLayerError (class in brainstate.interop) unregister_default_behavior (class in brainstate.environ) unregister_hook() (brainstate.HookManager method) (brainstate.State method) unregister_state_hook() (in module brainstate) UnsupportedLayerError (class in brainstate.interop) UnsupportedStructureError (class in brainstate.interop) 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.Map 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) user_checks (in module brainstate.transform) 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.InitHookContext attribute) (brainstate.MutableWriteHookContext attribute) (brainstate.ReadHookContext attribute) (brainstate.RestoreHookContext attribute) (brainstate.State attribute) (brainstate.State property) (brainstate.util.PrettyAttr attribute) (brainstate.WriteHookContext 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) 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) vjp() (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() (in module brainstate.nn) vmap_init_all_states() (in module brainstate.nn) vmap_new_states() (in module brainstate.transform) vmap_reset_all_states() (in module 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) WithTag (class in brainstate.util) write_its_value() (brainstate.StateTraceStack method) WriteHookContext (class in brainstate) Y yield_unique_pretty_repr_items() (in module brainstate.util) Z 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)
