fn
broadcast_back_for_reduce
→Storageint broadcast_back_for_reduce(const int & grad, const int & grad_shape, const int & input_shape, const int & axes, bool keepdims, Dtype dt, Device device)Broadcast grad back to input_shape after a forward reduction.
The backward pass for any reduction op (sum, mean, var,
...) receives a gradient at the reduced shape and must re-expand it
across the squeezed axes so that the result matches the original input
shape. This helper performs that expansion, taking the same
axes / keepdims parameters as the forward call to know which
dimensions to broadcast back along.
See Also
reduce_output_shape — the inverse direction.
Parameters
gradconst Storage&Incoming gradient at the reduced shape.
grad_shapeconst Shape&Shape of
grad; must equal reduce_output_shape(input_shape, axes, keepdims).input_shapeconst Shape&Shape to broadcast back to (the original forward input shape).
axesconst std::vector<int>&Axis list used in the forward reduction.
keepdimsboolSame value as passed to the forward reduction.
dtDtypeWorking dtype.
deviceDeviceTarget device.
Returns
StorageNew storage of shape input_shape carrying the expanded gradient.
Raises
ShapeMismatchgrad_shape does not equal the expected reduced shape.