An Intuitive Understanding on Tensor Dimension with Pytorch — Using torch.sum() as Example

Have you ever felt confused about how the tensor dimension works, especially when you need to apply operations like sum() over high dimension? The technique you learn in this post not only works for tensor dimension, but also works for Numpy ndarray dimension, because the logic of multi-dimensional array/tensor are the same.

This post is going to solve the puzzle in your head once and for good. It will provide you:

• An intuitive way to visualize how the operation sum() works on tensor dimension.
• An example using Pytorch to examine the tensor sum in code.

Shape (dimension) of the tensor

First, tensor is just another name for multi-dimensional array. When Mathematician has defined terms: scalar (0D), vector (1D), matrix (2D), we need a general term and concept for higher-D.

Here is a tensor, and we can find the shape of the tensor by counting the number of brackets in each level. (Note that the outermost bracket would be ignored, because it is there just to wrap the entire object.)

So here we have 2 elements in the outermost level (dim: 0), and 3 elements in the second level (dim: 1), and 3 elements in the innermost level (dim: 2). Thus, the shape of the tensor is (2,3,3).

Sum over Dimension 0

Next, we are going to sum over dimension 0, and we just need to follow the 3 easy steps:

• Identify elements in the dimension.
• Overlap these elements.
• Use X-ray eyes to sum up the numbers that appeared in the same position.

and that’s it! Step-by-step illustrations as follows:

Note: When you sum over a dimension i in a tensor, it would reduce that dimension i to 1, because you sum up every number in that dimension and reduce the multiple numbers into one sum. You will see that happening in every dimension.

Now, we already get the idea of how this work on dimension 0, let’s challenge the same steps on dimension 1.

Sum over Dimension 1

Note: The same dimensional reduction happens on dimension 1, after we sum over dimension 1, and the new shape after sum has reduced to (2,1,3).

We just need to carry the same steps to the last dimension.

Sum over Dimension 2

Note: The same dimensional reduction happens on dimension 2, after we sum over dimension 2, and the new shape after sum has reduced to (2,3,1).

Now we have an idea of how the tensor sum() operation works, let’s do it in code.

Using Pytorch to perform the tensor sum()

The following Jupyter Notebook shows how do we perform tensor sum() and examine our understanding on its dimension.

Note: In the function, you need to specify keepdim=True to retain its original dimension as we showcased in the first intuitive section. Otherwise, that dimension which was reduced to 1 would simply being removed.

This is the most intuitive visualization that I taught myself during the time I struggled understanding operators working on high tensor dimensions (such as sigmoid, max, sum etc.), and I hope this article fortifies your tensor intuition and makes you swift on identifying tensor dimension in your work! If you find this post to be useful, please leave a “clap” to this post. Your encouragement is what keeps me forward! Thanks for your reading!