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)

Thanks in advance for your help.

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)

Answer 2

import numpy as np

a = np.random.rand(128, 16, 8, 32)
b = np.random.rand(256, 8, 32)

# Reshape 'a' to (128, 16, 256, 32)
a_reshaped = a.transpose(0, 2, 1, 3).reshape(128, 16, -1)

# Reshape 'b' to (256, 8, 32)
b_reshaped = b.transpose(0, 2, 1)

# Perform matrix multiplication between reshaped 'a' and 'b'
output_reshaped = np.matmul(a_reshaped, b_reshaped)

# Reshape 'output' to (256, 128, 16)
output = output_reshaped.reshape(256, 128, 16).transpose(1, 2, 0)

This code reshapes the input arrays a and b and then performs a matrix multiplication using np.matmul. The result is reshaped to match the desired output shape.