Repository: dask/dask

Issue:
normalize_token pickles data of literal
While investigating a slow call to `tokenize` I discovered that when a `literal` is passed to a function that tokenizes args, it ends up pickling the quoted data instead of recursively normalising it. This is because there is no normalizer registered for `literal`, so it goes down the generic path for callables, which uses pickle.

Example:
```python
from dask.base import normalize_token
from dask.core import literal

d = {'a': [1, 2], 'b': [3]}
print(normalize_token(d))
print(normalize_token(literal(d)))
```

Output:
```
('list', [('tuple', ['a', ('list', [1, 2])]), ('tuple', ['b', ('list', [3])])])
b'cdask.core\nliteral\np0\n((dp1\nVa\np2\n(lp3\nL1L\naL2L\nasVb\np4\n(lp5\nL3L\nastp6\nRp7\n.'
```

This is on current master.

Discussion:
bmerry: I wasn't sure if the return value should include the class of the literal object, so that `tokenize(literal(x))` and `tokenize(x)` would be different. I'm assuming that if both are used then they're intended to mean the same thing, so I haven't.

Patch:
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 240baa35a..6c30e36e7 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -145,11 +145,15 @@ def blockwise(
         chunkss, arrays = unify_chunks(*args)
     else:
         arginds = [(a, i) for (a, i) in toolz.partition(2, args) if i is not None]
-        if arginds:
-            arg, ind = max(arginds, key=lambda ai: len(ai[1]))
-            chunkss = dict(zip(ind, arg.chunks))
-        else:
-            chunkss = {}
+        chunkss = {}
+        # For each dimension, use the input chunking that has the most blocks;
+        # this will ensure that broadcasting works as expected, and in
+        # particular the number of blocks should be correct if the inputs are
+        # consistent.
+        for arg, ind in arginds:
+            for c, i in zip(arg.chunks, ind):
+                if i not in chunkss or len(c) > len(chunkss[i]):
+                    chunkss[i] = c
         arrays = args[::2]
 
     for k, v in new_axes.items():
@@ -220,6 +224,14 @@ def blockwise(
                 elif isinstance(adjust_chunks[ind], numbers.Integral):
                     chunks[i] = tuple(adjust_chunks[ind] for _ in chunks[i])
                 elif isinstance(adjust_chunks[ind], (tuple, list)):
+                    if len(adjust_chunks[ind]) != len(chunks[i]):
+                        raise ValueError(
+                            "Dimension {0} has {1} blocks, "
+                            "adjust_chunks specified with "
+                            "{2} blocks".format(
+                                i, len(chunks[i]), len(adjust_chunks[ind])
+                            )
+                        )
                     chunks[i] = tuple(adjust_chunks[ind])
                 else:
                     raise NotImplementedError(
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
index a34e1b555..d9e182c2b 100644
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -86,7 +86,7 @@ with ignoring(AttributeError):
     nanstd = np.nanstd
 
 
-def coarsen(reduction, x, axes, trim_excess=False):
+def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
     """ Coarsen array by applying reduction to fixed size neighborhoods
 
     Parameters
@@ -140,7 +140,7 @@ def coarsen(reduction, x, axes, trim_excess=False):
     # (10, 10) -> (5, 2, 5, 2)
     newshape = tuple(concat([(x.shape[i] // axes[i], axes[i]) for i in range(x.ndim)]))
 
-    return reduction(x.reshape(newshape), axis=tuple(range(1, x.ndim * 2, 2)))
+    return reduction(x.reshape(newshape), axis=tuple(range(1, x.ndim * 2, 2)), **kwargs)
 
 
 def trim(x, axes=None):
diff --git a/dask/array/core.py b/dask/array/core.py
index 84fcbf620..ff59dedef 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -612,6 +612,17 @@ def map_blocks(
         out_ind = tuple(out_ind)
         if max(new_axis) > max(out_ind):
             raise ValueError("New_axis values do not fill in all dimensions")
+
+    if chunks is not None:
+        if len(chunks) != len(out_ind):
+            raise ValueError(
+                "Provided chunks have {0} dims, expected {1} "
+                "dims.".format(len(chunks), len(out_ind))
+            )
+        adjust_chunks = dict(zip(out_ind, chunks))
+    else:
+        adjust_chunks = None
+
     out = blockwise(
         func,
         out_ind,
@@ -621,6 +632,7 @@ def map_blocks(
         dtype=dtype,
         concatenate=True,
         align_arrays=False,
+        adjust_chunks=adjust_chunks,
         meta=meta,
         **kwargs,
     )
@@ -639,30 +651,6 @@ def map_blocks(
             v.dsk[key] = task
             dsk[k] = (v,) + vv[1:]
 
-    if chunks is not None:
-        if len(chunks) != len(out.numblocks):
-            raise ValueError(
-                "Provided chunks have {0} dims, expected {1} "
-                "dims.".format(len(chunks), len(out.numblocks))
-            )
-        chunks2 = []
-        for i, (c, nb) in enumerate(zip(chunks, out.numblocks)):
-            if isinstance(c, tuple):
-                # We only check cases where numblocks > 1. Because of
-                # broadcasting, we can't (easily) validate the chunks
-                # when the number of blocks is 1.
-                # See https://github.com/dask/dask/issues/4299 for more.
-                if nb > 1 and len(c) != nb:
-                    raise ValueError(
-                        "Dimension {0} has {1} blocks, "
-                        "chunks specified with "
-                        "{2} blocks".format(i, nb, len(c))
-                    )
-                chunks2.append(c)
-            else:
-                chunks2.append(nb * (c,))
-        out._chunks = normalize_chunks(chunks2)
-
     # If func has block_info as an argument, add it to the kwargs for each call
     if has_keyword(func, "block_info"):
         starts = {}
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 182322658..03ea568df 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -9,6 +9,7 @@ from distutils.version import LooseVersion
 import numpy as np
 from toolz import concat, sliding_window, interleave
 
+from ..compatibility import apply
 from ..core import flatten
 from ..base import tokenize
 from ..highlevelgraph import HighLevelGraph
@@ -1292,7 +1293,7 @@ def piecewise(x, condlist, funclist, *args, **kw):
 
 
 @wraps(chunk.coarsen)
-def coarsen(reduction, x, axes, trim_excess=False):
+def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
     if not trim_excess and not all(
         bd % div == 0 for i, div in axes.items() for bd in x.chunks[i]
     ):
@@ -1304,16 +1305,17 @@ def coarsen(reduction, x, axes, trim_excess=False):
 
     name = "coarsen-" + tokenize(reduction, x, axes, trim_excess)
     dsk = {
-        (name,) + key[1:]: (chunk.coarsen, reduction, key, axes, trim_excess)
+        (name,)
+        + key[1:]: (apply, chunk.coarsen, [reduction, key, axes, trim_excess], kwargs)
         for key in flatten(x.__dask_keys__())
     }
     chunks = tuple(
         tuple(int(bd // axes.get(i, 1)) for bd in bds) for i, bds in enumerate(x.chunks)
     )
 
-    dt = reduction(np.empty((1,) * x.ndim, dtype=x.dtype)).dtype
+    meta = reduction(np.empty((1,) * x.ndim, dtype=x.dtype), **kwargs)
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=[x])
-    return Array(graph, name, chunks, dtype=dt)
+    return Array(graph, name, chunks, meta=meta)
 
 
 def split_at_breaks(array, breaks, axis=0):
diff --git a/dask/base.py b/dask/base.py
index b7c96f728..1612b8408 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -15,7 +15,7 @@ from toolz.functoolz import Compose
 
 from .compatibility import is_dataclass, dataclass_fields
 from .context import thread_state
-from .core import flatten, quote, get as simple_get
+from .core import flatten, quote, get as simple_get, literal
 from .hashing import hash_buffer_hex
 from .utils import Dispatch, ensure_dict, apply
 from . import config, local, threaded
@@ -678,6 +678,11 @@ def normalize_seq(seq):
     return type(seq).__name__, list(map(normalize_token, seq))
 
 
+@normalize_token.register(literal)
+def normalize_literal(lit):
+    return "literal", normalize_token(lit())
+
+
 @normalize_token.register(object)
 def normalize_object(o):
     method = getattr(o, "__dask_tokenize__", None)
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 5fc7cf5d6..28d575a4f 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -581,6 +581,10 @@ def rewrite_blockwise(inputs):
     --------
     optimize_blockwise
     """
+    if len(inputs) == 1:
+        # Fast path: if there's only one input we can just use it as-is.
+        return inputs[0]
+
     inputs = {inp.output: inp for inp in inputs}
     dependencies = {
         inp.output: {d for d, v in inp.indices if v is not None and d in inputs}
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 57989c866..53298c698 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1645,7 +1645,7 @@ Dask Name: {name}, {task} tasks"""
             # Need the astype(int) for empty dataframes, which sum to float dtype
             result = self.reduction(
                 M.count,
-                aggregate=lambda x: x.sum().astype("int64"),
+                aggregate=_count_aggregate,
                 meta=meta,
                 token=token,
                 split_every=split_every,
@@ -5928,6 +5928,10 @@ def idxmaxmin_agg(x, fn=None, skipna=True, scalar=False):
     return res
 
 
+def _count_aggregate(x):
+    return x.sum().astype("int64")
+
+
 def safe_head(df, n):
     r = M.head(df, n)
     if len(r) != n:
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 8c16b8326..1a9ce204f 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -131,7 +131,10 @@ class HighLevelGraph(Mapping):
                     with ignoring(AttributeError):
                         deps[name] |= set(collection.__dask_layers__())
                 else:
-                    key = id(graph)
+                    try:
+                        [key] = collection.__dask_layers__()
+                    except AttributeError:
+                        key = id(graph)
                     layers[key] = graph
                     deps[name].add(key)
                     deps[key] = set()


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 39b2607d1..9b861cd34 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1370,6 +1370,14 @@ def test_map_blocks_with_kwargs():
     assert_eq(result, np.array([4, 9]))
 
 
+def test_map_blocks_infer_chunks_broadcast():
+    dx = da.from_array([[1, 2, 3, 4]], chunks=((1,), (2, 2)))
+    dy = da.from_array([[10, 20], [30, 40]], chunks=((1, 1), (2,)))
+    result = da.map_blocks(lambda x, y: x + y, dx, dy)
+    assert result.chunks == ((1, 1), (2, 2),)
+    assert_eq(result, np.array([[11, 22, 13, 24], [31, 42, 33, 44]]))
+
+
 def test_map_blocks_with_chunks():
     dx = da.ones((5, 3), chunks=(2, 2))
     dy = da.ones((5, 3), chunks=(2, 2))
@@ -2873,6 +2881,9 @@ def test_map_blocks_with_changed_dimension():
     with pytest.raises(ValueError):
         d.map_blocks(lambda b: b.sum(axis=0), chunks=((4, 4, 4),), drop_axis=0)
 
+    with pytest.raises(ValueError):
+        d.map_blocks(lambda b: b.sum(axis=1), chunks=((3, 4),), drop_axis=1)
+
     d = da.from_array(x, chunks=(4, 8))
     e = d.map_blocks(lambda b: b.sum(axis=1), drop_axis=1, dtype=d.dtype)
     assert e.chunks == ((4, 3),)
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
index 1e22396d9..dbac96c26 100644
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -587,3 +587,12 @@ def test_atop_legacy():
     z = da.blockwise(inc, "i", x, "i", dtype=x.dtype)
     assert_eq(y, z)
     assert y.name == z.name
+
+
+def test_non_hlg():
+    # Regression test for https://github.com/dask/dask/issues/5850
+    a = da.from_array(np.ones(1, np.float64), chunks=(1,))
+    a.dask = dict(a.dask)  # Convert from HighLevelGraph to plain dict
+    b = da.from_array(np.zeros(1, np.float64), chunks=(1,))
+    x = a + b
+    assert_eq(x, a)
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 6e16e9d44..a716fabf2 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1391,6 +1391,10 @@ def test_coarsen():
     assert_eq(
         da.chunk.coarsen(np.sum, x, {0: 2, 1: 4}), da.coarsen(da.sum, d, {0: 2, 1: 4})
     )
+    assert_eq(
+        da.chunk.coarsen(np.mean, x, {0: 2, 1: 4}, dtype="float32"),
+        da.coarsen(da.mean, d, {0: 2, 1: 4}, dtype="float32"),
+    )
 
 
 def test_coarsen_with_excess():
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 7d06bb4d3..663a77846 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -25,6 +25,7 @@ from dask.base import (
     named_schedulers,
     get_scheduler,
 )
+from dask.core import literal
 from dask.delayed import Delayed
 from dask.utils import tmpdir, tmpfile, ignoring
 from dask.utils_test import inc, dec
@@ -406,6 +407,10 @@ def test_tokenize_base_types(x):
     assert tokenize(x) == tokenize(x), x
 
 
+def test_tokenize_literal():
+    assert tokenize(literal(["x", 1])) == tokenize(literal(["x", 1]))
+
+
 @pytest.mark.skipif("not np")
 @pytest.mark.filterwarnings("ignore:the matrix:PendingDeprecationWarning")
 def test_tokenize_numpy_matrix():