Layer normalization — mlx_layer_norm • Rmlx

Normalizes inputs across the feature dimension.

Usage

mlx_layer_norm(normalized_shape, eps = 1e-05, device = mlx_default_device())

Arguments

normalized_shape: Size of the feature dimension to normalize.
eps: Small constant for numerical stability (default: 1e-5).
device: Execution target: supply "gpu", "cpu", or an mlx_stream created via mlx_new_stream(). Defaults to the current mlx_default_device() unless noted otherwise (helpers that act on an existing array typically reuse that array's device or stream).

Value

An mlx_module applying layer normalization.

See also

mlx.nn.LayerNorm

Examples

set.seed(1)
ln <- mlx_layer_norm(4)
x <- as_mlx(matrix(rnorm(12), 3, 4))
mlx_forward(ln, x)
#> mlx array [3 x 4]
#>   dtype: float32
#>   device: gpu
#>   values:
#>            [,1]       [,2]       [,3]       [,4]
#> [1,] -1.0668312  1.5259182 0.23306273 -0.6921500
#> [2,] -0.9833397 -0.7005272 0.09211669  1.5917501
#> [3,] -1.0064702 -0.9834564 1.13609314  0.8538334