`broadcast_like()` doesn't copy chunking structure

[slevang]

What is your issue?

import dask.array
import xarray as xr

da1 = xr.DataArray(dask.array.ones((3,3), chunks=(1, 1)), dims=["x", "y"])
da2 = xr.DataArray(dask.array.ones((3,), chunks=(1,)), dims=["x"])

da2.broadcast_like(da1).chunksizes

Frozen({'x': (1, 1, 1), 'y': (3,)})

Was surprised to not find any other issues around this. Feels like a major limitation of the method for a lot of use cases. Is there an easy hack around this?

[dcherian]

The fundamental issue here is that we use np.broadcast_to(da2.data, da1.shape) (in Variable.set_dims) instead of np.broadcast_arrays(da1.data, da2.data). In the former chunking info from da1 isn't propagated.

This would be a decent lift, involving some refactoring of the core broadcasting logic.

[slevang]

Following up, a workaround that seems effective is this:

da2 * xr.ones_like(da1)

ones_like dispatches to dask.array.full, so we get appropriate chunk structure and a trivial root task (instead of a giant numpy array) if we want to use this on a distributed client.