python - Summing over ellipsis broadcast dimension in numpy.einsum -

in numpy, have array can either 2-d or 3-d, , reduce 2-d while squaring each element. tried , doesn't work:

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

it returns error:

valueerror: output has more dimensions subscripts given in einstein sum, no '...' ellipsis provided broadcast dimensions. 

i suppose einsum doesn't assume when ellipsis goes away in right hand side, want sum on ellipsis dimension(s), if exist. there "elegant" way (i.e. without checking number of dimensions , using if statement) tell want 3-d:

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

and 2-d?

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

sometimes 'elegant' way handle variable dimensions use set of if tests, , hide them in function call. example @ np.atleast_3d; has 4way if/else clause. i'd recommend here, except adds dimension @ end, not start. if clauses using reshape not expensive (time wise), don't afraid use them. if find magical function, @ code; may surprised hidden.

---

ellipsis used dimensions 'go along ride', not ones want specific control. here want sum on initial dimension, need index 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 can handle both cases 1 expression.

another way 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 patch fixed handling of ellipsis, couple of years ago. details aren't super fresh in mind. part of reverse engineered parsing string expression (the original compiled), , review code (or refer it), if need 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 subscripts given in   einstein sum, no '...' ellipsis provided broadcast   dimensions. in [173]: np.einsum('...ij,...ij->...ij',a,a).shape out[173]: (5, 3, 3) 

the error message says trying pass ... dimensions output, , can't - because output missing dimensions or .... in other words, not perform summation on ... dimensions. pass output unchanged (broadcasting rules apply).