class

BatchNormNdBackward

BatchNormNdBackward

source edit

Python wrappers

Autograd node for training-mode Batch Normalization with N spatial axes.

Computes per-channel statistics over the (B, S_0, ..., S_{N-1}) axes of the input, normalises the input by them, and applies the learnable per-channel affine $y = \gamma \hat{x} + \beta$ . The per-channel mean $\mu_c$ and reciprocal standard deviation $\text{rstd}_c = 1 / \sqrt{\sigma_c^2 + \epsilon}$ are saved in saved_mean_ / saved_rstd_ for the backward closed-form.

Inherits FuncOp<BatchNormNdBackward<N>, 3> with three saved-input slots — x, gamma, and beta. The template is explicitly instantiated for N = 1, 2, 3 in the .cpp file and exposed via the aliases BatchNorm1dBackward, BatchNorm2dBackward, and BatchNorm3dBackward.

Math

\mu_c = \frac{1}{B\,\prod_i S_i} \sum_{b,\,\mathbf{s}} x_{b,c,\mathbf{s}}, \qquad \sigma_c^2 = \frac{1}{B\,\prod_i S_i} \sum_{b,\,\mathbf{s}} (x_{b,c,\mathbf{s}} - \mu_c)^2

\hat{x}_{b,c,\mathbf{s}} = (x_{b,c,\mathbf{s}} - \mu_c) \cdot \frac{1}{\sqrt{\sigma_c^2 + \epsilon}}, \qquad y = \gamma_c\,\hat{x} + \beta_c

The Python-side running-stats update applies, in order,

\hat{\mu}_t = (1 - m)\,\hat{\mu}_{t-1} + m\,\mu_c, \qquad \hat{\sigma}^2_t = (1 - m)\,\hat{\sigma}^2_{t-1} + m \cdot \frac{n}{n-1}\,\sigma_c^2,

where $n = B \prod_i S_i$ is the per-channel reduction count (the Bessel correction $n/(n-1)$ is applied only to the running variance buffer, never to the normalisation itself).

Shape

Input x : (B, C, S_0, ..., S_{N-1})
Output : same shape as x

References

Ioffe & Szegedy, "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift" (ICML 2015).

Attributes

schema_v1OpSchema

Registered op schema. name is "batch_norm1d", "batch_norm", or "batch_norm3d" depending on N; version=1; amp_policy=ForceFP32; produces_grad=true.

saved_mean_Storage

Per-channel mean

\mu_c

of shape (C,), computed by the backend and saved for backward.

saved_rstd_Storage

Per-channel reciprocal standard deviation

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

, shape (C,).

B_, C_int

Batch and channel counts captured from x.shape() for the backward.

S_int[N]

Spatial sizes captured from x.shape()[2:] (length-1 dummy array for the N==0 corner case, which is never instantiated).

eps_double

The forward's eps; preserved so that the GPU MPSGraph backward can derive variance from

\sigma_c^2 = 1 / \text{rstd}_c^2 - \epsilon

Notes

Thread safety: an instance is created once during forward and consumed once during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

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

source edit

Run the training-mode forward pass.

Computes per-channel batch statistics, normalises x, applies the affine, saves $\mu_c$ and the reciprocal std for backward, and (when grad mode is on) registers this as the grad_fn of the returned output. Dispatches to IBackend::batch_norm_forward which returns [y, mean, rstd] in a single fused call.

Parameters

xTensorImplPtr

Input tensor of shape (B, C, S_0, ..., S_{N-1}).

gammaTensorImplPtr

Per-channel scale of shape (C,).

betaTensorImplPtr

Per-channel shift of shape (C,).

epsdouble

Numerical-stability constant added inside the square root.

Returns

TensorImplPtr

Normalised output, same shape as x. Forward. When running_mean / running_var are non-null, the EMA update happens INSIDE this call using the same mean+rstd the backend already computed for the saved tensors — no second reduction over x. momentum is the EMA coefficient (must be finite; cumulative mode is handled in the Python wrapper).

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 using IBackend::batch_norm_backward and the cached saved_mean_ / saved_rstd_. Closed-form summary (per channel $c$ ):

