Broadcast size summation 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 by squaring each element. So I tried this and it does not 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 believe that einsum does not assume that when the ellipsis goes off on the right side, I want to summarize the dimension of the ellipsis, if they exist. Is there some "elegant" way (i.e., without checking the number of measurements and using the if statement) to say that I want to do this for 3-D:

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

and is it 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. See an example at np.atleast_3d ; it has a 4way if / else condition. I would recommend it here, except that it adds an extra dimension to the end rather than the beginning. if offers using reshape are not expensive (temporary), so don't be afraid to use them. Even if you find a magic function, look at its code; You may be surprised at what is hidden.

---

Ellipsis is used for dimensions that "go ahead for the ride," not those where you want some control. Here you want to summarize the original dimension, so you need to explicitly index it:

 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 do not think that you can handle both cases with one expression.

Another way to get the first amount

 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 made a fix correcting the ellipsis processing, but that was a couple of years ago. So the details are not super fresh in my mind. As part of this, I drew the parsing construct of the syntax expression (the original is compiled) and could look at this code (or pass it to him) 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 he is trying to transfer the sizes ... to the output and cannot - because there are no sizes on the output or ... In other words, it does not perform a summation over the dimensions ... They are transmitted to the exit unchanged (broadcast rules apply).