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.error package
  • AlreadyExistsError
  • ApplyModuleInvalidMethodError
  • ApplyScopeInvalidVariablesStructureError
  • ApplyScopeInvalidVariablesTypeError
  • AssignSubModuleError
  • CallCompactUnboundModuleError
  • CallSetupUnboundModuleError
  • CallUnbindOnUnboundModuleError
  • DescriptorAttributeError
  • IncorrectPostInitOverrideError
  • InvalidCheckpointError
  • InvalidFilterError
  • InvalidInstanceModuleError
  • InvalidRngError
  • InvalidScopeError
  • JaxTransformError
  • LazyInitError
  • MPACheckpointingRequiredError
  • MPARestoreDataCorruptedError
  • MPARestoreTargetRequiredError
  • MPARestoreTypeNotMatchError
  • 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 package
  • Module
    • Module
  • Init/Apply
    • apply()
    • init()
    • init_with_output()
  • Variable dictionary
    • Variable
  • Compact methods
    • compact()
  • No wrap methods
    • nowrap()
  • Profiling
    • flax.linen.enable_named_call
    • flax.linen.disable_named_call
    • flax.linen.override_named_call
  • Inspection
    • flax.linen.tabulate
  • Transformations
    • flax.linen.vmap
    • flax.linen.scan
    • flax.linen.jit
    • flax.linen.remat
    • flax.linen.remat_scan
    • flax.linen.map_variables
    • flax.linen.jvp
    • flax.linen.vjp
    • flax.linen.custom_vjp
    • flax.linen.while_loop
    • flax.linen.cond
    • flax.linen.switch
  • SPMD
    • flax.linen.Partitioned
    • flax.linen.with_partitioning
    • flax.linen.get_partition_spec
    • flax.linen.get_sharding
    • flax.linen.LogicallyPartitioned
    • flax.linen.logical_axis_rules
    • flax.linen.set_logical_axis_rules
    • flax.linen.get_logical_axis_rules
    • flax.linen.logical_to_mesh_axes
    • flax.linen.logical_to_mesh
    • flax.linen.logical_to_mesh_sharding
    • flax.linen.with_logical_constraint
    • flax.linen.with_logical_partitioning
  • Linear Modules
    • flax.linen.Dense
    • flax.linen.DenseGeneral
    • flax.linen.Conv
    • flax.linen.ConvTranspose
    • flax.linen.ConvLocal
    • flax.linen.Embed
  • Normalization
    • flax.linen.BatchNorm
    • flax.linen.LayerNorm
    • flax.linen.GroupNorm
  • Pooling
    • flax.linen.max_pool
    • flax.linen.avg_pool
    • flax.linen.pool
  • Activation functions
    • flax.linen.activation.PReLU
    • flax.linen.activation.celu
    • flax.linen.activation.elu
    • flax.linen.activation.gelu
    • flax.linen.activation.glu
    • flax.linen.activation.hard_sigmoid
    • flax.linen.activation.hard_silu
    • flax.linen.activation.hard_swish
    • flax.linen.activation.hard_tanh
    • flax.linen.activation.leaky_relu
    • flax.linen.activation.log_sigmoid
    • flax.linen.activation.log_softmax
    • flax.linen.activation.logsumexp
    • flax.linen.activation.one_hot
    • flax.linen.activation.relu
    • flax.linen.activation.relu6
    • flax.linen.activation.selu
    • flax.linen.activation.sigmoid
    • flax.linen.activation.silu
    • flax.linen.activation.soft_sign
    • flax.linen.activation.softmax
    • flax.linen.activation.softplus
    • flax.linen.activation.standardize
    • flax.linen.activation.swish
    • flax.linen.activation.tanh
  • Initializers
    • flax.linen.initializers.constant
    • flax.linen.initializers.delta_orthogonal
    • flax.linen.initializers.glorot_normal
    • flax.linen.initializers.glorot_uniform
    • flax.linen.initializers.he_normal
    • flax.linen.initializers.he_uniform
    • flax.linen.initializers.kaiming_normal
    • flax.linen.initializers.kaiming_uniform
    • flax.linen.initializers.lecun_normal
    • flax.linen.initializers.lecun_uniform
    • flax.linen.initializers.normal
    • flax.linen.initializers.ones
    • flax.linen.initializers.ones_init
    • flax.linen.initializers.orthogonal
    • flax.linen.initializers.uniform
    • flax.linen.initializers.variance_scaling
    • flax.linen.initializers.xavier_normal
    • flax.linen.initializers.xavier_uniform
    • flax.linen.initializers.zeros
    • flax.linen.initializers.zeros_init
  • Combinators
    • flax.linen.Sequential
  • Attention primitives
    • flax.linen.dot_product_attention_weights
    • flax.linen.dot_product_attention
    • flax.linen.make_attention_mask
    • flax.linen.make_causal_mask
    • flax.linen.SelfAttention
    • flax.linen.MultiHeadDotProductAttention
  • Stochastic
    • flax.linen.Dropout
  • RNN primitives
    • flax.linen.LSTMCell
    • flax.linen.OptimizedLSTMCell
    • flax.linen.GRUCell
    • flax.linen.RNNCellBase
    • flax.linen.RNN
    • flax.linen.Bidirectional
• 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
  • Early Stopping
    • EarlyStopping
  • Common Utilities
    • shard()
    • shard_prng_key()
    • stack_forest()
    • get_metrics()
    • onehot()
• flax.traverse_util package
  • Traversal objects
    • Traversal
    • TraverseId
    • TraverseMerge
    • TraverseCompose
    • TraverseFilter
    • TraverseAttr
    • TraverseItem
    • TraverseEach
    • TraverseTree
  • Dict utils
    • flatten_dict()
    • unflatten_dict()
    • path_aware_map()
  • Model parameter traversal
    • ModelParamTraversal