Index

_

__call__() (flax.nnx.BatchNorm method) (flax.nnx.bridge.NNXMeta method) (flax.nnx.bridge.ToLinen method) (flax.nnx.bridge.ToNNX method) (flax.nnx.Conv method) (flax.nnx.ConvTranspose method) (flax.nnx.Einsum method) (flax.nnx.Embed method) (flax.nnx.GroupNorm method) (flax.nnx.InstanceNorm method) (flax.nnx.LayerNorm method) (flax.nnx.Linear method) (flax.nnx.LinearGeneral method) (flax.nnx.LoRA method) (flax.nnx.LoRALinear method) (flax.nnx.MultiHeadAttention method) (flax.nnx.nn.recurrent.Bidirectional method) (flax.nnx.nn.recurrent.GRUCell method) (flax.nnx.nn.recurrent.LSTMCell method) (flax.nnx.nn.recurrent.OptimizedLSTMCell method) (flax.nnx.nn.recurrent.RNN method) (flax.nnx.nn.recurrent.SimpleCell method) (flax.nnx.RMSNorm method) (flax.nnx.SpectralNorm method) (flax.nnx.WeightNorm method)

__init__() (flax.nnx.metrics.Average method) (flax.nnx.metrics.Metric method) (flax.nnx.metrics.MultiMetric method) (flax.nnx.metrics.Welford method) (flax.nnx.optimizer.Optimizer method) (flax.nnx.Rngs method)

A

Accuracy (class in flax.nnx.metrics) add_axis() (flax.nnx.bridge.NNXMeta method) All (class in flax.nnx) Any (class in flax.nnx)

append() (flax.nnx.List method) apply_gradients() (flax.training.train_state.TrainState method) attend() (flax.nnx.Embed method) Average (class in flax.nnx.metrics)

B

ball() (flax.nnx.RngStream method) BatchNorm (class in flax.nnx) BatchStat (class in flax.nnx) bernoulli() (flax.nnx.RngStream method)

beta() (flax.nnx.RngStream method) Bidirectional (class in flax.nnx.nn.recurrent) binomial() (flax.nnx.RngStream method) bits() (flax.nnx.RngStream method)

C

Cache (class in flax.nnx) cached_partial() (in module flax.nnx) call() (in module flax.nnx) canonicalize_dtype() (in module flax.nnx.nn.dtypes) Carry (class in flax.nnx) categorical() (flax.nnx.RngStream method) cauchy() (flax.nnx.RngStream method) celu() (in module flax.nnx) check_pytree() (in module flax.nnx) chisquare() (flax.nnx.RngStream method) choice() (flax.nnx.RngStream method) clone() (in module flax.nnx) combine_masks() (in module flax.nnx)

compute() (flax.nnx.metrics.Average method) (flax.nnx.metrics.Metric method) (flax.nnx.metrics.MultiMetric method) (flax.nnx.metrics.Welford method) cond() (in module flax.nnx) Config (class in flax.configurations) constant() (in module flax.nnx.initializers) Conv (class in flax.nnx) ConvTranspose (class in flax.nnx) copy() (flax.core.frozen_dict.FrozenDict method) (in module flax.core.frozen_dict) create() (flax.training.train_state.TrainState class method) current_update_context() (in module flax.nnx) custom_vjp() (in module flax.nnx)

D

Data (in module flax.nnx) data() (in module flax.nnx) dataclass() (in module flax.struct) del_metadata() (flax.nnx.Variable method) delta_orthogonal() (in module flax.nnx.initializers)

Dict (class in flax.nnx) dirichlet() (flax.nnx.RngStream method) display() (in module flax.nnx) dot_product_attention() (in module flax.nnx) double_sided_maxwell() (flax.nnx.RngStream method) Dropout (class in flax.nnx)

E

Einsum (class in flax.nnx) elu() (in module flax.nnx) Embed (class in flax.nnx)

