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.

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

Programmatic Tensor Reduction along an Axis

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.

Learn Before

Related