API Reference

• flax.config package
  • Config
    • Config.flax_filter_frames
    • Config.flax_fix_rng_separator
    • Config.flax_preserve_adopted_names
    • Config.flax_profile
    • Config.flax_return_frozendict
    • Config.flax_use_orbax_checkpointing
    • Config.update()
• flax.core.frozen_dict package
  • FrozenDict
    • FrozenDict.copy()
    • FrozenDict.pop()
    • FrozenDict.pretty_repr()
    • FrozenDict.unfreeze()
  • freeze()
  • unfreeze()
  • copy()
  • pop()
  • pretty_repr()
• flax.cursor package
  • cursor()
  • Cursor
    • Cursor.apply_update()
    • Cursor.build()
    • Cursor.find()
    • Cursor.find_all()
    • Cursor.set()
• flax.errors package
  • AlreadyExistsError
  • ApplyModuleInvalidMethodError
  • ApplyScopeInvalidVariablesStructureError
  • ApplyScopeInvalidVariablesTypeError
  • AssignSubModuleError
  • CallCompactUnboundModuleError
  • CallSetupUnboundModuleError
  • CallUnbindOnUnboundModuleError
  • CursorFindError
  • DescriptorAttributeError
  • IncorrectPostInitOverrideError
  • InvalidCheckpointError
  • InvalidFilterError
  • InvalidInstanceModuleError
  • InvalidRngError
  • InvalidScopeError
  • JaxTransformError
  • LazyInitError
  • MPACheckpointingRequiredError
  • MPARestoreDataCorruptedError
  • MPARestoreTargetRequiredError
  • ModifyScopeVariableError
  • MultipleMethodsCompactError
  • NameInUseError
  • PartitioningUnspecifiedError
  • ReservedModuleAttributeError
  • ScopeCollectionNotFound
  • ScopeParamNotFoundError
  • ScopeParamShapeError
  • ScopeVariableNotFoundError
  • SetAttributeFrozenModuleError
  • SetAttributeInModuleSetupError
  • TransformTargetError
  • TransformedMethodReturnValueError
  • TraverseTreeError
• flax.experimental.nnx
  • module
    • Module
      • Module.sow()
      • Module.iter_modules()
      • Module.eval()
      • Module.init
      • Module.is_initializing()
      • Module.iter_modules()
      • Module.set_attributes()
      • 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
      • Embed
      • Linear
      • LinearGeneral
      • Einsum
    • Normalization
      • BatchNorm
      • LayerNorm
      • RMSNorm
    • Stochastic
      • Dropout
  • rnglib
    • Rngs
    • RngStream
  • spmd
    • get_partition_spec()
    • get_named_sharding()
    • with_partitioning()
    • with_sharding_constraint()
  • training
    • Metrics
      • Metric
      • Average
      • Accuracy
      • MultiMetric
    • Optimizer
      • Optimizer
  • transforms
    • Remat
    • Scan
    • Vmap
    • grad()
    • jit()
    • remat()
    • scan()
    • value_and_grad()
    • vmap()
  • variables
    • BatchStat
    • Cache
    • Empty
    • Intermediate
    • Param
    • Variable
    • VariableMetadata
    • with_metadata()
  • helpers
    • Dict
    • List
    • Sequential
    • TrainState
      • TrainState.replace()
  • visualization
    • display()
• flax.jax_utils package
  • partial_eval_by_shape()
  • Multi device utilities
    • replicate()
    • unreplicate()
    • prefetch_to_device()
    • pmean()
    • pad_shard_unpad()
