API Reference

• flax.nnx
  • graph
    • split()
    • merge()
    • update()
    • pop()
    • state()
    • variables()
    • graph()
    • graphdef()
    • iter_graph()
    • clone()
    • call()
    • cached_partial()
    • GraphDef
    • UpdateContext
    • update_context()
    • current_update_context()
  • module
    • Module
      • Module.eval()
      • Module.iter_children()
      • Module.iter_modules()
      • Module.perturb()
      • 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
  • training
    • Metrics
      • Metric
      • Average
      • Accuracy
      • Welford
      • MultiMetric
    • Optimizer
      • Optimizer
  • transforms
    • grad()
    • jit()
    • shard_map()
    • 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()
    • variable_name_from_type()
    • variable_type_from_name()
  • 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()
• flax.core.frozen_dict package
  • FrozenDict
    • FrozenDict.copy()
    • FrozenDict.pop()
    • FrozenDict.pretty_repr()
    • FrozenDict.unfreeze()
  • freeze()
  • unfreeze()
  • copy()
  • pop()
  • pretty_repr()
• flax.struct package
  • dataclass()
  • PyTreeNode
• flax.training package
  • Train state
    • TrainState
      • TrainState.apply_gradients()
      • TrainState.create()