I need to replace einsum operation with standard numpy operations in the following code:
import numpy as np
a = np.random.rand(128, 16, 8, 32)
b = np.random.rand(256, 8, 32)
output = np.einsum('aijb,rjb->ira', a, b)
How would I do that?
One option would be to align to a similar shape and broadcast multiply, then sum and reorder the axes:
output2 = (b[None, None]*a[:,:,None]).sum(axis=(-1, -2)).transpose((1, 2, 0))
# assert np.allclose(output, output2)
But this is much less efficient as it's producing a large intermediate (shape (128, 16, 256, 8, 32)):
# np.einsum('aijb,rjb->ira', a, b)
68.9 ms ± 23.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# (b[None, None]*a[:,:,None]).sum(axis=(-1, -2)).transpose((1, 2, 0))
4.66 s ± 1.65 s per loop (mean ± std. dev. of 7 runs, 1 loop each)
Shapes:
# b[None, None].shape
#a i r j b
(1, 1, 256, 8, 32)
# a[:,:,None].shape
# a i r j b
(128, 16, 1, 8, 32)