Struct LayerNorm 

pub struct LayerNorm {
    pub weight: Option<Param>,
    pub bias: Option<Param>,
    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<Param>

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

bias: Option<Param>

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

eps: f32

Added to the variance (MLX’s default 1e-5).

Implementations

impl LayerNorm

pub fn new(dims: i32) -> Self

A layer over dims features with MLX’s defaults: affine with a bias, eps = 1e-5.

pub fn eps(self, eps: f32) -> Self

Sets eps.

pub fn affine(self, affine: bool) -> Self

With (true, the default) or without a learned scale and shift (MLX’s affine).

pub fn without_bias(self) -> Self

Without the learned shift (MLX’s bias=False).

pub fn forward(&self, x: &Array) -> 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 Leaves for LayerNorm

fn visit(&self, __v: &mut Visit<'_>)

Reports every leaf (and child tree, enum variant, mode flag and static field) to v, in order.

impl Tree for LayerNorm

where Self: Clone,

fn visit_mut(&mut self, __v: &mut VisitMut<'_>)

Reports every leaf (and child tree and mode flag) to v mutably, in the order of Leaves::visit.

fn for_each(&self, f: impl FnMut(&str, &Array))

Calls f(path, leaf) for every leaf, in order.

fn map(&self, f: impl FnMut(&str, &Array) -> Array) -> Self

The tree with every leaf replaced by f(path, leaf) (same structure; parameters keep their frozen flag).

fn leaves(&self) -> Vec<Array>

The leaves, in order.

fn paths(&self) -> Vec<String>

The leaves’ paths, in order.

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

The leaves with their paths, in order (MLX’s tree_flatten).

fn unflatten( &self, leaves: impl IntoIterator<Item = Array>, ) -> Result<Self, Error>

A tree of this structure (self is the template) with leaves in order: exactly Tree::len of them, else an error naming both counts.

fn load_flat( &mut self, named: impl IntoIterator<Item = (String, Array)>, mode: LoadMode, ) -> Result<(), Error>

Replaces leaves by path. LoadMode::Strict requires exactly the tree’s paths with matching shapes; LoadMode::Partial assigns the matching names and ignores unknown ones (shapes must still match). Loaded dtypes are kept. On any problem the error lists every one (missing, unknown and mismatched names, sorted) and the tree is unchanged. Duplicate paths (in the tree or in named) are an error.

fn len(&self) -> usize

The number of leaves.

fn is_empty(&self) -> bool

Whether the tree has no leaves.

fn def(&self) -> TreeDef

The tree’s structure.