Meaning / equivalent of numpy.einsum expression

Question

I'm desperately trying to find the python built-in equivalent of the following numpy.einsum expression:

>>> a = np.array((((1, 2), (3, 4)), ((5, 6), (7, 8))))
>>> a
array([[[1, 2],
        [3, 4]],

       [[5, 6],
        [7, 8]]])

>>> b = np.array((((9, 10), (11, 12)), ((13, 14), (15, 16))))
>>> b
array([[[ 9, 10],
        [11, 12]],

       [[13, 14],
        [15, 16]]])

>>> np.einsum("abc,abd->dc", a, b)
array([[212, 260],
       [228, 280]])

Do you mean standard-library Python operations (e.g. using `for` loops), or non-einsum NumPy (e.g. broadcasting multiplication)? If it's the latter, the expression `(a[:, :, None, :] * b[:, :, :, None]).sum(axis=(0, 1))` may be what you want. — Alex Riley, Sep 16 '19 at 21:07

score 0 · Accepted Answer · answered Sep 16 '19 at 21:54

As @AlexRiley comments the direct translation is something like:

(a[...,None,:]*b[...,None]).sum((0,1))

Let's parse the spec string 'abc,abd->dc' and let's rename the terms to x and y so they do not clash with the indices:

This is read as result_dc = &Sum;_ab x_abc y_abd

As you can see the indices are taken verbatim from the spec string. Indices that do not occur in result spec are summed over. And that's it.

Side note: We can do better than that: Merging the first two axes the expression can be read as a matrix product for which numpy uses a highly optimized code path:

b.reshape(-1,b.shape[-1]).T@a.reshape(-1,a.shape[-1])

This is more than twice as fast as the direct translation and also a bit faster than the original einsum.

Meaning / equivalent of numpy.einsum expression

1 Answers1