[AstroPy] Issue with dask and astropy.units (including in astropy 4.3 and 5.0 development versions)

[Livesey, Nathaniel J (US 3290)]

Hi there,

   Thanks so much for astropy.units, I use it heavily (though I’m not an astronomer).  I’ve been having trouble using units with dask, but gather from the astropy website, and other searches, that that’s been addressed in version 4.3.  However, I fear some may still remain.  The MWE pasted at the bottom of this Email demonstrates the problem.

   With astropy version 4.3rc2.dev10+ga094d524a, and dask version 2021.04.0, I get multiple complaints of:

Exception: UnitConversionError("Can only apply 'add' function to quantities with compatible dimensions”)

 I also checked out the latest version in github (5.0.dev382+ga97b8331c) and had the same result.  When I set “parallel" to “False” in the code below, it works fine.  Thanks for any pointers (or a recommendation to wait for a future version will be fine), and apologies if I’m making an elementary mistake. 

Yours,

Nathaniel

Nathaniel Livesey
Mail Stop 183-701, Jet Propulsion Laboratory
4800 Oak Grove Drive, Pasadena, California 91109.
Phone: +1 818 354-4214   Fax: +1 818 393-5065


#####################################

from dask.distributed import Client
import astropy
import astropy.units as units
import dask
import numpy as np

my_unit = units.def_unit("me")


def do_stuff(p):
    return p + 1.0 * my_unit


parallel = True
if __name__ == "__main__" and parallel:
    dask_client = Client()

print(astropy.__version__)
print(dask.__version__)

if parallel:
    routine = dask.delayed(do_stuff)
else:
    routine = do_stuff

p = np.arange(5) << my_unit
result = []
for p_i in p:
    result.append(routine(p_i))

if parallel:
    result = dask.compute(*result)

print(units.Quantity(result))