Autograd node for Lp normalisation along a single axis.
Computes along axis, dividing
the input by its Lp-norm per slice while clamping the denominator
from below by eps to avoid division by zero. axis is
resolved to a non-negative index inside forward so the
saved value is always normalised.
Inherits FuncOp<LpNormalizeBackward, 1> with a single
saved-input slot — x. The per-slice norm
saved_norm_ is cached for the backward rule.
Math
Let denote a slice along the chosen axis (the elements summed over):
Attributes
schema_v1OpSchemaname="lp_normalize", version=1, amp_policy=ForceFP32, produces_grad=true).ord_double2.0 (Euclidean L2 norm).axis_intforward are converted before storage.eps_double1e-12.saved_norm_StorageStatic methods
1Run the forward pass.
Computes the per-slice Lp norm along axis, divides x by
the clamped norm, caches saved_norm_, and (when grad mode is
on) registers this as the grad_fn of the returned
output. axis is normalised to a non-negative index before
being stored in axis_.
Parameters
xTensorImplPtrorddouble2.0 for the L2 norm).axisintepsdoubleReturns
TensorImplPtrNormalised output, same shape as x.
Methods
1Compute the gradient with respect to x.
Returns {dx}. The closed-form involves a projection of
grad_out along the normalised slice direction; see the .cpp
for the explicit formula.
Parameters
grad_outStorageReturns
std::vector<Storage>Single-element vector {dx}.