Autograd node for inference-mode Batch Normalization.
Uses caller-supplied running mean and variance — not computed from
the current batch — to normalise the input. This is the path the
Python BatchNormNd module dispatches to once
model.eval() has been called (and track_running_stats=True).
Computes
per channel using the supplied
mean and var tensors, then applies the per-channel affine
. The reciprocal standard deviation
rstd_ is cached for the backward closed-form.
Inherits FuncOp<BatchNormEvalBackward, 5> with five saved-input
slots — x, mean, var, gamma, beta — because
during training-time fine-tuning users may want gradients with
respect to all five operands. In typical inference usage only x
requires grad, in which case the four remaining gradient slots are
trivially zero.
Math
See Also
BatchNorm2dBackward : Training-mode counterpart that computes statistics from the current batch.
Attributes
schema_v1OpSchemaname="batch_norm_eval", version=1, amp_policy=ForceFP32, produces_grad=true).eps_doubleeps; retained for the backward.rstd_Storage(C,).Static methods
1forward
→TensorImplPtrint forward(const int & x, const int & mean, const int & var, const int & gamma, const int & beta, double eps)Run the inference-mode forward pass.
Normalises x channel-wise using the supplied mean and
var (typically the buffers running_mean / running_var
accumulated during training), then applies the affine. Caches
rstd_ for backward and (when grad mode is on) registers
this as the grad_fn of the returned output.
Parameters
xTensorImplPtr(B, C, ...).meanTensorImplPtr(C,).varTensorImplPtr(C,).gammaTensorImplPtr(C,).betaTensorImplPtr(C,).epsdoubleReturns
TensorImplPtrNormalised output, same shape as x.
Methods
1Compute gradients for the five saved inputs.
Returns {dx, d_mean, d_var, d_gamma, d_beta} in the
saved-input order. In standard eval usage the gradients with
respect to mean, var, gamma, and beta are zero
(running stats and the affine are treated as frozen constants),
but the kernel computes them anyway so that fine-tuning use
cases work.
Parameters
grad_outStorageReturns
std::vector<Storage>Five-element vector {dx, d_mean, d_var, d_gamma, d_beta}.