neural_tangents.stax.LayerNorm(axis=-1, eps=1e-12, batch_axis=0, channel_axis=-1)[source]

Layer normalisation.

  • axis (Union[int, Sequence[int]]) – dimensions over which to normalize.

  • eps (float) – (small) positive constant to be added to the variance estimates in order to prevent division by zero.

  • batch_axis (int) – batch dimension. Defaults to 0, the leading axis.

  • channel_axis (int) – channel / feature dimension. Defaults to -1, the trailing axis. For kernel_fn, channel size is considered to be infinite.

Return type:

tuple[InitFn, ApplyFn, LayerKernelFn]


(init_fn, apply_fn, kernel_fn).