API Reference

• flax.config package
  • define_bool_state()
  • static_bool_env()
• 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.set()
• flax.error package
  • AlreadyExistsError
  • ApplyModuleInvalidMethodError
  • ApplyScopeInvalidVariablesStructureError
  • ApplyScopeInvalidVariablesTypeError
  • AssignSubModuleError
  • CallCompactUnboundModuleError
  • CallSetupUnboundModuleError
  • CallUnbindOnUnboundModuleError
  • 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
• 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.init()
      • Module.init_with_output()
      • Module.is_initializing()
      • Module.make_rng()
      • Module.param()
      • Module.perturb()
      • Module.setup()
      • Module.sow()
      • Module.tabulate()
      • Module.unbind()
      • Module.variable()
      • Module.variables
  • Init/Apply
    • apply()
    • init()
    • init_with_output()
    • flax.linen.apply
      • apply()
    • flax.linen.init
      • init()
    • flax.linen.init_with_output
      • init_with_output()
  • Layers
    • Linear Modules
      • Dense
      • DenseGeneral
      • Conv
      • ConvTranspose
      • ConvLocal
      • Embed
    • Pooling
      • max_pool()
      • avg_pool()
      • pool()
    • Normalization
      • BatchNorm
      • LayerNorm
      • GroupNorm
    • Combinators
      • Sequential
    • Stochastic
      • Dropout
    • Attention
      • SelfAttention
      • MultiHeadDotProductAttention
      • dot_product_attention_weights()
      • dot_product_attention()
      • make_attention_mask()
      • make_causal_mask()
    • Recurrent
      • RNNCellBase
      • LSTMCell
      • OptimizedLSTMCell
      • GRUCell
      • RNN
      • Bidirectional
      • flax.linen.Dense
      • flax.linen.DenseGeneral
      • flax.linen.Conv
      • flax.linen.ConvTranspose
      • flax.linen.ConvLocal
      • flax.linen.Embed
      • flax.linen.BatchNorm
      • flax.linen.LayerNorm
      • flax.linen.GroupNorm
      • flax.linen.RMSNorm
      • flax.linen.Sequential
      • flax.linen.Dropout
      • flax.linen.SelfAttention
      • flax.linen.MultiHeadDotProductAttention
      • flax.linen.RNNCellBase
      • flax.linen.LSTMCell
      • flax.linen.OptimizedLSTMCell
      • flax.linen.GRUCell
      • flax.linen.RNN
      • flax.linen.Bidirectional
      • flax.linen.max_pool
      • flax.linen.avg_pool
      • flax.linen.pool
      • flax.linen.dot_product_attention_weights
      • flax.linen.dot_product_attention
      • flax.linen.make_attention_mask
      • flax.linen.make_causal_mask
  • 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()
    • flax.linen.activation.PReLU
      • PReLU
    • flax.linen.activation.celu
      • celu()
    • flax.linen.activation.elu
      • elu()
    • flax.linen.activation.gelu
      • gelu()
    • flax.linen.activation.glu
      • glu()
    • flax.linen.activation.hard_sigmoid
      • hard_sigmoid()
    • flax.linen.activation.hard_silu
      • hard_silu()
    • flax.linen.activation.hard_swish
      • hard_swish()
    • flax.linen.activation.hard_tanh
      • hard_tanh()
    • flax.linen.activation.leaky_relu
      • leaky_relu()
    • flax.linen.activation.log_sigmoid
      • log_sigmoid()
    • flax.linen.activation.log_softmax
      • log_softmax()
    • flax.linen.activation.logsumexp
      • logsumexp()
    • flax.linen.activation.one_hot
      • one_hot()
    • flax.linen.activation.relu
      • relu
    • flax.linen.activation.relu6
      • relu6
    • flax.linen.activation.selu
      • selu()
    • flax.linen.activation.sigmoid
      • sigmoid()
    • flax.linen.activation.silu
      • silu()
    • flax.linen.activation.soft_sign
      • soft_sign()
    • flax.linen.activation.softmax
      • softmax()
    • flax.linen.activation.softplus
      • softplus()
    • flax.linen.activation.standardize
      • standardize()
    • flax.linen.activation.swish
      • swish()
    • flax.linen.activation.tanh
      • tanh()
  • Initializers
    • constant()
    • delta_orthogonal()
    • glorot_normal()
    • glorot_uniform()
    • he_normal()
    • he_uniform()
    • kaiming_normal()
    • kaiming_uniform()
    • lecun_normal()
    • lecun_uniform()
    • normal()
    • ones()
    • ones_init()
    • orthogonal()
    • uniform()
    • variance_scaling()
    • xavier_normal()
    • xavier_uniform()
    • zeros()
    • zeros_init()
    • flax.linen.initializers.constant
      • constant()
    • flax.linen.initializers.delta_orthogonal
      • delta_orthogonal()
    • flax.linen.initializers.glorot_normal
      • glorot_normal()
    • flax.linen.initializers.glorot_uniform
      • glorot_uniform()
    • flax.linen.initializers.he_normal
      • he_normal()
    • flax.linen.initializers.he_uniform
      • he_uniform()
    • flax.linen.initializers.kaiming_normal
      • kaiming_normal()
    • flax.linen.initializers.kaiming_uniform
      • kaiming_uniform()
    • flax.linen.initializers.lecun_normal
      • lecun_normal()
    • flax.linen.initializers.lecun_uniform
      • lecun_uniform()
    • flax.linen.initializers.normal
      • normal()
    • flax.linen.initializers.ones
      • ones()
    • flax.linen.initializers.ones_init
      • ones_init()
    • flax.linen.initializers.orthogonal
      • orthogonal()
    • flax.linen.initializers.uniform
      • uniform()
    • flax.linen.initializers.variance_scaling
      • variance_scaling()
    • flax.linen.initializers.xavier_normal
      • xavier_normal()
    • flax.linen.initializers.xavier_uniform
      • xavier_uniform()
    • flax.linen.initializers.zeros
      • zeros()
    • flax.linen.initializers.zeros_init
      • zeros_init()
  • Transformations
    • vmap()
    • scan()
    • jit()
    • remat()
    • remat_scan()
    • map_variables()
    • jvp()
    • vjp()
    • custom_vjp()
    • while_loop()
    • cond()
    • switch()
    • flax.linen.vmap
      • vmap()
    • flax.linen.scan
      • scan()
    • flax.linen.jit
      • jit()
    • flax.linen.remat
      • remat()
    • flax.linen.remat_scan
      • remat_scan()
    • flax.linen.map_variables
      • map_variables()
    • flax.linen.jvp
      • jvp()
    • flax.linen.vjp
      • vjp()
    • flax.linen.custom_vjp
      • custom_vjp()
    • flax.linen.while_loop
      • while_loop()
    • flax.linen.cond
      • cond()
    • flax.linen.switch
      • switch()
  • Inspection
    • tabulate()
    • flax.linen.tabulate
      • tabulate()
  • Variable dictionary
    • Variable
    • flax.linen.Variable
      • 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()
    • flax.linen.Partitioned
      • Partitioned
    • flax.linen.with_partitioning
      • with_partitioning()
    • flax.linen.get_partition_spec
      • get_partition_spec()
    • flax.linen.get_sharding
      • get_sharding()
    • flax.linen.LogicallyPartitioned
      • LogicallyPartitioned
    • flax.linen.logical_axis_rules
      • logical_axis_rules()
    • flax.linen.set_logical_axis_rules
      • set_logical_axis_rules()
    • flax.linen.get_logical_axis_rules
      • get_logical_axis_rules()
    • flax.linen.logical_to_mesh_axes
      • logical_to_mesh_axes()
    • flax.linen.logical_to_mesh
      • logical_to_mesh()
    • flax.linen.logical_to_mesh_sharding
      • logical_to_mesh_sharding()
    • flax.linen.with_logical_constraint
      • with_logical_constraint()
    • flax.linen.with_logical_partitioning
      • with_logical_partitioning()
  • Decorators
    • compact()
    • nowrap()
    • flax.linen.compact
      • compact()
    • flax.linen.nowrap
      • nowrap()
  • Profiling
    • enable_named_call()
    • disable_named_call()
    • override_named_call()
    • flax.linen.enable_named_call
      • enable_named_call()
    • flax.linen.disable_named_call
      • disable_named_call()
    • flax.linen.override_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.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__()