eval() (flax.nnx.Module method) eval_shape() (in module flax.nnx) Everything (class in flax.nnx) exponential() (flax.nnx.RngStream method)

F

f() (flax.nnx.RngStream method) Filter filter_state() (in module flax.nnx) find_duplicates() (in module flax.nnx) FlatState (class in flax.nnx) flatten() (in module flax.nnx) flatten_dict() (in module flax.traverse_util) flax.configurations module flax.nnx module, [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], [12], [13], [14], [15], [16], [17] flax.nnx.bridge module flax.nnx.initializers module flax.nnx.metrics module flax.nnx.nn.dtypes module flax.nnx.nn.recurrent module flax.nnx.optimizer module

flax.struct module flax.traverse_util module flax_always_shard_variable (flax.configurations.Config property) flax_array_ref (flax.configurations.Config property) flax_filter_frames (flax.configurations.Config property) flax_fix_rng_separator (flax.configurations.Config property) flax_hijax_variable (flax.configurations.Config property) flax_max_repr_depth (flax.configurations.Config property) flax_preserve_adopted_names (flax.configurations.Config property) flax_profile (flax.configurations.Config property) flax_pytree_module (flax.configurations.Config property) flax_return_frozendict (flax.configurations.Config property) flax_use_flaxlib (flax.configurations.Config property) flax_use_orbax_checkpointing (flax.configurations.Config property) flip_sequences() (in module flax.nnx.nn.recurrent) Folding in fori_loop() (in module flax.nnx) fork_rngs() (in module flax.nnx) freeze() (in module flax.core.frozen_dict) from_flat_state() (in module flax.nnx) FrozenDict (class in flax.core.frozen_dict)

G

gamma() (flax.nnx.RngStream method) gelu() (in module flax.nnx) generalized_normal() (flax.nnx.RngStream method) geometric() (flax.nnx.RngStream method) get_metadata() (flax.nnx.Variable method) get_named_sharding() (in module flax.nnx) get_partition_spec() (flax.nnx.bridge.NNXMeta method) (in module flax.nnx) glorot_normal() (in module flax.nnx.initializers)

glorot_uniform() (in module flax.nnx.initializers) glu() (in module flax.nnx) grad() (in module flax.nnx) graph() (in module flax.nnx) GraphDef (class in flax.nnx) graphdef() (in module flax.nnx) GroupNorm (class in flax.nnx) GRUCell (class in flax.nnx.nn.recurrent) gumbel() (flax.nnx.RngStream method)

H

hard_sigmoid() (in module flax.nnx) hard_silu() (in module flax.nnx) hard_swish() (in module flax.nnx)

hard_tanh() (in module flax.nnx) has_metadata() (flax.nnx.Variable method) he_normal() (in module flax.nnx.initializers) he_uniform() (in module flax.nnx.initializers)

I

identity() (in module flax.nnx) init_cache() (flax.nnx.MultiHeadAttention method) initialize_carry() (flax.nnx.nn.recurrent.GRUCell method) (flax.nnx.nn.recurrent.LSTMCell method) (flax.nnx.nn.recurrent.OptimizedLSTMCell method) (flax.nnx.nn.recurrent.SimpleCell method) insert() (flax.nnx.List method) InstanceNorm (class in flax.nnx)

int_flag() (in module flax.configurations) Intermediate (class in flax.nnx) is_data() (in module flax.nnx) iter_children() (flax.nnx.Module method) (in module flax.nnx) iter_graph() (in module flax.nnx) iter_modules() (flax.nnx.Module method) (in module flax.nnx)

J

jit() (in module flax.nnx)

jvp() (in module flax.nnx)

K

kaiming_normal() (in module flax.nnx.initializers)

kaiming_uniform() (in module flax.nnx.initializers)

L

laplace() (flax.nnx.RngStream method) LayerNorm (class in flax.nnx) lazy_init() (flax.nnx.bridge.ToNNX method) leaky_relu() (in module flax.nnx) lecun_normal() (in module flax.nnx.initializers) lecun_uniform() (in module flax.nnx.initializers) Linear (class in flax.nnx) LinearGeneral (class in flax.nnx) List (class in flax.nnx)

