neural_tangents.stax.LayerNorm
- neural_tangents.stax.LayerNorm(axis=-1, eps=1e-12, batch_axis=0, channel_axis=-1)[source]
Layer normalisation.
- Parameters
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 to0
, the leading axis.channel_axis (
int
) – channel / feature dimension. Defaults to-1
, the trailing axis. Forkernel_fn
, channel size is considered to be infinite.
- Return type
- Returns
(init_fn, apply_fn, kernel_fn)
.