flax.linen.LayerNorm#
- class flax.linen.LayerNorm(epsilon=1e-06, dtype=None, param_dtype=<class 'jax.numpy.float32'>, use_bias=True, use_scale=True, bias_init=<function zeros>, scale_init=<function ones>, reduction_axes=-1, feature_axes=-1, axis_name=None, axis_index_groups=None, use_fast_variance=True, parent=<flax.linen.module._Sentinel object>, name=None)[source]#
Layer normalization (https://arxiv.org/abs/1607.06450).
LayerNorm normalizes the activations of the layer for each given example in a batch independently, rather than across a batch like Batch Normalization. i.e. applies a transformation that maintains the mean activation within each example close to 0 and the activation standard deviation close to 1.
- epsilon#
A small float added to variance to avoid dividing by zero.
- Type
float
- dtype#
the dtype of the result (default: infer from input and params).
- Type
Optional[Any]
- param_dtype#
the dtype passed to parameter initializers (default: float32).
- Type
Any
- use_bias#
If True, bias (beta) is added.
- Type
bool
- use_scale#
If True, multiply by scale (gamma). When the next layer is linear (also e.g. nn.relu), this can be disabled since the scaling will be done by the next layer.
- Type
bool
- bias_init#
Initializer for bias, by default, zero.
- Type
Callable[[Any, Tuple[int, …], Any], Any]
- scale_init#
Initializer for scale, by default, one.
- Type
Callable[[Any, Tuple[int, …], Any], Any]
- reduction_axes#
Axes for computing normalization statistics.
- Type
Union[int, Sequence[int]]
- feature_axes#
Feature axes for learned bias and scaling.
- Type
Union[int, Sequence[int]]
- axis_name#
the axis name used to combine batch statistics from multiple devices. See jax.pmap for a description of axis names (default: None). This is only needed if the model is subdivided across devices, i.e. the array being normalized is sharded across devices within a pmap.
- Type
Optional[str]
- axis_index_groups#
groups of axis indices within that named axis representing subsets of devices to reduce over (default: None). For example, [[0, 1], [2, 3]] would independently batch-normalize over the examples on the first two and last two devices. See jax.lax.psum for more details.
- Type
Any
- use_fast_variance#
If true, use a faster, but less numerically stable, calculation for the variance.
- Type
bool
- __call__(x)[source]#
Applies layer normalization on the input.
- Parameters
x – the inputs
- Returns
Normalized inputs (the same shape as inputs).
Methods