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