class

LayerNormBackward

LayerNormBackward

source edit

Python wrappers

Autograd node for Layer Normalization.

Computes $\hat{x} = (x - \mu) / \sqrt{\sigma^2 + \epsilon}$ over the last $D$ dimensions of x (one (\mu, \sigma^2) pair per leading-axis slice), then applies the element-wise affine $y = \gamma\,\hat{x} + \beta$ . The per-slice mean saved_mean_ and reciprocal standard deviation saved_rstd_ are cached for the backward closed-form.

Inherits FuncOp<LayerNormBackward, 3> with three saved-input slots — x, gamma, and beta. outer_ * N_ always equals numel(x): outer_ is the product of the leading (non-normalised) dims and N_ is the product of the trailing (normalised) dims.

Math

For each leading-axis slice of x (flattened length $N$ ):

\mu = \frac{1}{N}\sum_i x_i, \qquad \sigma^2 = \frac{1}{N}\sum_i (x_i - \mu)^2

\hat{x}_i = (x_i - \mu) \cdot \frac{1}{\sqrt{\sigma^2 + \epsilon}}, \qquad y_i = \gamma_i\,\hat{x}_i + \beta_i

Mean and variance use the biased estimator (correction=0), as is standard for layer norm.

References

[Ba, Kiros & Hinton, "Layer Normalization" (arXiv:1607.06450, 2016)](https://arxiv.org/abs/1607.06450).

Attributes

schema_v1OpSchema

Registered schema (name="layer_norm", version=1, amp_policy=ForceFP32, produces_grad=true).

saved_mean_Storage

Per-slice mean

\mu

, shape (outer_,).

saved_rstd_Storage

Per-slice reciprocal standard deviation

1 / \sqrt{\sigma^2 + \epsilon}

, shape (outer_,).

outer_std::size_t

Product of all leading (non-normalised) dimensions of x.

N_std::size_t

Product of all normalised (trailing) dimensions of x (== numel(gamma)).

Notes

Thread safety: instances are created during forward and consumed during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

int forward(const int & x, const int & gamma, const int & beta, double eps)

source edit

Run the forward pass.

Reduces x over its last $D$ dimensions (where $D$ is the rank of gamma), normalises, applies the affine, and (when grad mode is on) registers this as the grad_fn of the returned output. Caches saved_mean_ and saved_rstd_ for backward.

Parameters

xTensorImplPtr

Input tensor of any shape (*, normalized_shape).

gammaTensorImplPtr

Per-element scale; its shape must match the trailing dims of x.

betaTensorImplPtr

Per-element shift; must have the same shape as gamma.

epsdouble

Small constant added to the variance for numerical stability.

Returns

TensorImplPtr

Normalised output, same shape as x.

Raises

ShapeMismatch

If gamma.ndim > x.ndim, if the trailing dims of x do not match gamma, or if gamma.shape != beta.shape.

Methods

apply

→std::vector<Storage>

int apply(int grad_out)

source edit

Compute gradients for the three saved inputs.

Returns {dx, d_gamma, d_beta} in the saved-input order. Closed-form (per leading slice of length $N$ ):

d\gamma_i = \sum_{\text{slice}} \hat{x}_i\,\text{grad\_out}_i, \qquad d\beta_i = \sum_{\text{slice}} \text{grad\_out}_i

and the input gradient follows the standard LN backward expansion (a difference of the centred grad and two reduction terms scaled by $\text{rstd}$ ).

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

Three-element vector {dx, d_gamma, d_beta}.

class

LayerNormBackward

LayerNormBackward

source edit

Python wrappers

Autograd node for Layer Normalization.

Math

For each leading-axis slice of x (flattened length $N$ ):

\mu = \frac{1}{N}\sum_i x_i, \qquad \sigma^2 = \frac{1}{N}\sum_i (x_i - \mu)^2

\hat{x}_i = (x_i - \mu) \cdot \frac{1}{\sqrt{\sigma^2 + \epsilon}}, \qquad y_i = \gamma_i\,\hat{x}_i + \beta_i

Mean and variance use the biased estimator (correction=0), as is standard for layer norm.

References

[Ba, Kiros & Hinton, "Layer Normalization" (arXiv:1607.06450, 2016)](https://arxiv.org/abs/1607.06450).

Attributes

schema_v1OpSchema

Registered schema (name="layer_norm", version=1, amp_policy=ForceFP32, produces_grad=true).

saved_mean_Storage

Per-slice mean

\mu

, shape (outer_,).

saved_rstd_Storage

Per-slice reciprocal standard deviation

1 / \sqrt{\sigma^2 + \epsilon}

, shape (outer_,).

outer_std::size_t

Product of all leading (non-normalised) dimensions of x.

N_std::size_t

Product of all normalised (trailing) dimensions of x (== numel(gamma)).

Notes

Thread safety: instances are created during forward and consumed during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

int forward(const int & x, const int & gamma, const int & beta, double eps)

source edit

Run the forward pass.

Parameters

xTensorImplPtr

Input tensor of any shape (*, normalized_shape).

gammaTensorImplPtr

Per-element scale; its shape must match the trailing dims of x.

betaTensorImplPtr

Per-element shift; must have the same shape as gamma.

epsdouble

Small constant added to the variance for numerical stability.

Returns

TensorImplPtr

Normalised output, same shape as x.

Raises

ShapeMismatch

If gamma.ndim > x.ndim, if the trailing dims of x do not match gamma, or if gamma.shape != beta.shape.

Methods

apply

→std::vector<Storage>

int apply(int grad_out)

source edit

Compute gradients for the three saved inputs.

Returns {dx, d_gamma, d_beta} in the saved-input order. Closed-form (per leading slice of length $N$ ):

d\gamma_i = \sum_{\text{slice}} \hat{x}_i\,\text{grad\_out}_i, \qquad d\beta_i = \sum_{\text{slice}} \text{grad\_out}_i

and the input gradient follows the standard LN backward expansion (a difference of the centred grad and two reduction terms scaled by $\text{rstd}$ ).

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

Three-element vector {dx, d_gamma, d_beta}.