log_sigmoid() (in module flax.nnx) log_softmax() (in module flax.nnx) loggamma() (flax.nnx.RngStream method) logistic() (flax.nnx.RngStream method) lognormal() (flax.nnx.RngStream method) logsumexp() (in module flax.nnx) LoRA (class in flax.nnx) LoRALinear (class in flax.nnx) LSTMCell (class in flax.nnx.nn.recurrent)

M

make_attention_mask() (in module flax.nnx) make_causal_mask() (in module flax.nnx) map_state() (in module flax.nnx) maxwell() (flax.nnx.RngStream method) Merge merge() (in module flax.nnx) merge_state() (in module flax.nnx) Metric (class in flax.nnx.metrics) Module module (class in flax.nnx) flax.configurations flax.nnx, [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], [12], [13], [14], [15], [16], [17] flax.nnx.bridge flax.nnx.initializers flax.nnx.metrics flax.nnx.nn.dtypes flax.nnx.nn.recurrent flax.nnx.optimizer flax.struct flax.traverse_util

MultiHeadAttention (class in flax.nnx) MultiMetric (class in flax.nnx.metrics) multinomial() (flax.nnx.RngStream method) multivariate_normal() (flax.nnx.RngStream method)

N

nnx_graph_mode (flax.configurations.Config property) nnx_graph_updates (flax.configurations.Config property) NNXMeta (class in flax.nnx.bridge)

normal() (flax.nnx.RngStream method) (in module flax.nnx.initializers) Not (class in flax.nnx) Nothing (class in flax.nnx)

O

Object (class in flax.nnx) OfType (class in flax.nnx) one_hot() (in module flax.nnx) ones() (in module flax.nnx.initializers) ones_init() (in module flax.nnx.initializers)

opt_state (flax.nnx.optimizer.Optimizer attribute) OptimizedLSTMCell (class in flax.nnx.nn.recurrent) Optimizer (class in flax.nnx.optimizer) orthogonal() (flax.nnx.RngStream method) (in module flax.nnx.initializers)

P

Param (class in flax.nnx) Params / parameters pareto() (flax.nnx.RngStream method) path_aware_map() (in module flax.traverse_util) PathContains (class in flax.nnx) permutation() (flax.nnx.RngStream method) perturb() (flax.nnx.Module method) poisson() (flax.nnx.RngStream method) pop() (flax.core.frozen_dict.FrozenDict method) (in module flax.core.frozen_dict) (in module flax.nnx)

pretty_repr() (flax.core.frozen_dict.FrozenDict method) (in module flax.core.frozen_dict) PRNG states promote_dtype() (in module flax.nnx.nn.dtypes) pure() (in module flax.nnx) Pytree (class in flax.nnx) PyTreeNode (class in flax.struct)

R

rademacher() (flax.nnx.RngStream method) randint() (flax.nnx.RngStream method) rayleigh() (flax.nnx.RngStream method) recursive_map() (in module flax.nnx) register_data_type() (in module flax.nnx) register_variable_name() (in module flax.nnx) relu() (in module flax.nnx) remat() (in module flax.nnx) remove_axis() (flax.nnx.bridge.NNXMeta method) replace() (flax.nnx.bridge.NNXMeta method) (flax.nnx.TrainState method)

replace_boxed() (flax.nnx.bridge.NNXMeta method) replace_by_pure_dict() (in module flax.nnx) reseed() (in module flax.nnx) reset() (flax.nnx.metrics.Average method) (flax.nnx.metrics.Metric method) (flax.nnx.metrics.MultiMetric method) (flax.nnx.metrics.Welford method) restore_int_paths() (in module flax.nnx) RMSNorm (class in flax.nnx) Rngs (class in flax.nnx) RngStream (class in flax.nnx) RNN (class in flax.nnx.nn.recurrent)

