flax.nnx# Experimental API. See the NNX page for more details. graph split() merge() update() pop() state() variables() graph() graphdef() iter_graph() clone() call() GraphDef UpdateContext UpdateContext.merge() UpdateContext.split() update_context() current_update_context() module Module Module.eval() Module.iter_children() Module.iter_modules() Module.set_attributes() Module.sow() Module.train() nn Activation functions celu() elu() gelu() glu() hard_sigmoid() hard_silu() hard_swish() hard_tanh() leaky_relu() log_sigmoid() log_softmax() logsumexp() one_hot() relu() selu() sigmoid() silu() soft_sign() softmax() softplus() standardize() swish() tanh() Attention MultiHeadAttention combine_masks() dot_product_attention() make_attention_mask() make_causal_mask() Dtypes canonicalize_dtype() promote_dtype() Initializers constant() delta_orthogonal() glorot_normal() glorot_uniform() he_normal() he_uniform() kaiming_normal() kaiming_uniform() lecun_normal() lecun_uniform() normal() truncated_normal() ones() ones_init() orthogonal() uniform() variance_scaling() xavier_normal() xavier_uniform() zeros() zeros_init() Linear Conv ConvTranspose Embed Linear LinearGeneral Einsum LoRA LoRA LoRALinear Normalization BatchNorm LayerNorm RMSNorm GroupNorm Recurrent LSTMCell OptimizedLSTMCell SimpleCell GRUCell RNN Bidirectional flip_sequences() Stochastic Dropout rnglib Rngs Rngs.__init__() RngStream reseed() spmd get_partition_spec() get_named_sharding() with_partitioning() with_sharding_constraint() state State State.filter() State.merge() State.split() training Metrics Metric Average Accuracy Welford MultiMetric Optimizer Optimizer transforms Jit Remat Scan Vmap grad() jit() remat() scan() value_and_grad() vmap() eval_shape() custom_vjp() cond() switch() while_loop() fori_loop() variables BatchStat Cache Intermediate Param Variable VariableMetadata VariableState with_metadata() helpers Dict Sequential TrainState TrainState.replace() visualization display() filterlib to_predicate() WithTag PathContains OfType Any All Not Everything Nothing bridge ToNNX ToNNX.__call__() ToNNX.lazy_init() ToLinen ToLinen.__call__() to_linen() NNXMeta NNXMeta.__call__() NNXMeta.add_axis() NNXMeta.get_partition_spec() NNXMeta.remove_axis() NNXMeta.replace() NNXMeta.replace_boxed() NNXMeta.to_nnx_variable() NNXMeta.unbox()