Struct LayerNorm 

pub struct LayerNorm {
    pub weight: Option<Array>,
    pub bias: Option<Array>,
    pub eps: f32,
}

Layer normalisation over the last axis, optionally with a learned scale and bias (MLX’s nn.LayerNorm, the fused fast::layer_norm).

Fields

weight: Option<Array>

The scale, [dims] (ones), if affine.

bias: Option<Array>

The shift, [dims] (zeros), if affine with a bias.

eps: f32

Added to the variance.

Implementations

impl LayerNorm

pub fn new(dims: usize, eps: f32, affine: bool, bias: bool) -> Result<Self>

A layer over dims features: affine adds a scale, and bias (with affine) a shift. MLX’s defaults: eps = 1e-5, affine with a bias.

pub fn forward(&self, x: &Array) -> Result<Array>

The normalised x.

Trait Implementations

impl Clone for LayerNorm

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for LayerNorm

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Module for LayerNorm

fn visit(&mut self, p: &mut Parameters<'_>)

Reports every parameter and submodule to p, frozen ones included.

fn parameters(&self) -> Vec<(String, Array)>

Every parameter with its path (layers.0.weight), in visiting order.

fn trainable_parameters(&self) -> Vec<(String, Array)>

The parameters that are trained (not frozen), with their paths.

fn update_trainable(&mut self, values: &[Array]) -> Result<()>

Replaces the trainable parameters, given in the order of Module::trainable_parameters (as optimizers return them).

fn load_weights( &mut self, weights: impl IntoIterator<Item = (String, Array)>, strict: bool, ) -> Result<()>

Replaces parameters by path, as MLX’s load_weights: each weight must have its parameter’s shape (the dtype is kept as loaded). With strict, every parameter must be given and no unknown names are allowed.