Numpy einsum fails for matrix product with unspecified leading and trailing axes


I want to use dot product between the leading and trailing axes of arrays with unspecified dimensions. einsum doesn’t seem to accommodate this use case. User error?

With prnd=np.random.RandomState(),

np.einsum('...i,i...', prnd.rand(5,2), prnd.rand(2,3)) 
Traceback (most recent call last):

  File "/tmp/ipykernel_1106807/", line 1, in <module>
    np.einsum('...i,i...', prnd.rand(5,2), prnd.rand(2,3))

  File "<__array_function__ internals>", line 5, in einsum

  File "/home/elliot/anaconda3/envs/current/lib/python3.9/site-packages/numpy/core/", line 1359, in einsum
    return c_einsum(*operands, **kwargs)

ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (5,2)->(5,2) (2,3)->(3,2) 

should be the same as…

np.einsum('ai,ib', prnd.rand(5,2), prnd.rand(2,3)).shape
Out[69]: (5, 3)

Out[70]: '1.21.4'
Asked By: MathManM



First of all, it has nothing to do with np.random.RandomState(), it works just the same with e.g. np.ones. Note that if you use ellipsis ... as part of the input shapes, the ellipsis will stand for the same (up to the rules of broadcasting) shape throughout the whole expression. In your case you try (slightly rewritten)

a = np.ones((5, 2))
b = np.ones((2, 3))
np.einsum('...i,i...->...', a, b)

But in a the ellipsis stands for 5 while in b the ellipsis stands for 3 which cannot be broadcast together, that is there is no possible output shape satisfying the broadcasting rules.

What would work is

a = np.ones((5, 2))
b = np.ones((2, 5))


a = np.ones((5, 2))
b = np.ones((2, 1))


a = np.ones((5, 2))
b = np.ones((2,))

But you can’t use ellipsis to "glue together" the trailing or leading dimensions to a new shape.

Answered By: flawr
Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.