Ellipsis broadcasting in numpy.einsum

前端 未结 1 1855
灰色年华
灰色年华 2020-12-21 03:19

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 a

相关标签:
1条回答
  • 2020-12-21 04:05

    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 "\nnditer 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.

    0 讨论(0)
提交回复
热议问题