CODEWITHSUNDEEP
pythonnumpy
pang
pang

Reputation: 43

Numpy : What does this einsum expression mean and is there an alternative?

How do I understand following expression (A is an [200, 2] array) :

B = numpy.einsum('...i,...j->...ij',A，A)

And how to write it in another way not using numpy.einsum?

Upvotes: 2

Views: 162

Answers (1)

Divakar
Divakar

Reputation: 221634

Well that basically does element-wise multiplication between the elements along the last axis for A for all pairs. Now, since that einsum expression isn't doing any sum-reduction and is simply doing the job of broacasting, we can avoid it by manually extending the input arrays to have one more dimension and let the broadcasting do its job. This would be more performant than using einsum and that's a good reason for look for an alternative to einsum for such a case.

The implementation for a 2D and even a n-dim A (using the ellipsis notation) would be -

A[...,None,:]*A[...,None]

Sample run -

In [71]: A = np.random.rand(3,4,5,6)

In [72]: np.allclose(np.einsum('...i,...j->...ij',A,A), A[...,None,:]*A[...,None])
Out[72]: True

The performance aspect discussed earlier makes more sense, when the axis of expansion is of decent length. So, with A of shape (200,2) i.e. 2 as the axis length for the last axis, the improvement with extension of dims and using broadcasting would be less/not-at-all noticeable, but only seen with decent lengths. Let's verify these aspects -

In [76]: A = np.random.rand(20000,2)

In [77]: %timeit np.einsum('...i,...j->...ij',A,A)
1000 loops, best of 3: 207 µs per loop

In [78]: %timeit A[...,None,:]*A[...,None]
1000 loops, best of 3: 364 µs per loop

In [79]: A = np.random.rand(200,200)

In [80]: %timeit np.einsum('...i,...j->...ij',A,A)
100 loops, best of 3: 12.1 ms per loop

In [81]: %timeit A[...,None,:]*A[...,None]
100 loops, best of 3: 9.74 ms per loop

Upvotes: 2

Related Questions