flax.nnx

Experimental API. See the NNX page for more details.

• graph
  • split()
  • merge()
  • update()
  • pop()
  • state()
  • graphdef()
  • iter_graph()
  • clone()
  • call()
  • GraphDef
  • UpdateContext
    • UpdateContext.merge()
    • UpdateContext.split()
  • update_context()
  • current_update_context()
• module
  • Module
    • Module.sow()
    • Module.iter_modules()
    • 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()
  • 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
  • Normalization
    • BatchNorm
    • LayerNorm
    • RMSNorm
    • GroupNorm
  • 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()
  • cond()
• 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()