numpy einsum collapse all but first dimension

Question

I have an array of dimension (5,...) (... can be anything) and I would like to form the dot product of all dimension after the 5 such that the resulting array has shape (5,). I thought the einsum

i...,i...->i

looked promising, but alas

import numpy as np

a = np.random.rand(5)
b = np.einsum("i...,i...->i", a, a)
assert b.shape == (5,)  # okay


a = np.random.rand(5, 2)
b = np.einsum("i...,i...->i", a, a)  # fail
assert b.shape == (5,)


a = np.random.rand(5, 2, 3)
b = np.einsum("i...,i...->i", a, a)  # fail
assert b.shape == (5,)

in einsum
    return c_einsum(*operands, **kwargs)
ValueError: output has more dimensions than subscripts given in einstein sum, but no '...' ellipsis provided to broadcast the extra dimensions.

I could perhaps reshape a into

b = np.einsum("ij,ij->i", a.reshape(a.shape[0], -1), a.reshape(a.shape[0], -1))

but this looks so messy. Any better suggestion?

I was thinking that `np.tensordot(x,x,axes=([1,2,3],[1,2,3])` would do it, but for a (5,3,3,3) shape, it produces a (5,5) output. Can that then be summed? — Tim Roberts, Mar 05 '21 at 20:36
`x = a.reshape(a.shape[0], -1); np.einsum("ij,ij->i", x,x)` looks good to me. Aslo `(a*a).reshape(a.shape[0],-1).sum(-1)` would work. — Quang Hoang, Mar 05 '21 at 20:42

hpaulj · Answer 1 · 2021-03-05T21:30:03.787

According to the docs, ellipsis "enable and control broadcasting". You want to do the dot sum-of-products on these dimensions. Those dimensions require explicit entires. So for a 4d with reduction on the last 2 dimensions, '...ij->...'

In [15]: arr = np.arange(24).reshape(2,3,4)

reduction on last or last 2 dimensions:

In [17]: np.einsum('...i->...',arr).shape
Out[17]: (2, 3)
In [18]: np.einsum('...ij->...',arr).shape
Out[18]: (2,)

sum equivalents:

In [20]: arr.sum(axis=(-1))
Out[20]: 
array([[ 6, 22, 38],
       [54, 70, 86]])
In [21]: arr.sum(axis=(-1)).shape
Out[21]: (2, 3)
In [22]: arr.sum(axis=(-2,-1)).shape
Out[22]: (2,)
In [23]: arr.sum(axis=(-2,-1))
Out[23]: array([ 66, 210])
In [24]: np.einsum('...ij->...',arr)
Out[24]: array([ 66, 210])

This expression from sympy may help visualize einsum

In [37]: f = Function('f')
    ...: Sum(f(x,y,z), (y, 0, 3),(z,0,2))
Out[37]: 
  2     3             
 ___   ___            
 ╲     ╲              
  ╲     ╲             
  ╱     ╱   f(x, y, z)
 ╱     ╱              
 ‾‾‾   ‾‾‾            
z = 0 y = 0

f is function of 3 variables (3d). einsum notation emulates Einstein notation, where dimensions that should be summed are explicitly listed. Ones that aren't summed, that in a sense, are just carried along, don't need to be (as) explicit.

PG11 · Answer 2 · 2021-03-05T23:05:55.920

0

For illustration, assume tensor A is a 2-by-2 matrix. The dot-product you need can be written as

C_i = A_{kj} A_{ij} D_{ki},

where D is the identity matrix.

Using np.einsum, this can be written as

D= np.identity(2)
X= np.einsum("k...,i...,ki->i...", A, A, D)
C= np.sum(X.reshape(X.shape[0],-1),axis=-1)

We need np.sum in the above to sum over all j indices.

edited Mar 05 '21 at 23:05

answered Mar 05 '21 at 22:02

PG11

155
5

"kj,ij,ki->i" would sum over `j` directly. – hpaulj Mar 05 '21 at 22:39
Yes but that requires us to manually enter indices for all dimensions and hence not generalizable. The code for a 2x2x3 matrix, for example, will be different from that for a 2x2 matrix. – PG11 Mar 05 '21 at 22:48
But the `sum` still nneds the axes. – hpaulj Mar 05 '21 at 22:52
You are right. I updated the answer to generalize it by flattening the array and summing over the first dimension – PG11 Mar 05 '21 at 23:07
I think with this modification my solution looks less elegant than the one posted in the question – PG11 Mar 05 '21 at 23:18

numpy einsum collapse all but first dimension

2 Answers2