API Reference#
- flax.linen package
- Module
- Init/Apply
- Variable dictionary
- Compact methods
- No wrap methods
- Profiling
- Inspection
- Transformations
- SPMD
- flax.linen.Partitioned
- flax.linen.with_partitioning
- flax.linen.get_partition_spec
- 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.with_logical_constraint
- flax.linen.with_logical_partitioning
- Linear Modules
- Normalization
- Pooling
- 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
- Attention primitives
- Stochastic
- RNN primitives
- flax.serialization package
- flax.core.frozen_dict package
- flax.struct package
- flax.jax_utils package
- flax.traceback_util package
- flax.traverse_util package
- flax.training package
- flax.config package
- 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