• flax.linen
  • Module
    • Module
      • Module.__setattr__()
      • Module.apply()
      • Module.bind()
      • Module.copy()
      • Module.get_variable()
      • Module.has_rng()
      • Module.has_variable()
      • Module.init()
      • Module.init_with_output()
      • Module.is_initializing()
      • Module.is_mutable_collection()
      • Module.lazy_init()
      • Module.make_rng()
      • Module.module_paths()
      • Module.param()
      • Module.perturb()
      • Module.put_variable()
      • Module.setup()
      • Module.sow()
      • Module.tabulate()
      • Module.unbind()
      • Module.variable()
      • Module.variables
    • apply()
    • init()
    • init_with_output()
    • intercept_methods()
  • Init/Apply
    • apply()
    • init()
    • init_with_output()
  • Layers
    • Linear Modules
      • Dense
      • DenseGeneral
      • Conv
      • ConvTranspose
      • ConvLocal
      • Einsum
      • Embed
    • Pooling
      • max_pool()
      • avg_pool()
      • pool()
    • Normalization
      • BatchNorm
      • LayerNorm
      • GroupNorm
      • RMSNorm
      • InstanceNorm
      • SpectralNorm
      • WeightNorm
    • Combinators
      • Sequential
    • Stochastic
      • Dropout
    • Attention
      • MultiHeadDotProductAttention
      • MultiHeadAttention
      • SelfAttention
      • dot_product_attention_weights()
      • dot_product_attention()
      • make_attention_mask()
      • make_causal_mask()
    • Recurrent
      • RNNCellBase
      • LSTMCell
      • OptimizedLSTMCell
      • ConvLSTMCell
      • SimpleCell
      • GRUCell
      • MGUCell
      • RNN
      • Bidirectional
    • BatchApply
      • BatchApply
  • Activation functions
    • PReLU
      • PReLU.param_dtype
      • PReLU.negative_slope_init
      • PReLU.__call__()
    • 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()
  • 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()
  • Transformations
    • vmap()
    • scan()
    • jit()
    • remat()
    • remat_scan()
    • map_variables()
    • jvp()
    • vjp()
    • custom_vjp()
    • while_loop()
    • cond()
    • switch()
  • Inspection
    • tabulate()
  • Variable dictionary
    • Variable
  • SPMD
    • Partitioned()
    • with_partitioning()
    • get_partition_spec()
    • get_sharding()
    • LogicallyPartitioned()
    • logical_axis_rules()
    • set_logical_axis_rules()
    • get_logical_axis_rules()
    • logical_to_mesh_axes()
    • logical_to_mesh()
    • logical_to_mesh_sharding()
    • with_logical_constraint()
    • with_logical_partitioning()
  • Decorators
    • compact()
    • nowrap()
  • Profiling
    • enable_named_call()
    • disable_named_call()
    • override_named_call()
• flax.serialization package
  • State dicts
    • from_state_dict()
    • to_state_dict()
    • register_serialization_state()
  • Serialization with MessagePack
    • msgpack_serialize()
    • msgpack_restore()
    • to_bytes()
    • from_bytes()
• flax.struct package
  • dataclass()
  • PyTreeNode
• flax.traceback_util package
  • Traceback filtering utils
    • hide_flax_in_tracebacks()
    • show_flax_in_tracebacks()
• flax.training package
  • Checkpoints
    • save_checkpoint()
    • save_checkpoint_multiprocess()
    • latest_checkpoint()
    • restore_checkpoint()
    • convert_pre_linen()
  • Learning rate schedules
    • create_constant_learning_rate_schedule()
    • create_stepped_learning_rate_schedule()
    • create_cosine_learning_rate_schedule()
  • Train state
    • TrainState
      • TrainState.apply_gradients()
      • TrainState.create()
  • Early Stopping
    • EarlyStopping
      • EarlyStopping.min_delta
      • EarlyStopping.patience
      • EarlyStopping.best_metric
      • EarlyStopping.patience_count
      • EarlyStopping.should_stop
      • EarlyStopping.has_improved
      • EarlyStopping.update()
  • Common Utilities
    • shard()
    • shard_prng_key()
    • stack_forest()
    • get_metrics()
    • onehot()
• flax.traverse_util package
  • Traversal objects
    • Traversal
      • Traversal.compose()
      • Traversal.each()
      • Traversal.filter()
      • Traversal.iterate()
      • Traversal.merge()
      • Traversal.set()
      • Traversal.tree()
      • Traversal.update()
    • TraverseId
      • TraverseId.iterate()
      • TraverseId.update()
    • TraverseMerge
      • TraverseMerge.iterate()
      • TraverseMerge.update()
    • TraverseCompose
      • TraverseCompose.iterate()
      • TraverseCompose.update()
    • TraverseFilter
      • TraverseFilter.iterate()
      • TraverseFilter.update()
    • TraverseAttr
      • TraverseAttr.iterate()
      • TraverseAttr.update()
    • TraverseItem
      • TraverseItem.iterate()
      • TraverseItem.update()
    • TraverseEach
      • TraverseEach.iterate()
      • TraverseEach.update()
    • TraverseTree
      • TraverseTree.iterate()
      • TraverseTree.update()
  • Dict utils
    • flatten_dict()
    • unflatten_dict()
    • path_aware_map()
  • Model parameter traversal
    • ModelParamTraversal
      • ModelParamTraversal.__init__()