replacing einsum with normal operations

Question

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?

Answer 1

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)