Summing over ellipsis broadcast dimension in numpy.einsum

Question

In numpy, I have an array that can be either 2-D or 3-D, and I would like to reduce it to 2-D while squaring each element. So I tried this and it doesn't work:

A = np.random.rand(5, 3, 3)
np.einsum('...ij,...ij->ij', A, A)

It returns this error:

ValueError: output has more dimensions than subscripts given in einstein sum, but no '...' ellipsis provided to broadcast the extra dimensions.

I suppose einsum doesn't assume that when the ellipsis goes away in the right hand side, I want to sum over the ellipsis dimension(s), if they exist. Is there some "elegant" way (i.e. without checking the number of dimensions and using an if statement) to tell it that I want to do this for 3-D:

A = np.random.rand(5, 3, 3)
np.einsum('aij,aij->ij', A, A)

and this for 2-D?

A = np.random.rand(3, 3)
np.einsum('ij,ij->ij', A, A)

Answer 1

Sometimes the 'elegant' way to handle variable dimensions is to use a set of if tests, and hide them in a function call. Look for example at np.atleast_3d; it has a 4way if/else clause. I'd recommend it here, except it adds the extra dimension at the end, not the start. if clauses using reshape are not expensive (time wise), so don't be afraid to use them. Even if you find some magical function, look at its code; you may be surprised what is hidden.

---

ellipsis is used for dimensions that 'go along for the ride', not ones where you want specific control. Here you want to sum over the initial dimension, so you need to index it explicitly:

In [161]: np.einsum('i...,i...',A,A)
Out[161]: 
array([[ 1.26942035,  1.32052776,  1.74118617],
       [ 1.59679765,  1.49331565,  2.04573002],
       [ 2.29027005,  1.48351522,  1.36679208]])
In [162]: np.einsum('aij,aij->ij',A,A)
Out[162]: 
array([[ 1.26942035,  1.32052776,  1.74118617],
       [ 1.59679765,  1.49331565,  2.04573002],
       [ 2.29027005,  1.48351522,  1.36679208]])

For 2d array:

In [165]: np.einsum('ij,ij->ij',A[0],A[0])
Out[165]: 
array([[ 0.20497776,  0.11632197,  0.65396968],
       [ 0.0529767 ,  0.24723351,  0.27559647],
       [ 0.62806525,  0.33081124,  0.57070406]])
In [166]: A[0]*A[0]
Out[166]: 
array([[ 0.20497776,  0.11632197,  0.65396968],
       [ 0.0529767 ,  0.24723351,  0.27559647],
       [ 0.62806525,  0.33081124,  0.57070406]])
In [167]: 
In [167]: np.einsum('...,...',A[0],A[0])
Out[167]: 
array([[ 0.20497776,  0.11632197,  0.65396968],
       [ 0.0529767 ,  0.24723351,  0.27559647],
       [ 0.62806525,  0.33081124,  0.57070406]])

I don't think you can handle both cases with one expression.

Another way to get the first sum

In [168]: (A*A).sum(axis=0)
Out[168]: 
array([[ 1.26942035,  1.32052776,  1.74118617],
       [ 1.59679765,  1.49331565,  2.04573002],
       [ 2.29027005,  1.48351522,  1.36679208]])

I contributed the patch that fixed the handling of ellipsis, but that was a couple of years ago. So the details aren't super fresh in my mind. As part of that I reverse engineered the parsing the string expression (the original is compiled), and could review that code (or refer you to it), if we need a more definitive answer.

---

In [172]: np.einsum('...ij,...ij->ij',A,A)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-172-dfe39e268402> in <module>()
----> 1 np.einsum('...ij,...ij->ij',A,A)

ValueError: output has more dimensions than subscripts given in 
 einstein sum, but no '...' ellipsis provided to broadcast 
 the extra dimensions.
In [173]: np.einsum('...ij,...ij->...ij',A,A).shape
Out[173]: (5, 3, 3)

The error message says that it is trying to pass the ... dimensions to the output, and can't - because the output is missing dimensions or .... In other words, it does not perform summation over ... dimensions. They pass to the output unchanged (broadcasting rules apply).