Search code examples
numpyarray-broadcastingnumpy-einsum

numpy.einsum with ellipses of different dimensionality

I often find that I'd like like to do an operation between the last few dimensions of two arrays, where the first dimensions don't necessarily match. As an example I'd like to do something like:

a = np.random.randn(10, 10, 3, 3)
b = np.random.randn(5, 3)
c = np.einsum('...ij, ,,,j -> ...,,,i', a, b)

and the result should satisfy c.shape = (10, 10, 5, 3) and c[i, j, k] = a[i, j] @ b[k]. Is there a way to achieve this with the existing interface?

Solution

  • ended up getting around the issue with the following helper function

    def batch_matvec(A, b):
    product = np.einsum('...ij, ...kj->...ki', A, b.reshape(-1, b.shape[-1]))
    return product.reshape((*A.shape[:-2], *b.shape[:-1], -1))