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