Make dask names change when chunking Variables by different amounts.

[dcherian]

When rechunking by the current chunk size, name should not change.
Add a __dask_tokenize__ method for ReprObject so that this behaviour is present
when DataArrays are converted to temporary Datasets and back.

• Closes assert_equal and dask #3350
• Tests added
• Passes black . && mypy . && flake8
• Fully documented, including whats-new.rst for all changes and api.rst for new API

[dcherian]

The tests fail on dask == 2.8.1 with this interesting bug. Here's a reproducible example.

import dask
import xarray as xr

ds = xr.Dataset({'x': (('y',), dask.array.ones(10, chunks=(3,)))})
mapped = ds.map_blocks(lambda x: x)
mapped.compute()  # works

xr.testing.assert_equal(mapped, ds)  # does not work
xr.testing.assert_equal(mapped, ds.compute()) # works
xr.testing.assert_equal(mapped.compute(), ds)  # works
xr.testing.assert_equal(mapped.compute(), ds.compute())  # works

The traceback is

~/miniconda3/envs/dcpy/lib/python3.7/site-packages/dask/array/optimization.py in optimize(dsk, keys, fuse_keys, fast_functions, inline_functions_fast_functions, rename_fused_keys, **kwargs)
     41     if isinstance(dsk, HighLevelGraph):
     42         dsk = optimize_blockwise(dsk, keys=keys)
---> 43         dsk = fuse_roots(dsk, keys=keys)
     44 
     45     # Low level task optimizations

~/miniconda3/envs/dcpy/lib/python3.7/site-packages/dask/blockwise.py in fuse_roots(graph, keys)
    819             isinstance(layer, Blockwise)
    820             and len(deps) > 1
--> 821             and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
    822             and all(len(dependents[dep]) == 1 for dep in deps)
    823         ):

~/miniconda3/envs/dcpy/lib/python3.7/site-packages/dask/blockwise.py in <genexpr>(.0)
    819             isinstance(layer, Blockwise)
    820             and len(deps) > 1
--> 821             and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
    822             and all(len(dependents[dep]) == 1 for dep in deps)
    823         ):

KeyError: 'lambda-6720ab0e3639d5c63fc06dfc66a3ce47-x'

This key is not in dependencies. From https://github.com/dask/dask/blob/67fb5363009c583c175cb577776a4f2f4e811410/dask/blockwise.py#L816-L826

    for name, layer in graph.layers.items():
        deps = graph.dependencies[name]
        if (
            isinstance(layer, Blockwise)
            and len(deps) > 1
            and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
            and all(len(dependents[dep]) == 1 for dep in deps)
        ):
            new = toolz.merge(layer, *[layers[dep] for dep in deps])
            new, _ = fuse(new, keys, ave_width=len(deps))

I'm not sure whether this is a bug in fuse_roots, HighLevelGraph.from_collections or in how map_blocks calls HighLevelGraph.from_collections here:

xarray/xarray/core/parallel.py

Line 315 in 69c85b8

graph = HighLevelGraph.from_collections(gname, graph, dependencies=[dataset])

cc @mrocklin

[TomAugspurger]

So this is enough to fix this in Dask

diff --git a/dask/blockwise.py b/dask/blockwise.py
index 52a36c246..84e0ecc08 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -818,7 +818,7 @@ def fuse_roots(graph: HighLevelGraph, keys: list):
         if (
             isinstance(layer, Blockwise)
             and len(deps) > 1
-            and not any(dependencies[dep] for dep in deps)  # no need to fuse if 0 or 1
+            and not any(dependencies.get(dep, {}) for dep in deps)  # no need to fuse if 0 or 1
             and all(len(dependents[dep]) == 1 for dep in deps)
         ):
             new = toolz.merge(layer, *[layers[dep] for dep in deps])

I'm trying to understand why we're getting this KeyError though. I want to make sure that we have a valid HighLevelGraph before making that change.

[TomAugspurger]

@mrocklin if you get a chance, can you confirm that the values in HighLevelGraph.depedencies should be a subset of the keys of layers?

So in the following, the lambda-<...>-x is problematic, because it's not a key in layers?

(Pdb) pp list(self.layers)
['eq-e98e52fb2b8e27b4b5158d399330c72d',
 'lambda-0f1d0bc5e7df462d7125839aed006e04',
 'ones-c4a83f4b990021618d55e0fa61a351d6']
(Pdb) pp self.dependencies
{'eq-e98e52fb2b8e27b4b5158d399330c72d': {'lambda-0f1d0bc5e7df462d7125839aed006e04-x',
                                         'ones-c4a83f4b990021618d55e0fa61a351d6'},
 'lambda-0f1d0bc5e7df462d7125839aed006e04': {'ones-c4a83f4b990021618d55e0fa61a351d6'},
 'ones-c4a83f4b990021618d55e0fa61a351d6': set()}

That's coming from the name of the DataArray / the dask arary in DataArray.data.

[mrocklin]

@mrocklin if you get a chance, can you confirm that the values in HighLevelGraph.depedencies should be a subset of the keys of layers?

That sounds like a reasonable expectation, but honestly it's been a while, so I don't fully trust my knowledge here. It might be worth adding some runtime checks into the HighLevelGraph constructor to see where this might be occurring.

[dcherian]

gentle ping @crusaderky

[dcherian]

Thanks @crusaderky