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

Question

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/647899105.py", 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/einsumfunc.py", 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)

np.__version__
Out[70]: '1.21.4'

Asked By: MathManM

||

Source

Answer 1

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

or

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

or

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

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

Question:

Answers: