fn
reduce_grad_to_shape
→Storageint reduce_grad_to_shape(const int & grad, const int & grad_shape, const int & target_shape, Dtype dtype, Device device)Reduce grad back to target_shape by summing over broadcast axes.
Used in the backward pass of any binary op that broadcast one operand
against the other. The gradient initially has the broadcast (output)
shape; this helper walks the trailing dimensions and sums along any
axis where target_shape is size 1 or absent, so the result matches
the original input shape ready to be accumulated into .grad.
Math
For each output index i,
Parameters
gradconst Storage&Incoming gradient at the broadcast shape.
grad_shapeconst Shape&Shape of
grad.target_shapeconst Shape&Original (pre-broadcast) input shape to reduce back to.
dtypeDtypeWorking dtype.
deviceDeviceTarget device.
Returns
StorageStorage of shape target_shape with broadcast axes summed.
Notes
No-op when grad_shape == target_shape.