python - Ellipsis broadcasting in numpy.einsum

Question

I'm having a problem understanding why the following doesn't work:

I have an array prefactor that can be three-dimensional or six-dimensional. I have an array dipoles that has four dimensions. The first three dimensions of dipoles match the last three dimensions of prefactor.

As I don't know the shape of prefactor, I'm using an Ellipsis to account for the three optional dimensions in prefactor:

numpy.einsum('...lmn,lmno->...o', prefactor, dipoles)

(In the example here, prefactor.shape is (1, 1, 1, 160, 160, 128) and dipoles.shape is (160, 160, 128, 3). When executing, I get the error:

operand 1 did not have enough dimensions to match the broadcasting, and couldn't be extended because einstein sum subscripts were specified at both the start and end

It does work, however, when I add an ellipsis to the second term as well:

numpy.einsum('...lmn,...lmno->...o', prefactor, dipoles)

Just that I don't understand why, because there should be no need for an ellipsis there. Does someone know what's going on here?

The same question has been asked at http://comments.gmane.org/gmane.comp.python.numeric.general/53705 but there is no satisfactory answer yet.

Answer 1

There is a github issue for this problem:

https://github.com/numpy/numpy/issues/2455 improvement of index notation in einsum (Trac #1862)

Error case:

einsum('ij...,j->ij...',A,B)

Current work around requires (empty) ellipsis:

einsum('ij...,j...->ij...',A,B)

It looks like einsum loops through the string argument and the ops several times, identifying the indexes, and broadcast types (right, left, middle, none), and op dimensions. With this it constructs an numpy.nditer. It's while constructing op_axes for the nditer that einsum raises this error. I don't know if the test criteria is too tight (ibroadcast >= ndim), or if it needs to take an additional step to construct the right op_axes for this argument.

https://github.com/numpy/numpy/issues/2619 shows how nditer can be used to replicate einsum behavior. Working from this I can replicate your calculation thus:

prefactor = np.random.random((1, 1, 1, 160, 160, 128))
dipoles = np.random.random((160, 160, 128, 3))
x = numpy.einsum('...lmn,...lmno->...o', prefactor, dipoles)
#numpy.einsum('...lmn,lmno->...o', prefactor, dipoles)  # not work

op_axes = [[0,1,2,3,4,5,-1], [-1,-1,-1,0,1,2,3], [0,1,2,-1,-1,-1,3]]
flags = ['reduce_ok','buffered', 'external_loop', 'delay_bufalloc', 'grow_inner']
op_flags = [['readonly']]*nops + [['allocate','readwrite']]
it = np.nditer([prefactor,dipoles,None], flags, op_flags, op_axes=op_axes)
it.operands[nops][...] = 0
it.reset()
#it.debug_print()
for (x,y,w) in it:
    w[...] += x*y
print "
nditer usage:"
print it.operands[nops] # == x
print it.operands[nops].shape # (1, 1, 1, 3)

The op_axes line is indicative of what einsum deduces from '...lmn,...lmno->...o'.

---

I am exploring this issue on https://github.com/hpaulj/numpy-einsum.

There I have a einsum_py.py which emulates np.einsum with Python code. The part that is relevant to this issue is parse_subscripts(), and in particular prepare_op_axes(). It appears that only the BROADCAST_RIGHT iteration (starting from the end) is needed to correctly create op_axes, regardless of where ellipses are in the subscripts. It also removes the error message that is at the core of this issue.

The einsum.c.src file on that repository has this change, and compiles correctly with the current master distribution (just replace the file and build). It tests fine against test_einsum.py, as well as examples from this issue.

I've submitted a pull request for this change.