The function of sum() in python

Question

I am reading some code with python, which is about probability distribution. But in many cases, instead of using a framework like PyTorch to compute distribution, it uses the following code

def sample_energy_0(self, y, M):
    device = next(self.parameters()).device
    x = torch.randn(M, y.shape[0], self.latent_dim).to(device)
    return x

def energy_0(self, x, y):
    return (x**2).sum(axis=2, keepdims=True) / 2

How does the function sum(axis=2) work here?

Or maybe it's using numpy. Either way, `x` is a numpy array, and it has a `sum()` method that calculates the sum of that axis of the array. — Barmar, Jan 06 '23 at 00:56

score 1 · Answer 1 · answered Jan 06 '23 at 01:03

x.sum(axis=2) means for a tensor of shape (A, B, C, ...), the elements along the C dimension are summed to yield a new tensor of shape (A, B, ...). keepdims=True means that the C dimension is "kept" as a dummy dimension of 1:

>>> import torch
>>> x = torch.arange(0, 30).reshape(2,3,5)
>>> x
tensor([[[ 0,  1,  2,  3,  4],
         [ 5,  6,  7,  8,  9],
         [10, 11, 12, 13, 14]],

        [[15, 16, 17, 18, 19],
         [20, 21, 22, 23, 24],
         [25, 26, 27, 28, 29]]])
>>> x.shape
torch.Size([2, 3, 5])
>>> x.sum(axis=2).shape
torch.Size([2, 3])
>>> x.sum(axis=2, keepdims=True).shape
torch.Size([2, 3, 1])

score 1 · Answer 2 · answered Jan 06 '23 at 01:04

The expression (x**2).sum(axis=2, keepdims=True) / 2 isn't calling the function sum(), it's calling a method named .sum() on some object.

The type of the object is unknown from looking at your code alone, but it seems likely that x is expected to be a Tensor, which has a .sum() method, which you can find in the documentation. However, since a keyword parameter called axis is passed, it's also quite possible the function expects a numpy array, and it is calling the numpy.sum().

The (x**2) applies the ** (raising to a power) operation on x and 2 and the result is typically the same type as a type of x. So, the .sum() method can be called on the result.

Note that you could call energy_0() with x being either a Tensor or a numpy array and it would work just fine, as long as the code in the function is compatible with the interface of the class of the object being passed.

The function of sum() in python

2 Answers2