d\gamma_c = \sum_{b,\mathbf{s}} \hat{x}_{b,c,\mathbf{s}}\, \text{grad\_out}_{b,c,\mathbf{s}}, \qquad d\beta_c = \sum_{b,\mathbf{s}} \text{grad\_out}_{b,c,\mathbf{s}}

and the input gradient follows the standard BN backward expansion.

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

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

class

BatchNormNdBackward

BatchNormNdBackward

source edit

Python wrappers

Autograd node for training-mode Batch Normalization with N spatial axes.

Math

\mu_c = \frac{1}{B\,\prod_i S_i} \sum_{b,\,\mathbf{s}} x_{b,c,\mathbf{s}}, \qquad \sigma_c^2 = \frac{1}{B\,\prod_i S_i} \sum_{b,\,\mathbf{s}} (x_{b,c,\mathbf{s}} - \mu_c)^2

\hat{x}_{b,c,\mathbf{s}} = (x_{b,c,\mathbf{s}} - \mu_c) \cdot \frac{1}{\sqrt{\sigma_c^2 + \epsilon}}, \qquad y = \gamma_c\,\hat{x} + \beta_c

The Python-side running-stats update applies, in order,

\hat{\mu}_t = (1 - m)\,\hat{\mu}_{t-1} + m\,\mu_c, \qquad \hat{\sigma}^2_t = (1 - m)\,\hat{\sigma}^2_{t-1} + m \cdot \frac{n}{n-1}\,\sigma_c^2,

where $n = B \prod_i S_i$ is the per-channel reduction count (the Bessel correction $n/(n-1)$ is applied only to the running variance buffer, never to the normalisation itself).

Shape

Input x : (B, C, S_0, ..., S_{N-1})
Output : same shape as x

References

Ioffe & Szegedy, "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift" (ICML 2015).

Attributes

schema_v1OpSchema

Registered op schema. name is "batch_norm1d", "batch_norm", or "batch_norm3d" depending on N; version=1; amp_policy=ForceFP32; produces_grad=true.

saved_mean_Storage

Per-channel mean

\mu_c

of shape (C,), computed by the backend and saved for backward.

saved_rstd_Storage

Per-channel reciprocal standard deviation

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

, shape (C,).

B_, C_int

Batch and channel counts captured from x.shape() for the backward.

S_int[N]

Spatial sizes captured from x.shape()[2:] (length-1 dummy array for the N==0 corner case, which is never instantiated).

eps_double

The forward's eps; preserved so that the GPU MPSGraph backward can derive variance from

\sigma_c^2 = 1 / \text{rstd}_c^2 - \epsilon

Notes

Thread safety: an instance is created once during forward and consumed once during backward — no concurrent access expected.

Static methods

static

forward

→TensorImplPtr

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

source edit

Run the training-mode forward pass.

Parameters

xTensorImplPtr

Input tensor of shape (B, C, S_0, ..., S_{N-1}).

gammaTensorImplPtr

Per-channel scale of shape (C,).

betaTensorImplPtr

Per-channel shift of shape (C,).

epsdouble

Numerical-stability constant added inside the square root.

Returns

TensorImplPtr

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 using IBackend::batch_norm_backward and the cached saved_mean_ / saved_rstd_. Closed-form summary (per channel $c$ ):

d\gamma_c = \sum_{b,\mathbf{s}} \hat{x}_{b,c,\mathbf{s}}\, \text{grad\_out}_{b,c,\mathbf{s}}, \qquad d\beta_c = \sum_{b,\mathbf{s}} \text{grad\_out}_{b,c,\mathbf{s}}

and the input gradient follows the standard BN backward expansion.

Parameters

grad_outStorage

Incoming gradient, same shape as the forward output.

Returns

std::vector<Storage>

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