Ultimate flexibility with einsum ellipses

Question

I have a question about einsum ellipsis that I thought would be somewhere on StackExchange for sure, but somehow I can't seem to find.

Essentially I have some code which does lots of matrix and vector contractions using numpy's einsum. The input is usually some parameters that are then used to create vectors and matrices. The code works well, but now I would like to generalize it so that the input parameters can be scanned over a certain range. The nicest thing to do would be to make them vectors and modify my einsum expressions such that they accept an arbitrary number of additional dimensions that are simply carried through. This question is to ask if this is possible and if so how.

So in my view, this problem boils down to the following. Say I have an einsum expression that creates that does some kind of matrix multiplication, e.g.

c = np.einsum('ij,jk->ik', a, b)

Now I want to add an arbitrary number of indices to both a and b and simply add them as extra indices in the final matrix, e.g.

c = np.einsum('ijabc,jkde->ikabcde', a, b)

Now when you only do that for one of either a or b, you can do this easily by an ellipsis

c = np.einsum('ij...,jk->ik...', a, b)

So my question is whether you can have multiple ellipses in an einsum somehow, e.g.

c = np.einsum('ij...,jk...->ik...', a, b)

This will throw an error of course, but hopefully, it is clear what I mean from the examples.

Does einsum support this kind of 'multi-ellipsis' notation? Or is there any other way to implement this without looping?

My guess is that there is no such way because one would have to tell einsum in what order to put the remaining indices, i.e. one would have to label the ellipses somehow.

`np.einsum('ij...,jk...->ik...', a, b)` works fine when the ellipsis is used to represent the trailing axes that are to be aligned. Do you plan to use ellipsis for a different purpose? If you intend to let them "spread-out" as with `np.einsum('ijabc,jkde->ikabcde', a, b)`, no you can't use ellipsis and have to use the explicit mention of indexing string chars. — Divakar, Apr 06 '19 at 13:04
@Divakar thanks for the insight! I would indeed like to them to “spread-out”, as you say. Or in different words: to keep them as extra dimensions. — Wolpertinger, Apr 06 '19 at 13:35

Divakar · Accepted Answer · 2019-04-06T16:34:05.640

Since, there are no axes that are to be aligned, we can simply use tensordot that lets the axes not participating in sum-reductions to be "spread-out" with an additional rollaxis, like so -

np.rollaxis(np.tensordot(a,b,axes=(1,0)),a.ndim-1,1)

If you want to use einsum, we can reshape those to 3D such that the last axis of them are the merged ones (third axis onwards are merged into one) and then go ahead with einsum and finally reshape back to their ndim-1 shapes being spread out in the output, something like this -

shp_a = a.shape
shp_b = b.shape
shp_a[:1] + shp_a[2:]
out_shp = shp_a[:1] + (shp_b[1],) + shp_a[2:] + shp_b[2:]

a3D = a.reshape(shp_a[:2]+(-1,))
b3D = b.reshape(shp_b[:2]+(-1,))
out = np.einsum('ijk,jlm->ilkm',a3D,b3D).reshape(out_shp)

We could also generate the corresponding einsum string notation itself and hence skip all the array manipulation and hence focus on the string manipulation itself to get something like this -

import string

def einsum_spreadout(a,b,a_axes,b_axes,a_spread_axis,b_spread_axis):
    from numpy.core import numerictypes as nt

    if isinstance(a_axes, (int, nt.integer)):
        a_axes = (a_axes,)

    if isinstance(b_axes, (int, nt.integer)):
        b_axes = (b_axes,)

    s = string.ascii_letters

    a_str = s[:a.ndim]
    b_str = s[a.ndim:a.ndim+b.ndim]

    b_str_ar = np.frombuffer(b_str,dtype='S1').copy()
    for (i,j) in zip(a_axes,b_axes):
        b_str_ar[j] = a_str[i]
    b_str = ''.join(b_str_ar)    

    out_str = a_str[:a_spread_axis] + b_str[:b_spread_axis]
    out_str += a_str[a_spread_axis:] + b_str[b_spread_axis:]

    out_str_ar = np.frombuffer(out_str,dtype='S1').copy()
    out_str = ''.join(out_str_ar[~np.isin(out_str_ar,np.take(b_str_ar,b_axes))])
    einsum_str = a_str+','+b_str+'->'+out_str

    return np.einsum(einsum_str,a,b)

Few sample case runs to to showcase its usage -

>>> a = np.random.rand(3,4,6,7,8)
>>> b = np.random.rand(4,5,9,10)
>>> einsum_spreadout(a,b,a_axes=1,b_axes=0,a_spread_axis=2,b_spread_axis=2).shape
(3, 5, 6, 7, 8, 9, 10)

>>> b = np.random.rand(4,5,6,10)
>>> einsum_spreadout(a,b,a_axes=(1,2),b_axes=(0,2),a_spread_axis=2,b_spread_axis=2).shape
(3, 5, 7, 8, 10)

>>> einsum_spreadout(a,b,a_axes=(1,2),b_axes=(0,2),a_spread_axis=4,b_spread_axis=4).shape
(3, 7, 5, 10, 8)

This looks nice already, thanks! I would, however, prefer an einsum implementation (if possible, which I do realize it may not be). The simple matrix contraction in the question was meant to be an example for illustration. For more complicated einsum operations your solution would not generalize. — Wolpertinger, Apr 06 '19 at 14:53
@Wolpertinger And added einsum string notation manipulation one. — Divakar, Apr 06 '19 at 16:34

Ultimate flexibility with einsum ellipses

1 Answers1