dask einsum failing within dask blockwise due to Ellipsis TypeError

31 Views Asked by At

I'm not able to understand why the first one (ex1) of these two examples of code works, and the second one (ex2) throws an error:

A = da.rechunk(da.ones([3,10,100]), chunks=(1,10,100))
B = da.rechunk(da.ones([3,10,100]), chunks=(1,10,100))
J = da.ones(10)
K = da.ones(100)
ex1 = da.einsum('ajk,     bjk,    j,    k -> ab',
                 A,       B,      J,     K,
                          optimize = 'greedy')
ex1.compute()

gives:

array([[1000., 1000., 1000.],
       [1000., 1000., 1000.],
       [1000., 1000., 1000.]])

whereas

ex2 = da.blockwise(da.einsum, 'ab',
             A, 'ajk',
            B, 'bjk',
            J,'j',
            K,'k',
            dtype=A.dtype,
            optimize='greedy')
ex2.compute()

gives

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-207-2906c4cf36cd> in <module>()
----> 1 ex2.compute()

/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/dask/array/einsumfuncs.py in parse_einsum_input()
     83                 else:
     84                     raise TypeError(
---> 85                         "For this input type lists must contain "
     86                         "either int or Ellipsis"
     87                     )

TypeError: For this input type lists must contain either int or Ellipsis
0

There are 0 best solutions below