API Reference

• flax.nnx
  • graph
    • split()
    • merge()
    • update()
    • pop()
    • state()
    • variables()
    • graph()
    • graphdef()
    • iter_graph()
    • recursive_map()
    • clone()
    • call()
    • set_metadata()
    • cached_partial()
    • GraphDef
    • UpdateContext
    • update_context()
    • current_update_context()
    • find_duplicates()
    • pure()
    • flatten()
    • unflatten()
  • object
    • Pytree
    • Object
    • data()
    • Data
    • static()
    • Static
    • is_data()
    • register_data_type()
    • check_pytree()
  • module
    • iter_children()
    • iter_modules()
    • Module
      • Module.eval()
      • Module.iter_children()
      • Module.iter_modules()
      • Module.perturb()
      • Module.set_attributes()
      • Module.sow()
      • Module.train()
    • view()
    • view_info()
    • with_attributes()
  • 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()
      • identity()
      • 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
      • InstanceNorm
      • SpectralNorm
      • WeightNorm
    • Recurrent
      • LSTMCell
      • OptimizedLSTMCell
      • SimpleCell
      • GRUCell
      • RNN
      • Bidirectional
      • flip_sequences()
    • Stochastic
      • Dropout
  • rnglib
    • Rngs
      • Rngs.__init__()
    • RngStream
      • RngStream.ball()
      • RngStream.bernoulli()
      • RngStream.beta()
      • RngStream.binomial()
      • RngStream.bits()
      • RngStream.categorical()
      • RngStream.cauchy()
      • RngStream.chisquare()
      • RngStream.choice()
      • RngStream.dirichlet()
      • RngStream.double_sided_maxwell()
      • RngStream.exponential()
      • RngStream.f()
      • RngStream.gamma()
      • RngStream.generalized_normal()
      • RngStream.geometric()
      • RngStream.gumbel()
      • RngStream.laplace()
      • RngStream.loggamma()
      • RngStream.logistic()
      • RngStream.lognormal()
      • RngStream.maxwell()
      • RngStream.multinomial()
      • RngStream.multivariate_normal()
      • RngStream.normal()
      • RngStream.orthogonal()
      • RngStream.pareto()
      • RngStream.permutation()
      • RngStream.poisson()
      • RngStream.rademacher()
      • RngStream.randint()
      • RngStream.rayleigh()
      • RngStream.t()
      • RngStream.triangular()
      • RngStream.truncated_normal()
      • RngStream.uniform()
      • RngStream.wald()
      • RngStream.weibull_min()
    • split_rngs()
    • fork_rngs()
    • reseed()
  • spmd
    • get_partition_spec()
    • get_named_sharding()
    • with_partitioning()
  • state
    • State
    • FlatState
    • filter_state()
    • from_flat_state()
    • map_state()
    • merge_state()
    • replace_by_pure_dict()
    • restore_int_paths()
    • to_flat_state()
    • to_pure_dict()
    • split_state()
  • training
    • Metrics
      • Metric
      • Average
      • Accuracy
      • Welford
      • MultiMetric
    • Optimizer
      • Optimizer
  • transforms
    • grad()
    • jit()
    • shard_map()
    • remat()
    • scan()
    • Carry
    • value_and_grad()
    • vmap()
    • eval_shape()
    • custom_vjp()
    • vjp()
    • jvp()
    • cond()
    • switch()
    • while_loop()
    • fori_loop()
  • variables
    • BatchStat
    • Cache
    • Intermediate
    • Param
    • Variable
      • Variable.del_metadata()
      • Variable.get_metadata()
      • Variable.has_metadata()
      • Variable.set_metadata()
      • Variable.type
    • VariableMetadata
    • with_metadata()
    • variable_name_from_type()
    • variable_type_from_name()
    • register_variable_name()
  • helpers
    • Sequential
    • List
      • List.append()
      • List.insert()
    • Dict
    • 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()