S

scan() (in module flax.nnx) selu() (in module flax.nnx) Sequential (class in flax.nnx) set_attributes() (flax.nnx.Module method) set_metadata() (flax.nnx.Variable method) (in module flax.nnx) set_view() (flax.nnx.BatchNorm method) (flax.nnx.Dropout method) (flax.nnx.MultiHeadAttention method) shard_map() (in module flax.nnx) sigmoid() (in module flax.nnx) silu() (in module flax.nnx) SimpleCell (class in flax.nnx.nn.recurrent) soft_sign() (in module flax.nnx) softmax() (in module flax.nnx)

softplus() (in module flax.nnx) sow() (flax.nnx.Module method) SpectralNorm (class in flax.nnx) Split and merge split() (in module flax.nnx) split_rngs() (in module flax.nnx) split_state() (in module flax.nnx) standardize() (in module flax.nnx) State (class in flax.nnx) state() (in module flax.nnx) Static (in module flax.nnx) static() (in module flax.nnx) static_int_env() (in module flax.configurations) step (flax.nnx.optimizer.Optimizer attribute) swish() (in module flax.nnx) switch() (in module flax.nnx)

T

t() (flax.nnx.RngStream method) tabulate() (in module flax.nnx) tanh() (in module flax.nnx) temp_flip_flag() (flax.configurations.Config method) to_flat_state() (in module flax.nnx) to_linen() (in module flax.nnx.bridge) to_nnx_variable() (flax.nnx.bridge.NNXMeta method) to_predicate() (in module flax.nnx.filterlib) to_pure_dict() (in module flax.nnx) ToLinen (class in flax.nnx.bridge)

ToNNX (class in flax.nnx.bridge) train() (flax.nnx.Module method) TrainState (class in flax.nnx) (class in flax.training.train_state) Transformation triangular() (flax.nnx.RngStream method) truncated_normal() (flax.nnx.RngStream method) (in module flax.nnx.initializers) tx (flax.nnx.optimizer.Optimizer attribute) type (flax.nnx.Variable property)

U

unbox() (flax.nnx.bridge.NNXMeta method) unflatten() (in module flax.nnx) unflatten_dict() (in module flax.traverse_util) unfreeze() (flax.core.frozen_dict.FrozenDict method) (in module flax.core.frozen_dict) uniform() (flax.nnx.RngStream method) (in module flax.nnx.initializers) update() (flax.configurations.Config method) (flax.nnx.metrics.Accuracy method) (flax.nnx.metrics.Average method) (flax.nnx.metrics.Metric method) (flax.nnx.metrics.MultiMetric method) (flax.nnx.metrics.Welford method) (flax.nnx.optimizer.Optimizer method) (in module flax.nnx)

update_context() (in module flax.nnx) UpdateContext (class in flax.nnx)

V

value_and_grad() (in module flax.nnx) Variable (class in flax.nnx) variable_name_from_type() (in module flax.nnx) variable_type_from_name() (in module flax.nnx) VariableMetadata (class in flax.nnx)

variables() (in module flax.nnx) variance_scaling() (in module flax.nnx.initializers) view() (in module flax.nnx) view_info() (in module flax.nnx) vjp() (in module flax.nnx) vmap() (in module flax.nnx)

W

wald() (flax.nnx.RngStream method) weibull_min() (flax.nnx.RngStream method) WeightNorm (class in flax.nnx) Welford (class in flax.nnx.metrics)

while_loop() (in module flax.nnx) with_attributes() (in module flax.nnx) with_metadata() (in module flax.nnx) with_partitioning() (in module flax.nnx) WithTag (class in flax.nnx)

X

xavier_normal() (in module flax.nnx.initializers)

xavier_uniform() (in module flax.nnx.initializers)

Z

zeros() (in module flax.nnx.initializers)

zeros_init() (in module flax.nnx.initializers)