Search code examples
pythonnumpynumpy-einsum

Numpy array addition made as simple as np.einsum()?


If I have a.shape = (3,4,5) and b.shape = (3,5), using np.einsum() makes broadcasting then multiplying the two arrays super easy and explicit:

result = np.einsum('abc, ac -> abc', a, b)

But if I want to add the two arrays, so far as I can tell, I need two separate steps so that the broadcasting happens properly, and the code feels less explicit.

b = np.expand_dims(b, 1)
result = a + b

Is there way out there that allows me to do this array addition with the clarity of np.einsum()?


Solution

  • Broadcasting can occur only on one extra dimension. For adding these two arrays one could expand them in a one-liner as follows:

    import numpy as np
    a = np.random.rand(3,4,5); b = np.random.rand(3,5);
    c = a + b[:, None, :] # c is shape of a, broadcasting occurs along 2nd dimension
    

    Note this is not any different than c = a + np.expand_dim(b, 1). In terms of clarity it is a personal style thing. I prefer broadcasting, others prefer einsum.