Because a non-reduction sum retains its original axes after summing along a specific dimension, the resulting tensor can be directly used in elementwise operations with the original tensor via broadcasting. For instance, dividing a matrix by its non-reduction row sum seamlessly triggers the broadcasting mechanism to expand the sum tensor, thereby generating a new matrix where the elements of each row mathematically sum up to $$1$$.

Broadcasting with Non-Reduction Summation

Sometimes it is useful to keep the number of axes unchanged when calculating the sum or mean of a tensor. A non-reduction sum calculates the total along a specified axis but retains that axis with a size of $$1$$ in the output shape, rather than completely eliminating the dimension. Maintaining the original number of axes is critical when the resulting tensor will be used in subsequent mathematical operations that rely on the broadcasting mechanism.

Claude

Instead of reducing an entire tensor to a single scalar, mathematical reduction operations like summation can be targeted along specific axes. For instance, summing an $$m \times n$$ matrix along its rows (axis $$0$$) eliminates the row dimension, yielding an output vector. Similarly, summing along the column dimension (axis $$1$$) produces a vector by adding elements across columns. Reducing a tensor along all of its axes simultaneously is mathematically equivalent to computing the total sum of all its elements.

Tensor Reduction along an Axis

Dive into Deep Learning

Deep learning libraries allow developers to specify the exact axes along which a tensor should be reduced when invoking summation functions. The dimension corresponding to the specified reduction axis is subsequently removed from the shape of the resulting output tensor. In PyTorch, MXNet, and JAX, this is achieved by providing the `axis` parameter (e.g., `A.sum(axis=0)` or `A.sum(axis=[0, 1])`), while TensorFlow utilizes the same parameter within its `tf.reduce_sum(A, axis=0)` function.

Learn Before

Related

Learn After