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

user_10170285 0
Viewed 221

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()?

numpy-einsum numpy python
1 Answers

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.

user_73437210
Related