Repository: dask/dask
Pull request #11568 (2024-11-28 12:18:42)

Blockwise uses `Task` class
This migrates `Blockwise` to the Task class (xref https://github.com/dask/dask/issues/9969)

fjetter: Tests all pass. Next up is perf testing. Review can start but I recommend starting with https://github.com/dask/dask/pull/11569 first since this one builds on that
fjetter: In a follow up PR I will remove the `SubgraphCallable` object. It's only use is in the `fuse_subgraph` code path of the fusion code that is effectively unused now (and off by default)

Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index d239f8393..c86c7252b 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -110,7 +110,7 @@ def parse_input(obj: Any) -> object:
             return _wrap_namedtuple_task(None, obj, parse_input)
         return Tuple(*(parse_input(o) for o in obj))
     if isinstance(obj, dict):
-        return Dict(**{k: parse_input(v) for k, v in obj.items()})  # type: ignore
+        return Dict({k: parse_input(v) for k, v in obj.items()})
 
     if isinstance(obj, TaskRef):
         return Alias(obj.key)
@@ -151,18 +151,6 @@ class _MultiContainer(Container):
         return False
 
 
-def _to_dict(*args: Any) -> dict:
-    vals = args[:-1]
-    keys = args[-1]
-    return {key: v for key, v in zip(keys, vals)}
-
-
-def _wrap_dict_in_task(d: dict) -> Task:
-    """Wrap a dictionary in a task such that all values will be properly
-    recognized as runnable tasks"""
-    return Task(None, _to_dict, *d.values(), d.keys())
-
-
 SubgraphType = None
 
 
@@ -182,6 +170,9 @@ def convert_legacy_task(
     task: _T,
     all_keys: Container,
 ) -> GraphNode | _T:
+    if isinstance(task, GraphNode):
+        return task
+
     global SubgraphType
 
     if SubgraphType is None:
@@ -189,9 +180,6 @@ def convert_legacy_task(
 
         SubgraphType = SubgraphCallable
 
-    if isinstance(task, GraphNode):
-        return task
-
     if type(task) is tuple and task and callable(task[0]):
         func, args = task[0], task[1:]
         if isinstance(func, SubgraphType):
@@ -301,11 +289,10 @@ def convert_legacy_graph(
         t = convert_legacy_task(k, arg, all_keys)
         if isinstance(t, Alias) and t.target.key == k:
             continue
+        elif not isinstance(t, GraphNode):
+            t = DataNode(k, t)
         new_dsk[k] = t
-    new_dsk2 = {
-        k: v if isinstance(v, GraphNode) else DataNode(k, v) for k, v in new_dsk.items()
-    }
-    return new_dsk2
+    return new_dsk
 
 
 def resolve_aliases(dsk: dict, keys: set, dependents: dict) -> dict:
@@ -419,6 +406,7 @@ class GraphNode:
     def __eq__(self, value: object) -> bool:
         if type(value) is not type(self):
             return False
+
         from dask.tokenize import tokenize
 
         return tokenize(self) == tokenize(value)
@@ -432,12 +420,80 @@ class GraphNode:
             sizeof(getattr(self, sl)) for sl in type(self).__slots__
         ) + sys.getsizeof(type(self))
 
+    def substitute(
+        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
+    ) -> GraphNode:
+        """Substitute a dependency with a new value. The new value either has to
+        be a new valid key or a GraphNode to replace the dependency entirely.
+
+        The GraphNode will not be mutated but instead a shallow copy will be
+        returned. The substitution will be performed eagerly.
+
+        Parameters
+        ----------
+        subs : dict[KeyType, KeyType  |  GraphNode]
+            The mapping describing the substitutions to be made.
+        key : KeyType | None, optional
+            The key of the new GraphNode object. If None provided, the key of
+            the old one will be reused.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def fuse(*tasks: GraphNode, key: KeyType | None = None) -> GraphNode:
+        """Fuse a set of tasks into a single task.
+
+        The tasks are fused into a single task that will execute the tasks in a
+        subgraph. The internal tasks are no longer accessible from the outside.
+
+        All provided tasks must form a valid subgraph that will reduce to a
+        single key. If multiple outputs are possible with the provided tasks, an
+        exception will be raised.
+
+        The tasks will not be rewritten but instead a new Task will be created
+        that will merely reference the old task objects. This way, Task objects
+        may be reused in multiple fused tasks.
+
+        Parameters
+        ----------
+        key : KeyType | None, optional
+            The key of the new Task object. If None provided, the key of the
+            final task will be used.
+
+        See also
+        --------
+        GraphNode.substitute : Easer substitution of dependencies
+        """
+        if any(t.key is None for t in tasks):
+            raise ValueError("Cannot fuse tasks with missing keys")
+        if len(tasks) == 1:
+            return tasks[0].substitute({}, key=key)
+        all_keys = set()
+        all_deps: set[KeyType] = set()
+        for t in tasks:
+            all_deps.update(t.dependencies)
+            all_keys.add(t.key)
+        external_deps = all_deps - all_keys
+        leafs = all_keys - all_deps
+        if len(leafs) > 1:
+            raise ValueError(f"Cannot fuse tasks with multiple outputs {leafs}")
+
+        outkey = leafs.pop()
+
+        return Task(
+            key or outkey,
+            _execute_subgraph,
+            {t.key: t for t in tasks},
+            outkey,
+            (Dict({k: Alias(k) for k in external_deps}) if external_deps else {}),
+            {},
+        )
+
 
 _no_deps: frozenset = frozenset()
 
 
 class Alias(GraphNode):
-    __weakref__: Any = None
     target: TaskRef
     __slots__ = tuple(__annotations__)
 
@@ -455,12 +511,35 @@ class Alias(GraphNode):
     def copy(self):
         return Alias(self.key, self.target)
 
+    def substitute(
+        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
+    ) -> GraphNode:
+        if self.key in subs or self.target.key in subs:
+            sub_key = subs.get(self.key, self.key)
+            val = subs.get(self.target.key, self.target.key)
+            if sub_key == self.key and val == self.target.key:
+                return self
+            if isinstance(val, GraphNode):
+                return val.substitute({}, key=key)
+            if key is None and isinstance(sub_key, GraphNode):
+                raise RuntimeError(
+                    f"Invalid substitution encountered {self.key!r} -> {sub_key}"
+                )
+            return Alias(key or sub_key, val)  # type: ignore
+        return self
+
+    def __dask_tokenize__(self):
+        return (type(self).__name__, self.key, self.target.key)
+
     def __call__(self, values=()):
         self._verify_values(values)
         return values[self.target.key]
 
     def __repr__(self):
-        return f"Alias({self.key}->{self.target})"
+        if self.key != self.target.key:
+            return f"Alias({self.key!r}->{self.target.key!r})"
+        else:
+            return f"Alias({self.key!r})"
 
     def __eq__(self, value: object) -> bool:
         if not isinstance(value, Alias):
@@ -492,7 +571,7 @@ class DataNode(GraphNode):
         return self.value
 
     def __repr__(self):
-        return f"DataNode({self.key}, type={self.typ}, {self.value})"
+        return f"DataNode({self.value!r})"
 
     def __reduce__(self):
         return (DataNode, (self.key, self.value))
@@ -502,6 +581,13 @@ class DataNode(GraphNode):
 
         return (type(self).__name__, tokenize(self.value))
 
+    def substitute(
+        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
+    ) -> DataNode:
+        if key is not None and key != self.key:
+            return DataNode(key, self.value)
+        return self
+
     def __iter__(self):
         return iter(self.value)
 
@@ -551,6 +637,8 @@ class Task(GraphNode):
     ):
         self.key = key
         self.func = func
+        if isinstance(func, Task):
+            raise TypeError("Cannot nest tasks")
         self.args = tuple(
             Alias(obj.key) if isinstance(obj, TaskRef) else obj for obj in args
         )
@@ -636,11 +724,16 @@ class Task(GraphNode):
     def __call__(self, values=()):
         self._verify_values(values)
         new_argspec = tuple(
-            a(values) if isinstance(a, GraphNode) else a for a in self.args
+            a({k: values[k] for k in a.dependencies}) if isinstance(a, GraphNode) else a
+            for a in self.args
         )
         if self.kwargs:
             kwargs = {
-                k: kw(values) if isinstance(kw, GraphNode) else kw
+                k: (
+                    kw({k: values[k] for k in kw.dependencies})
+                    if isinstance(kw, GraphNode)
+                    else kw
+                )
                 for k, kw in self.kwargs.items()
             }
             return self.func(*new_argspec, **kwargs)
@@ -658,28 +751,37 @@ class Task(GraphNode):
                 self._is_coro = False
         return self._is_coro
 
-    @staticmethod
-    def fuse(*tasks: Task, key: KeyType | None = None) -> Task:
-        all_keys = set()
-        all_deps: set[KeyType] = set()
-        for t in tasks:
-            all_deps.update(t.dependencies)
-            all_keys.add(t.key)
-        external_deps = all_deps - all_keys
-        leafs = all_keys - all_deps
-        if len(leafs) > 1:
-            raise ValueError("Cannot fuse tasks with multiple outputs")
-
-        outkey = leafs.pop()
-
-        return Task(
-            key or outkey,
-            _execute_subgraph,
-            {t.key: t for t in tasks},
-            outkey,
-            Dict({k: Alias(k) for k in external_deps}),
-            {},
-        )
+    def substitute(
+        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
+    ) -> Task:
+        subs_filtered = {
+            k: v for k, v in subs.items() if k in self.dependencies and k != v
+        }
+        if subs_filtered:
+            new_args = tuple(
+                a.substitute(subs_filtered) if isinstance(a, GraphNode) else a
+                for a in self.args
+            )
+            new_kwargs = {
+                k: v.substitute(subs_filtered) if isinstance(v, GraphNode) else v
+                for k, v in self.kwargs.items()
+            }
+            return type(self)(
+                key or self.key,
+                self.func,
+                *new_args,
+                **new_kwargs,
+            )
+        elif key is None or key == self.key:
+            return self
+        else:
+            # Rename
+            return type(self)(
+                key,
+                self.func,
+                *self.args,
+                **self.kwargs,
+            )
 
 
 class NestedContainer(Task):
@@ -703,6 +805,35 @@ class NestedContainer(Task):
             **kwargs,
         )
 
+    def __repr__(self):
+        return f"{type(self).__name__}({self.args})"
+
+    def substitute(
+        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
+    ) -> NestedContainer:
+        subs_filtered = {
+            k: v for k, v in subs.items() if k in self.dependencies and k != v
+        }
+        if not subs_filtered:
+            return self
+        return type(self)(
+            *(
+                a.substitute(subs_filtered) if isinstance(a, GraphNode) else a
+                for a in self.args
+            )
+        )
+
+    def __dask_tokenize__(self):
+        from dask.tokenize import tokenize
+
+        return (
+            type(self).__name__,
+            self.klass,
+            sorted(tokenize(a) for a in self.args),
+        )
+
+        return super().__dask_tokenize__()
+
     @classmethod
     def to_container(cls, *args, **kwargs):
         return cls.klass(args)
@@ -734,6 +865,25 @@ class Dict(NestedContainer):
             *(Tuple(*it) for it in kwargs.items()), _dependencies=_dependencies
         )
 
+    def __repr__(self):
+        values = ", ".join(f"{tup.args[0]}: {tup.args[1]}" for tup in self.args)
+        return f"Dict({values})"
+
+    def substitute(
+        self, subs: dict[KeyType, KeyType | GraphNode], key: KeyType | None = None
+    ) -> Dict:
+        subs_filtered = {
+            k: v for k, v in subs.items() if k in self.dependencies and k != v
+        }
+        if not subs_filtered:
+            return self
+
+        new = {}
+        for tup in self.args:
+            k, v = tup.substitute(subs_filtered).args
+            new[k] = v
+        return type(self)(new)
+
 
 class DependenciesMapping(MutableMapping):
     def __init__(self, dsk):
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 0bfe4412b..ff4fd0c0d 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -6,8 +6,8 @@ import warnings
 import tlz as toolz
 
 from dask import base, utils
+from dask.blockwise import _blockwise_unpack_collections_task_spec
 from dask.blockwise import blockwise as core_blockwise
-from dask.delayed import unpack_collections
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import ArrayBlockwiseDep
 
@@ -201,7 +201,7 @@ def blockwise(
             v = (v,)
         chunkss[k] = v
 
-    arginds = zip(arrays, args[1::2])
+    arginds = list(zip(arrays, args[1::2]))
     numblocks = {}
 
     dependencies = []
@@ -213,7 +213,8 @@ def blockwise(
     for arg, ind in arginds:
         if ind is None:
             arg = normalize_arg(arg)
-            arg, collections = unpack_collections(arg)
+            arg, collections = _blockwise_unpack_collections_task_spec(arg)
+
             dependencies.extend(collections)
         else:
             if (
@@ -235,7 +236,7 @@ def blockwise(
     kwargs2 = {}
     for k, v in kwargs.items():
         v = normalize_arg(v)
-        v, collections = unpack_collections(v)
+        v, collections = _blockwise_unpack_collections_task_spec(v)
         dependencies.extend(collections)
         kwargs2[k] = v
 
diff --git a/dask/array/core.py b/dask/array/core.py
index 34c559c9f..92229a9fb 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -28,6 +28,7 @@ from tlz.curried import pluck
 from toolz import frequencies
 
 from dask import compute, config, core
+from dask._task_spec import TaskRef
 from dask.array import chunk
 from dask.array.chunk import getitem
 from dask.array.chunk_types import is_valid_array_chunk, is_valid_chunk_type
@@ -317,7 +318,7 @@ def graph_from_arraylike(
             getitem,
             name,
             out_ind,
-            original_name,
+            TaskRef(original_name),
             None,
             ArraySliceDep(chunks),
             out_ind,
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index d3bd3a9d9..b947bd89f 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -142,8 +142,8 @@ def slice_array(out_name, in_name, blockdims, index, itemsize):
     >>> dsk, blockdims = slice_array('y', 'x', [(20, 20, 20, 20, 20)],
     ...                              (slice(10, 35),), 8)
     >>> pprint(dsk)  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
-    {('y', 0): <Task ('y', 0) getitem(Alias(('x', 0)->('x', 0)), (slice(10, 20, 1),))>,
-     ('y', 1): <Task ('y', 1) getitem(Alias(('x', 1)->('x', 1)), (slice(0, 15, 1),))>}
+    {('y', 0): <Task ('y', 0) getitem(Alias(('x', 0)), (slice(10, 20, 1),))>,
+     ('y', 1): <Task ('y', 1) getitem(Alias(('x', 1)), (slice(0, 15, 1),))>}
     >>> blockdims
     ((10, 15),)
 
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 4bd31a258..214599506 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -42,6 +42,7 @@ from dask._task_spec import (
     GraphNode,
     List,
     Task,
+    TaskRef,
     _execute_subgraph,
     convert_legacy_graph,
     cull,
@@ -56,10 +57,10 @@ from dask.base import (
     replace_name_in_key,
     tokenize,
 )
-from dask.blockwise import blockwise
+from dask.blockwise import _blockwise_unpack_collections_task_spec, blockwise
 from dask.context import globalmethod
 from dask.core import flatten, istask, quote
-from dask.delayed import Delayed, unpack_collections
+from dask.delayed import Delayed
 from dask.highlevelgraph import HighLevelGraph
 from dask.sizeof import sizeof
 from dask.typing import Graph, NestedKeys, no_default
@@ -101,6 +102,7 @@ def lazify_task(task, start=True):
     >>> lazify_task(task)  # doctest: +ELLIPSIS
     (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3]))
     """
+
     if isinstance(task, GraphNode):
         if isinstance(task, List) and len(task.args) < 50:
             return List(*[lazify_task(arg, False) for arg in task.args])
@@ -2099,14 +2101,12 @@ def unpack_scalar_dask_kwargs(kwargs):
     kwargs2 = {}
     dependencies = []
     for k, v in kwargs.items():
-        vv, collections = unpack_collections(v)
+        vv, collections = _blockwise_unpack_collections_task_spec(v)
         if not collections:
             kwargs2[k] = v
         else:
             kwargs2[k] = vv
             dependencies.extend(collections)
-    if dependencies:
-        kwargs2 = (dict, (zip, list(kwargs2), list(kwargs2.values())))
     return kwargs2, dependencies
 
 
@@ -2273,11 +2273,10 @@ def map_partitions(func, *args, **kwargs):
         if isinstance(a, Bag):
             bags.append(a)
             args2.append(a)
-        elif isinstance(a, (Item, Delayed)):
-            args2.append(a.key)
-            dependencies.append(a)
         else:
+            a, collections = _blockwise_unpack_collections_task_spec(a)
             args2.append(a)
+            dependencies.extend(collections)
 
     bag_kwargs = {}
     other_kwargs = {}
@@ -2299,26 +2298,29 @@ def map_partitions(func, *args, **kwargs):
         raise ValueError("All bags must have the same number of partitions.")
     npartitions = npartitions.pop()
 
-    def build_args(n):
-        return [(a.name, n) if isinstance(a, Bag) else a for a in args2]
+    if bag_kwargs:
 
-    def build_bag_kwargs(n):
-        if not bag_kwargs:
-            return {}
-        return (
-            dict,
-            (zip, list(bag_kwargs), [(b.name, n) for b in bag_kwargs.values()]),
-        )
+        def build_args(n):
+            return [TaskRef((a.name, n)) if isinstance(a, Bag) else a for a in args2]
+
+        def build_bag_kwargs(n) -> dict:
+            if not bag_kwargs:
+                return other_kwargs
+            rv = {
+                k: TaskRef((b.name, n)) if isinstance(b, Bag) else b
+                for k, b in bag_kwargs.items()
+            }
+            rv.update(other_kwargs)
+            return rv
 
-    if bag_kwargs:
         # Avoid using `blockwise` when a key-word
         # argument is being used to refer to a collection.
         dsk = {
-            (name, n): (
-                apply,
+            (name, n): Task(
+                (name, n),
                 func,
-                build_args(n),
-                (merge, build_bag_kwargs(n), other_kwargs),
+                *build_args(n),
+                **build_bag_kwargs(n),
             )
             for n in range(npartitions)
         }
@@ -2333,12 +2335,6 @@ def map_partitions(func, *args, **kwargs):
                 numblocks[arg.name] = (arg.npartitions,)
             else:
                 pairs.extend([arg, None])
-        if other_kwargs and isinstance(other_kwargs, tuple):
-            # `other_kwargs` is a nested subgraph,
-            # designed to generate the kwargs lazily.
-            # We need to convert this to a dictionary
-            # before passing to `blockwise`
-            other_kwargs = other_kwargs[0](other_kwargs[1][0](*other_kwargs[1][1:]))
         dsk = blockwise(
             func,
             name,
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 141eb3db0..9f825a059 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import itertools
 import os
 from collections.abc import Hashable, Iterable, Mapping, Sequence
+from dataclasses import fields, is_dataclass, replace
 from itertools import product
 from math import prod
 from typing import Any
@@ -10,12 +11,15 @@ from typing import Any
 import tlz as toolz
 
 import dask
+from dask import base, utils
+from dask._task_spec import Dict, GraphNode, List, Task, TaskRef, parse_input
 from dask.base import clone_key, get_name_from_key, tokenize
 from dask.core import flatten, ishashable, keys_in_tasks, reverse_dict
+from dask.delayed import Delayed, finalize
 from dask.highlevelgraph import HighLevelGraph, Layer
-from dask.optimization import SubgraphCallable, fuse
-from dask.typing import Graph, Key
-from dask.utils import _deprecated, apply, ensure_dict, homogeneous_deepmap
+from dask.optimization import fuse
+from dask.typing import Key
+from dask.utils import _deprecated, ensure_dict, homogeneous_deepmap
 
 
 class BlockwiseDep:
@@ -122,7 +126,7 @@ class BlockwiseDepDict(BlockwiseDep):
     >>> layer = Blockwise(
     ...     output="collection-name",
     ...     output_indices="i",
-    ...     dsk={"collection-name": (func, '_0')},
+    ...     task=Task("collection-name", func, TaskRef("_0")),
     ...     indices=[(dep, "i")],
     ...     numblocks={},
     ... )
@@ -236,17 +240,172 @@ def blockwise(
 ):
     """Create a Blockwise symbolic mutable mapping
 
-    This is like the ``make_blockwise_graph`` function, but rather than construct a
-    dict, it returns a symbolic Blockwise object.
+    Applies a function, ``func``, across blocks from many different input
+    collections.  We arrange the pattern with which those blocks interact with
+    sets of matching indices.  E.g.::
 
-    ``*arrind_pairs`` is similar to those in `make_blockwise_graph`, but in addition to
-    allowing for collections it can accept BlockwiseDep instances, which allows for lazy
-    evaluation of arguments to ``func`` which might be different for different
-    chunks/partitions.
+        blockwise(func, 'z', 'i', 'x', 'i', 'y', 'i')
+
+    yield an embarrassingly parallel communication pattern and is read as
+
+        $$ z_i = func(x_i, y_i) $$
+
+    More complex patterns may emerge, including multiple indices::
+
+        blockwise(func, 'z', 'ij', 'x', 'ij', 'y', 'ji')
+
+        $$ z_{ij} = func(x_{ij}, y_{ji}) $$
+
+    Indices missing in the output but present in the inputs results in many
+    inputs being sent to one function (see examples).
+
+    Examples
+    --------
+    Simple embarrassing map operation
+
+    >>> def inc(x):
+    ...    return x + 1
+    >>> dict(
+    ...     blockwise(
+    ...         inc,
+    ...         'z', 'ij',
+    ...         'x', 'ij',
+    ...         numblocks={'x': (2, 2)}
+    ...     )
+    ... )  # doctest: +NORMALIZE_WHITESPACE
+    {('z', 0, 0): <Task ('z', 0, 0) inc(Alias(('x', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) inc(Alias(('x', 0, 1)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) inc(Alias(('x', 1, 0)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) inc(Alias(('x', 1, 1)))>}
+
+    Simple operation on two datasets x and y
+
+    >>> def add(x, y):
+    ...     return x + y
+    >>> dict(
+    ...     blockwise(
+    ...         add,
+    ...         'z', 'ij',
+    ...         'x', 'ij',
+    ...         'y', 'ij',
+    ...         numblocks={'x': (2, 2), 'y': (2, 2)}
+    ...     )
+    ... )  # doctest: +NORMALIZE_WHITESPACE
+    {('z', 0, 0): <Task ('z', 0, 0) add(Alias(('x', 0, 0)), Alias(('y', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) add(Alias(('x', 0, 1)), Alias(('y', 0, 1)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) add(Alias(('x', 1, 0)), Alias(('y', 1, 0)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) add(Alias(('x', 1, 1)), Alias(('y', 1, 1)))>}
+
+    Operation that flips one of the datasets
+
+    >>> def addT(x, y):
+    ...     return x + y.T
+    >>> dict(
+    ...     blockwise(
+    ...         addT,
+    ...         'z', 'ij',
+    ...         'x', 'ij',
+    ...         'y', 'ji',
+    ...         numblocks={'x': (2, 2), 'y': (2, 2)}
+    ...     )
+    ... )  # doctest: +NORMALIZE_WHITESPACE
+    {('z', 0, 0): <Task ('z', 0, 0) addT(Alias(('x', 0, 0)), Alias(('y', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) addT(Alias(('x', 0, 1)), Alias(('y', 1, 0)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) addT(Alias(('x', 1, 0)), Alias(('y', 0, 1)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) addT(Alias(('x', 1, 1)), Alias(('y', 1, 1)))>}
+
+    Dot product with contraction over ``j`` index.  Yields list arguments
+
+    >>> from dask.array.core import dotmany
+    >>> dict(
+    ...     blockwise(
+    ...         dotmany,
+    ...         'z', 'ik',
+    ...         'x', 'ij',
+    ...         'y', 'jk',
+    ...         numblocks={'x': (2, 2), 'y': (2, 2)}
+    ...     )
+    ... )  # doctest: +SKIP
+    {('z', 0,  0):
+        <Task ('z', 0, 0) dotmany(
+        List((Alias(('x', 0, 0)), Alias(('x', 0, 1)))),
+        List((Alias(('y', 0, 0)), Alias(('y', 1, 0)))))
+        >,
+    ('z', 0,  1):
+        <Task ('z', 0, 1) dotmany(
+        List((Alias(('x', 0, 0)), Alias(('x', 0, 1)))),
+        List((Alias(('y', 0, 1)), Alias(('y', 1, 1)))))
+        >,
+    ('z', 1,  0):
+        <Task ('z', 1, 0) dotmany(
+        List((Alias(('x', 1, 0)), Alias(('x', 1, 1)))),
+        List((Alias(('y', 0, 0)), Alias(('y', 1, 0)))))
+        >,
+    ('z', 1,  1):
+        <Task ('z', 1, 1) dotmany(
+        List((Alias(('x', 1, 0)), Alias(('x', 1, 1)))),
+        List((Alias(('y', 0, 1)), Alias(('y', 1, 1)))))
+        >
+    }
+
+    Pass ``concatenate=True`` to concatenate arrays ahead of time.
+    ``concatenate_axes`` is then called on the lists before the arguments are
+    passed to the function.
+
+    Supports Broadcasting rules
+
+    >>> dict(
+    ...     blockwise(
+    ...         add,
+    ...         'z', 'ij',
+    ...         'x', 'ij',
+    ...         'y', 'ij',
+    ...         concatenate=True,
+    ...         numblocks={'x': (1, 2), 'y': (2, 2)}
+    ...     )
+    ... )  # doctest: +SKIP
+    {
+    ('z', 0, 0): <Task ('z', 0, 0) add(Alias(('x', 0, 0)), Alias(('y', 0, 0)))>,
+    ('z', 0, 1): <Task ('z', 0, 1) add(Alias(('x', 0, 1)), Alias(('y', 0, 1)))>,
+    ('z', 1, 0): <Task ('z', 1, 0) add(Alias(('x', 0, 0)), Alias(('y', 1, 0)))>,
+    ('z', 1, 1): <Task ('z', 1, 1) add(Alias(('x', 0, 1)), Alias(('y', 1, 1)))>
+    }
+
+    Include literals by indexing with ``None``
+
+    >>> dict(
+    ...     blockwise(
+    ...         add,
+    ...         'z', 'i',
+    ...         'x', 'i',
+    ...         100, None,
+    ...         numblocks={'x': (2, )}
+    ...     )
+    ... )  # doctest: +NORMALIZE_WHITESPACE
+    {('z', 0): <Task ('z', 0) add(Alias(('x', 0)), DataNode(100))>,
+    ('z', 1): <Task ('z', 1) add(Alias(('x', 1)), DataNode(100))>}
+
+    If the broadcasted value is a delayed object or other dask collection, the
+    key has to be wrapped appropriately as an Alias object.
+
+    >>> from dask import delayed
+    >>> from dask._task_spec import Alias
+    >>> delayed_inc = delayed(inc)(5, dask_key_name='dinc')
+    >>> dict(
+    ...     blockwise(
+    ...         add,
+    ...         'z', 'i',
+    ...         'x', 'i',
+    ...         Alias(delayed_inc.key), None,
+    ...         numblocks={'x': (2, )}
+    ...     )
+    ... )  # doctest: +NORMALIZE_WHITESPACE
+    {('z', 0): <Task ('z', 0) add(Alias(('x', 0)), Alias('dinc'))>,
+    ('z', 1): <Task ('z', 1) add(Alias(('x', 1)), Alias('dinc'))>}
 
     See Also
     --------
-    make_blockwise_graph
+    dask.array.blockwise
     Blockwise
     """
     new_axes = new_axes or {}
@@ -278,6 +437,7 @@ def blockwise(
         inputs.append(name)
         inputs_indices.append(index)
 
+    task_args = list(map(TaskRef, map(blockwise_token, range(len(inputs)))))
     # Unpack delayed objects in kwargs
     new_keys = {n for c in dependencies for n in c.__dask_layers__()}
     if kwargs:
@@ -286,28 +446,15 @@ def blockwise(
             blockwise_token(i) for i in range(len(inputs), len(inputs) + len(new_keys))
         )
         sub = dict(zip(new_keys, new_tokens))
-        inputs.extend(new_keys)
+        inputs.extend(map(TaskRef, new_keys))
         inputs_indices.extend((None,) * len(new_keys))
-        kwargs = subs(kwargs, sub)
 
     indices = [(k, v) for k, v in zip(inputs, inputs_indices)]
-    keys = map(blockwise_token, range(len(inputs)))
-
-    # Construct local graph
-    if not kwargs:
-        subgraph = {output: (func,) + tuple(keys)}
-    else:
-        _keys = list(keys)
-        if new_keys:
-            _keys = _keys[: -len(new_keys)]
-        kwargs2 = (dict, list(map(list, kwargs.items())))
-        subgraph = {output: (apply, func, _keys, kwargs2)}
-
-    # Construct final output
+    task = Task(output, func, *task_args, **kwargs)
     subgraph = Blockwise(
         output,
         output_indices,
-        subgraph,
+        task,
         indices,
         numblocks=numblocks,
         concatenate=concatenate,
@@ -372,8 +519,8 @@ class Blockwise(Layer):
 
     output: str
     output_indices: tuple[str, ...]
-    dsk: Graph
-    indices: tuple[tuple[str, tuple[str, ...] | None], ...]
+    task: Task
+    indices: tuple[tuple[str | TaskRef, tuple[str, ...] | None], ...]
     numblocks: Mapping[str, Sequence[int]]
     concatenate: bool | None
     new_axes: Mapping[str, int]
@@ -384,8 +531,8 @@ class Blockwise(Layer):
         self,
         output: str,
         output_indices: Iterable[str],
-        dsk: Graph,
-        indices: Iterable[tuple[str | BlockwiseDep, Iterable[str] | None]],
+        task: Task,
+        indices: Iterable[tuple[str | TaskRef | BlockwiseDep, Iterable[str] | None]],
         numblocks: Mapping[str, Sequence[int]],
         concatenate: bool | None = None,
         new_axes: Mapping[str, int] | None = None,
@@ -397,23 +544,29 @@ class Blockwise(Layer):
         self.output = output
         self.output_indices = tuple(output_indices)
         self.output_blocks = output_blocks
-        self.dsk = dsk
+        self.task = task
+        assert isinstance(task, Task)
 
         # Remove `BlockwiseDep` arguments from input indices
         # and add them to `self.io_deps`.
         # TODO: Remove `io_deps` and handle indexable objects
         # in `self.indices` throughout `Blockwise`.
         _tmp_indices = []
+        numblocks = dict(numblocks)
+        io_deps = dict(io_deps or {})
         if indices:
-            numblocks = ensure_dict(numblocks, copy=True)
-            io_deps = ensure_dict(io_deps or {}, copy=True)
             for dep, ind in indices:
+                if ind is not None:
+                    # FIXME: The Blockwise API is a little weird this way
+                    assert not isinstance(
+                        dep, TaskRef
+                    ), "TaskRef objects are only allowed for broadcasted inputs with None as index."
                 if isinstance(dep, BlockwiseDep):
                     name = tokenize(dep)
                     io_deps[name] = dep
                     numblocks[name] = dep.numblocks
                 else:
-                    name = dep
+                    name = dep  # type: ignore
                 _tmp_indices.append((name, tuple(ind) if ind is not None else ind))
         self.numblocks = numblocks
         self.io_deps = io_deps or {}
@@ -450,11 +603,9 @@ class Blockwise(Layer):
             return self._cached_dict["dsk"]
         else:
             keys = tuple(map(blockwise_token, range(len(self.indices))))
-            dsk, _ = fuse(self.dsk, [self.output])
-            func = SubgraphCallable(dsk, self.output, keys)
 
-            dsk = make_blockwise_graph(
-                func,
+            dsk = _make_blockwise_graph(
+                self.task,
                 self.output,
                 self.output_indices,
                 *list(toolz.concat(self.indices)),
@@ -464,6 +615,7 @@ class Blockwise(Layer):
                 output_blocks=self.output_blocks,
                 dims=self.dims,
                 io_deps=self.io_deps,
+                keys=keys,
             )
 
             self._cached_dict = {"dsk": dsk}
@@ -513,7 +665,6 @@ class Blockwise(Layer):
         # Generate coordinate map
         (coord_maps, concat_axes, dummies) = _get_coord_mapping(
             self.dims,
-            self.output,
             self.output_indices,
             self.numblocks,
             self.indices,
@@ -523,6 +674,8 @@ class Blockwise(Layer):
         # Gather constant dependencies (for all output keys)
         const_deps = set()
         for arg, ind in self.indices:
+            if isinstance(arg, TaskRef):
+                arg = arg.key
             if ind is None:
                 try:
                     if arg in all_hlg_keys:
@@ -539,7 +692,7 @@ class Blockwise(Layer):
                 if ind is not None and arg not in self.io_deps:
                     arg_coords = tuple(coords[c] for c in cmap)
                     if axes:
-                        tups = lol_product((arg,), arg_coords)
+                        tups = _lol_product((arg,), arg_coords)
                         deps.update(flatten(tups))
                         if concatenate:
                             tups = (concatenate, tups, axes)
@@ -554,6 +707,8 @@ class Blockwise(Layer):
                 for out_coords in output_blocks:
                     key = (self.output,) + out_coords
                     valid_key_dep = io_dep[out_coords]
+                    if isinstance(valid_key_dep, TaskRef):
+                        valid_key_dep = valid_key_dep.key
                     key_deps[key] |= {valid_key_dep}
 
         return key_deps
@@ -562,7 +717,7 @@ class Blockwise(Layer):
         return Blockwise(
             self.output,
             self.output_indices,
-            self.dsk,
+            self.task,
             self.indices,
             self.numblocks,
             concatenate=self.concatenate,
@@ -611,15 +766,19 @@ class Blockwise(Layer):
         is_leaf = True
 
         indices = []
-        k: Key
+        k: Key | TaskRef
         for k, idxv in self.indices:
             # Note: k may not be a key and thus not be hashable in the case where
             # one or more args of blockwise() are sequences of literals;
             # e.g. k = (list, [0, 1, 2])
             # See https://github.com/dask/dask/issues/8978
+
             if ishashable(k) and k in names:
                 is_leaf = False
-                k = clone_key(k, seed)
+                k = clone_key(k, seed)  # type: ignore
+            elif isinstance(k, TaskRef) and k.key in names:
+                is_leaf = False
+                k = TaskRef(clone_key(k.key, seed))
 
             indices.append((k, idxv))
 
@@ -630,22 +789,27 @@ class Blockwise(Layer):
                 k = clone_key(k, seed)
             numblocks[k] = nbv
 
-        dsk = {clone_key(k, seed): v for k, v in self.dsk.items()}
-
         if bind_to is not None and is_leaf:
             from dask.graph_manipulation import chunks
 
             # It's always a Delayed generated by dask.graph_manipulation.checkpoint;
             # the layer name always matches the key
             assert isinstance(bind_to, str)
-            dsk = {k: (chunks.bind, v, f"_{len(indices)}") for k, v in dsk.items()}
-            indices.append((bind_to, None))
+            newtask = Task(
+                clone_key(self.task.key, seed),
+                chunks.bind,
+                self.task,
+                TaskRef(blockwise_token(len(indices))),
+            )
+            indices.append((TaskRef(bind_to), None))
+        else:
+            newtask = self.task.substitute({}, key=clone_key(self.task.key, seed))
 
         return (
             Blockwise(
                 output=clone_key(self.output, seed),
                 output_indices=self.output_indices,
-                dsk=dsk,
+                task=newtask,
                 indices=indices,
                 numblocks=numblocks,
                 concatenate=self.concatenate,
@@ -660,7 +824,6 @@ class Blockwise(Layer):
 
 def _get_coord_mapping(
     dims,
-    output,
     out_indices,
     numblocks,
     argpairs,
@@ -682,8 +845,6 @@ def _get_coord_mapping(
         Mapping between each index specified in `argpairs` and
         the number of output blocks for that index. Corresponds
         to the Blockwise `dims` attribute.
-    output : str
-        Corresponds to the Blockwise `output` attribute.
     out_indices : tuple
         Corresponds to the Blockwise `output_indices` attribute.
     numblocks : dict
@@ -701,7 +862,7 @@ def _get_coord_mapping(
             block_names.add(name)
             for x in ind:
                 all_indices.add(x)
-    assert set(numblocks) == block_names
+    assert set(numblocks) == block_names, (numblocks, block_names)
 
     dummy_indices = all_indices - set(out_indices)
 
@@ -752,8 +913,8 @@ def _get_coord_mapping(
     return coord_maps, concat_axes, dummies
 
 
-def make_blockwise_graph(
-    func,
+def _make_blockwise_graph(
+    task,
     output,
     out_indices,
     *arrind_pairs,
@@ -763,111 +924,8 @@ def make_blockwise_graph(
     output_blocks=None,
     dims=None,
     io_deps=None,
+    keys=None,
 ):
-    """Tensor operation
-
-    Applies a function, ``func``, across blocks from many different input
-    collections.  We arrange the pattern with which those blocks interact with
-    sets of matching indices.  E.g.::
-
-        make_blockwise_graph(func, 'z', 'i', 'x', 'i', 'y', 'i')
-
-    yield an embarrassingly parallel communication pattern and is read as
-
-        $$ z_i = func(x_i, y_i) $$
-
-    More complex patterns may emerge, including multiple indices::
-
-        make_blockwise_graph(func, 'z', 'ij', 'x', 'ij', 'y', 'ji')
-
-        $$ z_{ij} = func(x_{ij}, y_{ji}) $$
-
-    Indices missing in the output but present in the inputs results in many
-    inputs being sent to one function (see examples).
-
-    Examples
-    --------
-    Simple embarrassing map operation
-
-    >>> inc = lambda x: x + 1
-    >>> make_blockwise_graph(inc, 'z', 'ij', 'x', 'ij', numblocks={'x': (2, 2)})  # doctest: +SKIP
-    {('z', 0, 0): (inc, ('x', 0, 0)),
-     ('z', 0, 1): (inc, ('x', 0, 1)),
-     ('z', 1, 0): (inc, ('x', 1, 0)),
-     ('z', 1, 1): (inc, ('x', 1, 1))}
-
-    Simple operation on two datasets
-
-    >>> add = lambda x, y: x + y
-    >>> make_blockwise_graph(add, 'z', 'ij', 'x', 'ij', 'y', 'ij', numblocks={'x': (2, 2),
-    ...                                                      'y': (2, 2)})  # doctest: +SKIP
-    {('z', 0, 0): (add, ('x', 0, 0), ('y', 0, 0)),
-     ('z', 0, 1): (add, ('x', 0, 1), ('y', 0, 1)),
-     ('z', 1, 0): (add, ('x', 1, 0), ('y', 1, 0)),
-     ('z', 1, 1): (add, ('x', 1, 1), ('y', 1, 1))}
-
-    Operation that flips one of the datasets
-
-    >>> addT = lambda x, y: x + y.T  # Transpose each chunk
-    >>> #                                        z_ij ~ x_ij y_ji
-    >>> #               ..         ..         .. notice swap
-    >>> make_blockwise_graph(addT, 'z', 'ij', 'x', 'ij', 'y', 'ji', numblocks={'x': (2, 2),
-    ...                                                       'y': (2, 2)})  # doctest: +SKIP
-    {('z', 0, 0): (add, ('x', 0, 0), ('y', 0, 0)),
-     ('z', 0, 1): (add, ('x', 0, 1), ('y', 1, 0)),
-     ('z', 1, 0): (add, ('x', 1, 0), ('y', 0, 1)),
-     ('z', 1, 1): (add, ('x', 1, 1), ('y', 1, 1))}
-
-    Dot product with contraction over ``j`` index.  Yields list arguments
-
-    >>> make_blockwise_graph(dotmany, 'z', 'ik', 'x', 'ij', 'y', 'jk', numblocks={'x': (2, 2),
-    ...                                                          'y': (2, 2)})  # doctest: +SKIP
-    {('z', 0, 0): (dotmany, [('x', 0, 0), ('x', 0, 1)],
-                            [('y', 0, 0), ('y', 1, 0)]),
-     ('z', 0, 1): (dotmany, [('x', 0, 0), ('x', 0, 1)],
-                            [('y', 0, 1), ('y', 1, 1)]),
-     ('z', 1, 0): (dotmany, [('x', 1, 0), ('x', 1, 1)],
-                            [('y', 0, 0), ('y', 1, 0)]),
-     ('z', 1, 1): (dotmany, [('x', 1, 0), ('x', 1, 1)],
-                            [('y', 0, 1), ('y', 1, 1)])}
-
-    Pass ``concatenate=True`` to concatenate arrays ahead of time
-
-    >>> make_blockwise_graph(f, 'z', 'i', 'x', 'ij', 'y', 'ij', concatenate=True,
-    ...     numblocks={'x': (2, 2), 'y': (2, 2,)})  # doctest: +SKIP
-    {('z', 0): (f, (concatenate_axes, [('x', 0, 0), ('x', 0, 1)], (1,)),
-                   (concatenate_axes, [('y', 0, 0), ('y', 0, 1)], (1,)))
-     ('z', 1): (f, (concatenate_axes, [('x', 1, 0), ('x', 1, 1)], (1,)),
-                   (concatenate_axes, [('y', 1, 0), ('y', 1, 1)], (1,)))}
-
-    Supports Broadcasting rules
-
-    >>> make_blockwise_graph(add, 'z', 'ij', 'x', 'ij', 'y', 'ij', numblocks={'x': (1, 2),
-    ...                                                      'y': (2, 2)})  # doctest: +SKIP
-    {('z', 0, 0): (add, ('x', 0, 0), ('y', 0, 0)),
-     ('z', 0, 1): (add, ('x', 0, 1), ('y', 0, 1)),
-     ('z', 1, 0): (add, ('x', 0, 0), ('y', 1, 0)),
-     ('z', 1, 1): (add, ('x', 0, 1), ('y', 1, 1))}
-
-    Support keyword arguments with apply
-
-    >>> def f(a, b=0): return a + b
-    >>> make_blockwise_graph(f, 'z', 'i', 'x', 'i', numblocks={'x': (2,)}, b=10)  # doctest: +SKIP
-    {('z', 0): (apply, f, [('x', 0)], {'b': 10}),
-     ('z', 1): (apply, f, [('x', 1)], {'b': 10})}
-
-    Include literals by indexing with ``None``
-
-    >>> make_blockwise_graph(add, 'z', 'i', 'x', 'i', 100, None,  numblocks={'x': (2,)})  # doctest: +SKIP
-    {('z', 0): (add, ('x', 0), 100),
-     ('z', 1): (add, ('x', 1), 100)}
-
-    See Also
-    --------
-    dask.array.blockwise
-    dask.blockwise.blockwise
-    """
-
     if numblocks is None:
         raise ValueError("Missing required numblocks argument.")
     new_axes = new_axes or {}
@@ -884,7 +942,6 @@ def make_blockwise_graph(
     # the actual graph
     (coord_maps, concat_axes, dummies) = _get_coord_mapping(
         dims,
-        output,
         out_indices,
         numblocks,
         argpairs,
@@ -900,46 +957,48 @@ def make_blockwise_graph(
     output_blocks = output_blocks or list(
         itertools.product(*[range(dims[i]) for i in out_indices])
     )
+    from dask._task_spec import DataNode
 
     dsk = {}
-    # Create argument lists
     for out_coords in output_blocks:
-        deps = set()
+        this_task = task
         coords = out_coords + dummies
         args = []
-        for cmap, axes, (arg, ind) in zip(coord_maps, concat_axes, argpairs):
+        for cmap, axes, (arg, ind), key in zip(
+            coord_maps, concat_axes, argpairs, keys, strict=True
+        ):
+            if key not in task.dependencies:
+                # FIXME: This feels like a bug
+                continue
             if ind is None:
-                args.append(arg)
+                if not isinstance(arg, (GraphNode, TaskRef)):
+                    this_task = this_task.substitute({key: DataNode(None, arg)})
+                else:
+                    this_task = this_task.substitute({key: arg})
+                continue
+            arg_coords = tuple(coords[c] for c in cmap)
+            if arg in io_deps:
+                val = parse_input(io_deps[arg].get(arg_coords, arg_coords))
+                if not isinstance(val, GraphNode):
+                    val = DataNode(None, val)
+                this_task = this_task.substitute({key: val})
             else:
-                arg_coords = tuple(coords[c] for c in cmap)
+                subs = {}
                 if axes:
-                    tups = lol_product((arg,), arg_coords)
-                    if arg not in io_deps:
-                        deps.update(flatten(tups))
-
+                    tups = _lol_product((arg,), arg_coords, as_taskref=True)
                     if concatenate:
-                        tups = (concatenate, tups, axes)
-                else:
-                    tups = (arg,) + arg_coords
-                    if arg not in io_deps:
-                        deps.add(tups)
-                # Replace "place-holder" IO keys with "real" args
-                if arg in io_deps:
-                    # We don't want to stringify keys for args
-                    # we are replacing here
-                    idx = tups[1:]
-                    args.append(io_deps[arg].get(idx, idx))
+                        tups = Task(key, concatenate, tups, axes)
+                    subs[key] = tups
                 else:
-                    args.append(tups)
-        out_key = (output,) + out_coords
-
-        args.insert(0, func)
-        dsk[out_key] = tuple(args)
-
+                    subs[key] = (arg, *arg_coords)
+                this_task = this_task.substitute(subs)
+        new_key = (output,) + out_coords
+        assert isinstance(this_task, Task)
+        dsk[new_key] = Task.fuse(this_task, *args, key=new_key)
     return dsk
 
 
-def lol_product(head, values):
+def _lol_product(head, values, as_taskref=False):
     """List of list of tuple keys, similar to `itertools.product`.
 
     Parameters
@@ -952,18 +1011,32 @@ def lol_product(head, values):
 
     Examples
     --------
-    >>> lol_product(('x',), (1, 2, 3))
+    >>> _lol_product(('x',), (1, 2, 3))
     ('x', 1, 2, 3)
-    >>> lol_product(('x',), (1, [2, 3], 4, [5, 6]))  # doctest: +NORMALIZE_WHITESPACE
+    >>> _lol_product(('x',), (1, [2, 3], 4, [5, 6]))  # doctest: +NORMALIZE_WHITESPACE
     [[('x', 1, 2, 4, 5), ('x', 1, 2, 4, 6)],
      [('x', 1, 3, 4, 5), ('x', 1, 3, 4, 6)]]
     """
     if not values:
+        if as_taskref:
+            return TaskRef(head)
         return head
     elif isinstance(values[0], list):
-        return [lol_product(head + (x,), values[1:]) for x in values[0]]
+        # FIXME: Constructor of List is odd
+        if as_taskref:
+            return List(
+                *(
+                    _lol_product(head + (x,), values[1:], as_taskref=as_taskref)
+                    for x in values[0]
+                )
+            )
+        else:
+            return list(
+                _lol_product(head + (x,), values[1:], as_taskref=as_taskref)
+                for x in values[0]
+            )
     else:
-        return lol_product(head + (values[0],), values[1:])
+        return _lol_product(head + (values[0],), values[1:], as_taskref=as_taskref)
 
 
 def lol_tuples(head, ind, values, dummies):
@@ -1236,13 +1309,14 @@ def rewrite_blockwise(inputs):
     indices = list(inputs[root].indices)
     new_axes = inputs[root].new_axes
     concatenate = inputs[root].concatenate
-    dsk = dict(inputs[root].dsk)
+    task = inputs[root].task
+    dsk = {task.key: task}
 
     changed = True
     while changed:
         changed = False
-        for i, (dep, ind) in enumerate(indices):
-            if ind is None:
+        for i, (dep, current_dep_indices) in enumerate(indices):
+            if current_dep_indices is None:
                 continue
             if dep not in inputs:
                 continue
@@ -1253,20 +1327,23 @@ def rewrite_blockwise(inputs):
             # (in different iteration orders) into a single
             # subgraph key/dependency
             # (see: https://github.com/dask/dask/issues/8535)
-            local_dep = dep if dep == root else _unique_dep(dep, ind)
+            local_dep = dep if dep == root else _unique_dep(dep, current_dep_indices)
 
             # Replace _n with dep name in existing tasks
             # (inc, _0) -> (inc, 'b')
-            dsk = {k: subs(v, {blockwise_token(i): local_dep}) for k, v in dsk.items()}
+            dsk = {
+                k: v.substitute({blockwise_token(i): local_dep}) for k, v in dsk.items()
+            }
 
             # Remove current input from input indices
             # [('a', 'i'), ('b', 'i')] -> [('a', 'i')]
-            _, current_dep_indices = indices.pop(i)
+            indices.pop(i)
+
             sub = {
                 blockwise_token(i): blockwise_token(i - 1)
                 for i in range(i + 1, len(indices) + 1)
             }
-            dsk = subs(dsk, sub)
+            dsk = {k: v.substitute(sub) for k, v in dsk.items()}
 
             # Change new input_indices to match give index from current computation
             # [('c', j')] -> [('c', 'i')]
@@ -1299,11 +1376,11 @@ def rewrite_blockwise(inputs):
                     index_map[id_key] = len(indices)
                     sub[blockwise_token(ii)] = blockwise_token(len(indices))
                     indices.append(index)
-            new_dsk = subs(inputs[dep].dsk, sub)
-
-            # Change new_dsk key to match local_dep
-            if dep != local_dep and dep in new_dsk:
-                new_dsk[local_dep] = new_dsk.pop(dep)
+            if dep != local_dep:
+                key = local_dep
+            else:
+                key = dep
+            new_dsk = {key: inputs[dep].task.substitute(sub, key=key)}
 
             # indices.extend(new_indices)
             dsk.update(new_dsk)
@@ -1323,8 +1400,8 @@ def rewrite_blockwise(inputs):
             new_indices.append(x)
 
     sub = {blockwise_token(k): blockwise_token(v) for k, v in sub.items()}
-    dsk = {k: subs(v, sub) for k, v in dsk.items() if k not in sub.keys()}
-
+    dsk = {k: v.substitute(sub) for k, v in dsk.items() if k not in sub.keys()}
+    task = Task.fuse(*dsk.values(), key=root)
     indices_check = {k for k, v in indices if v is not None}
     numblocks = toolz.merge([inp.numblocks for inp in inputs.values()])
     numblocks = {k: v for k, v in numblocks.items() if v is None or k in indices_check}
@@ -1337,7 +1414,7 @@ def rewrite_blockwise(inputs):
     return Blockwise(
         root,
         inputs[root].output_indices,
-        dsk,
+        task,
         new_indices,
         numblocks=numblocks,
         new_axes=new_axes,
@@ -1495,3 +1572,97 @@ def fuse_roots(graph: HighLevelGraph, keys: list):
             dependencies[name] = set()
 
     return HighLevelGraph(layers, dependencies)
+
+
+def _blockwise_unpack_collections_task_spec(expr):
+    # FIXME This is a copy of the delayed.unpack_collections function with the
+    # addition of the TaskSpec class. Eventually this should all be consolidated
+    # but to reduce the number of changes we'll vendor this here
+    # FIXME: There is also a dask.base version of unpack_collections that looks
+    # similar but is different. At the very least the names should be fixed
+    if isinstance(expr, Delayed):
+        return TaskRef(expr._key), (expr,)
+
+    if base.is_dask_collection(expr):
+        if hasattr(expr, "optimize"):
+            # Optimize dask-expr collections
+            expr = expr.optimize()
+
+        finalized = finalize(expr)
+        return TaskRef(finalized._key), (finalized,)
+
+    if type(expr) is type(iter(list())):
+        expr = list(expr)
+    elif type(expr) is type(iter(tuple())):
+        expr = tuple(expr)
+    elif type(expr) is type(iter(set())):
+        expr = set(expr)
+
+    typ = type(expr)
+
+    if typ in (list, tuple, set):
+        args, collections = utils.unzip(
+            (_blockwise_unpack_collections_task_spec(e) for e in expr), 2
+        )
+        collections = tuple(toolz.unique(toolz.concat(collections), key=id))
+        if not collections:
+            return expr, ()
+        args = List(*args)
+        # Ensure output type matches input type
+        if typ is not list:
+            args = Task(None, typ, args)
+        return args, collections
+
+    if typ is dict:
+        args, collections = _blockwise_unpack_collections_task_spec(
+            [[k, v] for k, v in expr.items()]
+        )
+        if not collections:
+            return expr, ()
+        return Dict(args), collections
+
+    if typ is slice:
+        args, collections = _blockwise_unpack_collections_task_spec(
+            [expr.start, expr.stop, expr.step]
+        )
+        if not collections:
+            return expr, ()
+        return Task(None, slice, *args), collections
+
+    if is_dataclass(expr):
+        args, collections = _blockwise_unpack_collections_task_spec(
+            [
+                [f.name, getattr(expr, f.name)]
+                for f in fields(expr)
+                if hasattr(expr, f.name)  # if init=False, field might not exist
+            ]
+        )
+        if not collections:
+            return expr, ()
+        try:
+            _fields = {
+                f.name: getattr(expr, f.name)
+                for f in fields(expr)
+                if hasattr(expr, f.name)
+            }
+            replace(expr, **_fields)
+        except (TypeError, ValueError) as e:
+            if isinstance(e, ValueError) or "is declared with init=False" in str(e):
+                raise ValueError(
+                    f"Failed to unpack {typ} instance. "
+                    "Note that using fields with `init=False` are not supported."
+                ) from e
+            else:
+                raise TypeError(
+                    f"Failed to unpack {typ} instance. "
+                    "Note that using a custom __init__ is not supported."
+                ) from e
+        return Task(None, typ, **dict(args)), collections
+
+    if utils.is_namedtuple_instance(expr):
+        args, collections = _blockwise_unpack_collections_task_spec([v for v in expr])
+        if not collections:
+            return expr, ()
+        return Task(None, typ, *args), collections
+
+    return expr, ()
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index e50035131..93014c43f 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -31,7 +31,13 @@ from dask.base import (
     named_schedulers,
     tokenize,
 )
-from dask.blockwise import Blockwise, BlockwiseDep, BlockwiseDepDict, blockwise
+from dask.blockwise import (
+    Blockwise,
+    BlockwiseDep,
+    BlockwiseDepDict,
+    _blockwise_unpack_collections_task_spec,
+    blockwise,
+)
 from dask.context import globalmethod
 from dask.dataframe import methods
 from dask.dataframe._compat import (
@@ -71,7 +77,7 @@ from dask.dataframe.utils import (
     raise_on_meta_error,
     valid_divisions,
 )
-from dask.delayed import Delayed, delayed, unpack_collections
+from dask.delayed import Delayed, delayed
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import DataFrameTreeReduction
 from dask.typing import Graph, NestedKeys, no_default
@@ -7050,7 +7056,7 @@ def map_partitions(
             dependencies.append(arg)
             continue
         arg = normalize_arg(arg)
-        arg2, collections = unpack_collections(arg)
+        arg2, collections = _blockwise_unpack_collections_task_spec(arg)
         if collections:
             args2.append(arg2)
             dependencies.extend(collections)
@@ -7061,7 +7067,7 @@ def map_partitions(
     simple = True
     for k, v in kwargs.items():
         v = normalize_arg(v)
-        v, collections = unpack_collections(v)
+        v, collections = _blockwise_unpack_collections_task_spec(v)
         dependencies.extend(collections)
         kwargs3[k] = v
         if collections:
diff --git a/dask/dataframe/io/csv.py b/dask/dataframe/io/csv.py
index d4d642787..fd9aca432 100644
--- a/dask/dataframe/io/csv.py
+++ b/dask/dataframe/io/csv.py
@@ -5,6 +5,8 @@ from collections.abc import Mapping
 from io import BytesIO
 from warnings import catch_warnings, simplefilter, warn
 
+from dask._task_spec import convert_legacy_task
+
 try:
     import psutil
 except ImportError:
@@ -368,7 +370,7 @@ def text_blocks_to_pandas(
 
     if path:
         colname, path_converter = path
-        paths = [b[1].path for b in blocks]
+        paths = [b.args[0].path for b in blocks]
         if path_converter:
             paths = [path_converter(p) for p in paths]
         head = head.assign(
@@ -663,7 +665,13 @@ def read_pandas(
             if is_integer_dtype(head[c].dtype) and c not in specified_dtypes:
                 head[c] = head[c].astype(float)
 
-    values = [[list(dsk.dask.values()) for dsk in block] for block in values]
+    values = [
+        [
+            [convert_legacy_task(k, ts, {}) for k, ts in dsk.dask.items()]
+            for dsk in block
+        ]
+        for block in values
+    ]
 
     return text_blocks_to_pandas(
         reader,
diff --git a/dask/dataframe/io/io.py b/dask/dataframe/io/io.py
index cbc145bc0..bfd330b5d 100644
--- a/dask/dataframe/io/io.py
+++ b/dask/dataframe/io/io.py
@@ -11,6 +11,7 @@ import numpy as np
 import pandas as pd
 
 import dask.array as da
+from dask._task_spec import TaskRef
 from dask.base import is_dask_collection, tokenize
 from dask.blockwise import BlockwiseDepDict, blockwise
 from dask.dataframe._compat import is_any_real_numeric_dtype
@@ -675,7 +676,7 @@ def from_delayed(
         name=name,
         columns=None,
         inputs=BlockwiseDepDict(
-            {(i,): inp.key for i, inp in enumerate(dfs)},
+            {(i,): TaskRef(inp.key) for i, inp in enumerate(dfs)},
             produces_keys=True,
         ),
         io_func=(
diff --git a/dask/dataframe/optimize.py b/dask/dataframe/optimize.py
index dd93fb12f..f4e7cded7 100644
--- a/dask/dataframe/optimize.py
+++ b/dask/dataframe/optimize.py
@@ -68,7 +68,7 @@ def optimize_dataframe_getitem(dsk, keys):
             return False
 
         # Callable must be `getitem`
-        if layer.dsk[layer.output][0] != operator.getitem:
+        if layer.task.func != operator.getitem:
             return False
 
         return True
@@ -216,7 +216,7 @@ def optimize_dataframe_getitem(dsk, keys):
                 new_block = Blockwise(
                     block.output,
                     block.output_indices,
-                    block.dsk,
+                    block.task,
                     new_indices,
                     numblocks,
                     block.concatenate,
diff --git a/dask/dataframe/rolling.py b/dask/dataframe/rolling.py
index de4dc38d9..b0a52152f 100644
--- a/dask/dataframe/rolling.py
+++ b/dask/dataframe/rolling.py
@@ -10,7 +10,7 @@ from pandas.core.window import Rolling as pd_Rolling
 
 from dask.array.core import normalize_arg
 from dask.base import tokenize
-from dask.blockwise import BlockwiseDepDict
+from dask.blockwise import BlockwiseDepDict, _blockwise_unpack_collections_task_spec
 from dask.dataframe import methods
 from dask.dataframe._compat import check_axis_keyword_deprecation
 from dask.dataframe.core import (
@@ -31,7 +31,6 @@ from dask.dataframe.utils import (
     is_dataframe_like,
     is_series_like,
 )
-from dask.delayed import unpack_collections
 from dask.highlevelgraph import HighLevelGraph
 from dask.typing import no_default
 from dask.utils import M, apply, derived_from, funcname, has_keyword
@@ -255,7 +254,7 @@ def map_overlap(
             dependencies.append(arg)
             continue
         arg = normalize_arg(arg)
-        arg2, collections = unpack_collections(arg)
+        arg2, collections = _blockwise_unpack_collections_task_spec(arg)
         if collections:
             args2.append(arg2)
             dependencies.extend(collections)
@@ -266,7 +265,7 @@ def map_overlap(
     simple = True
     for k, v in kwargs.items():
         v = normalize_arg(v)
-        v, collections = unpack_collections(v)
+        v, collections = _blockwise_unpack_collections_task_spec(v)
         dependencies.extend(collections)
         kwargs3[k] = v
         if collections:
diff --git a/dask/delayed.py b/dask/delayed.py
index e43fab1da..cf90b8afa 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -31,6 +31,7 @@ from dask.utils import (
     funcname,
     is_namedtuple_instance,
     methodcaller,
+    unzip,
 )
 
 __all__ = ["Delayed", "delayed"]
@@ -39,14 +40,6 @@ __all__ = ["Delayed", "delayed"]
 DEFAULT_GET = named_schedulers.get("threads", named_schedulers["sync"])
 
 
-def unzip(ls, nout):
-    """Unzip a list of lists into ``nout`` outputs."""
-    out = list(zip(*ls))
-    if not out:
-        out = [()] * nout
-    return out
-
-
 def finalize(collection):
     assert is_dask_collection(collection)
 
@@ -93,6 +86,7 @@ def unpack_collections(expr):
     >>> collections
     (Delayed('a'), Delayed('b'))
     """
+    # FIXME: See blockwise._unpack_collections for a TaskSpec version
     if isinstance(expr, Delayed):
         return expr._key, (expr,)
 
diff --git a/dask/graph_manipulation.py b/dask/graph_manipulation.py
index 93ed007ac..598ab9ec9 100644
--- a/dask/graph_manipulation.py
+++ b/dask/graph_manipulation.py
@@ -9,6 +9,7 @@ import uuid
 from collections.abc import Callable, Hashable
 from typing import Literal, TypeVar
 
+from dask._task_spec import List, TaskRef
 from dask.base import (
     clone_key,
     get_collection_names,
@@ -177,7 +178,11 @@ def _build_map_layer(
         except AttributeError:
             numblocks = (collection.npartitions,)
         indices = tuple(i for i, _ in enumerate(numblocks))
-        kwargs = {"_deps": [d.key for d in dependencies]} if dependencies else {}
+        kwargs = (
+            {"_deps": List(*[TaskRef(d.key) for d in dependencies])}
+            if dependencies
+            else {}
+        )
 
         return blockwise(
             func,
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
index 33bbb5fc6..3594eaead 100644
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -827,8 +827,8 @@ class HighLevelGraph(Graph):
         for k in dep_key1:
             if self.dependencies[k] != dependencies[k]:
                 raise ValueError(
-                    f"incorrect dependencies[{repr(k)}]: {repr(self.dependencies[k])} "
-                    f"expected {repr(dependencies[k])}"
+                    f"incorrect HLG dependencies[{repr(k)}]: {repr(self.dependencies[k])} "
+                    f"expected {repr(dependencies[k])} from task dependencies"
                 )
 
     def __repr__(self) -> str:
diff --git a/dask/layers.py b/dask/layers.py
index 5fe1bb482..ee02f1b25 100644
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any
 import tlz as toolz
 from tlz.curried import map
 
+from dask._task_spec import Task, TaskRef
 from dask.base import tokenize
 from dask.blockwise import Blockwise, BlockwiseDep, BlockwiseDepDict, blockwise_token
 from dask.core import flatten
@@ -1089,11 +1090,11 @@ class DataFrameIOLayer(Blockwise):
             io_arg_map = inputs
 
         # Use Blockwise initializer
-        dsk = {self.name: (io_func, blockwise_token(0))}
+        task = Task(self.name, io_func, TaskRef(blockwise_token(0)))
         super().__init__(
             output=self.name,
             output_indices="i",
-            dsk=dsk,
+            task=task,
             indices=[(io_arg_map, "i")],
             numblocks={},
             annotations=annotations,
diff --git a/dask/utils.py b/dask/utils.py
index 3f42bd6c0..f21ba3b2f 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -2308,3 +2308,11 @@ class shorten_traceback:
         #     return exc_tb.tb_next
 
         return exc_tb
+
+
+def unzip(ls, nout):
+    """Unzip a list of lists into ``nout`` outputs."""
+    out = list(zip(*ls))
+    if not out:
+        out = [()] * nout
+    return out


Test patch:
diff --git a/conftest.py b/conftest.py
index 988da63b6..f9eb31d77 100644
--- a/conftest.py
+++ b/conftest.py
@@ -61,3 +61,32 @@ def pytest_addoption(parser):
 def pytest_runtest_setup(item):
     if "slow" in item.keywords and not item.config.getoption("--runslow"):
         pytest.skip("need --runslow option to run")
+
+
+def pytest_assertrepr_compare(op, left, right):
+    import difflib
+
+    from dask._task_spec import Task
+
+    if isinstance(left, Task) and isinstance(right, Task):
+
+        def _get_attrs(node):
+            return (
+                [
+                    str(node.func),
+                ]
+                + sorted([str(a) for a in node.args])
+                + sorted([f"{k}: {v}" for k, v in node.kwargs.items()])
+            )
+
+        diff = list(
+            difflib.ndiff(
+                _get_attrs(left),
+                _get_attrs(right),
+            )
+        )
+        return [
+            "Comparing two dask graph nodes:",
+            f" left: {left.key} right: {right.key}",
+            " Diff:",
+        ] + diff
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index e3d8b8a2d..4ad9cf67a 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -8,7 +8,7 @@ import xml.etree.ElementTree
 
 import pytest
 
-from dask._task_spec import Task
+from dask._task_spec import DataNode, List, Task, TaskRef
 
 np = pytest.importorskip("numpy")
 
@@ -57,8 +57,8 @@ from dask.array.reshape import _not_implemented_message
 from dask.array.utils import assert_eq, same_keys
 from dask.base import compute_as_if_collection, tokenize
 from dask.blockwise import (
+    _make_blockwise_graph,
     broadcast_dimensions,
-    make_blockwise_graph,
     optimize_blockwise,
 )
 from dask.delayed import Delayed, delayed
@@ -94,55 +94,126 @@ def test_graph_from_arraylike(inline_array):
 
 
 def test_top():
-    assert make_blockwise_graph(inc, "z", "ij", "x", "ij", numblocks={"x": (2, 2)}) == {
-        ("z", 0, 0): (inc, ("x", 0, 0)),
-        ("z", 0, 1): (inc, ("x", 0, 1)),
-        ("z", 1, 0): (inc, ("x", 1, 0)),
-        ("z", 1, 1): (inc, ("x", 1, 1)),
+    t = Task("inc", inc, TaskRef("A"))
+    assert _make_blockwise_graph(
+        t, "z", "ij", "x", "ij", numblocks={"x": (2, 2)}, keys=("A",)
+    ) == {
+        ("z", 0, 0): Task(("z", 0, 0), inc, TaskRef(("x", 0, 0))),
+        ("z", 0, 1): Task(("z", 0, 1), inc, TaskRef(("x", 0, 1))),
+        ("z", 1, 0): Task(("z", 1, 0), inc, TaskRef(("x", 1, 0))),
+        ("z", 1, 1): Task(("z", 1, 1), inc, TaskRef(("x", 1, 1))),
     }
 
-    assert make_blockwise_graph(
-        add, "z", "ij", "x", "ij", "y", "ij", numblocks={"x": (2, 2), "y": (2, 2)}
+    keys = list(map(str, range(2)))
+    t = Task("add", add, *(TaskRef(k) for k in keys))
+    assert _make_blockwise_graph(
+        t,
+        "z",
+        "ij",
+        "x",
+        "ij",
+        "y",
+        "ij",
+        numblocks={"x": (2, 2), "y": (2, 2)},
+        keys=keys,
     ) == {
-        ("z", 0, 0): (add, ("x", 0, 0), ("y", 0, 0)),
-        ("z", 0, 1): (add, ("x", 0, 1), ("y", 0, 1)),
-        ("z", 1, 0): (add, ("x", 1, 0), ("y", 1, 0)),
-        ("z", 1, 1): (add, ("x", 1, 1), ("y", 1, 1)),
+        ("z", 0, 0): Task(("z", 0, 0), add, TaskRef(("x", 0, 0)), TaskRef(("y", 0, 0))),
+        ("z", 0, 1): Task(("z", 0, 1), add, TaskRef(("x", 0, 1)), TaskRef(("y", 0, 1))),
+        ("z", 1, 0): Task(("z", 1, 0), add, TaskRef(("x", 1, 0)), TaskRef(("y", 1, 0))),
+        ("z", 1, 1): Task(("z", 1, 1), add, TaskRef(("x", 1, 1)), TaskRef(("y", 1, 1))),
     }
+    keys = list(map(str, range(2)))
+    t = Task(
+        "dotmany",
+        dotmany,
+        TaskRef(keys[0]),
+        TaskRef(keys[1]),
+    )
+    out = _make_blockwise_graph(
+        t,
+        "z",
+        "ik",
+        "x",
+        "ij",
+        "y",
+        "jk",
+        numblocks={"x": (2, 2), "y": (2, 2)},
+        keys=keys,
+    )
+    expected = {
+        ("z", 0, 0): Task(
+            ("z", 0, 0),
+            dotmany,
+            List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),
+            List(TaskRef(("y", 0, 0)), TaskRef(("y", 1, 0))),
+        ),
+        ("z", 0, 1): Task(
+            ("z", 0, 1),
+            dotmany,
+            List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),
+            List(TaskRef(("y", 0, 1)), TaskRef(("y", 1, 1))),
+        ),
+        ("z", 1, 0): Task(
+            ("z", 1, 0),
+            dotmany,
+            List(TaskRef(("x", 1, 0)), TaskRef(("x", 1, 1))),
+            List(TaskRef(("y", 0, 0)), TaskRef(("y", 1, 0))),
+        ),
+        ("z", 1, 1): Task(
+            ("z", 1, 1),
+            dotmany,
+            List(TaskRef(("x", 1, 0)), TaskRef(("x", 1, 1))),
+            List(TaskRef(("y", 0, 1)), TaskRef(("y", 1, 1))),
+        ),
+    }
+    assert out[("z", 0, 0)] == expected[("z", 0, 0)]
+    assert out == expected
 
-    assert make_blockwise_graph(
-        dotmany, "z", "ik", "x", "ij", "y", "jk", numblocks={"x": (2, 2), "y": (2, 2)}
+    t = Task("identity", identity, TaskRef("0"))
+    assert _make_blockwise_graph(
+        t, "z", "", "x", "ij", numblocks={"x": (2, 2)}, keys=("0",)
     ) == {
-        ("z", 0, 0): (dotmany, [("x", 0, 0), ("x", 0, 1)], [("y", 0, 0), ("y", 1, 0)]),
-        ("z", 0, 1): (dotmany, [("x", 0, 0), ("x", 0, 1)], [("y", 0, 1), ("y", 1, 1)]),
-        ("z", 1, 0): (dotmany, [("x", 1, 0), ("x", 1, 1)], [("y", 0, 0), ("y", 1, 0)]),
-        ("z", 1, 1): (dotmany, [("x", 1, 0), ("x", 1, 1)], [("y", 0, 1), ("y", 1, 1)]),
+        ("z",): Task(
+            ("z",),
+            identity,
+            List(
+                List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),
+                List(TaskRef(("x", 1, 0)), TaskRef(("x", 1, 1))),
+            ),
+        )
     }
 
-    assert make_blockwise_graph(
-        identity, "z", "", "x", "ij", numblocks={"x": (2, 2)}
-    ) == {("z",): (identity, [[("x", 0, 0), ("x", 0, 1)], [("x", 1, 0), ("x", 1, 1)]])}
-
 
 def test_top_supports_broadcasting_rules():
-    assert make_blockwise_graph(
-        add, "z", "ij", "x", "ij", "y", "ij", numblocks={"x": (1, 2), "y": (2, 1)}
+
+    t = Task("add", add, TaskRef("A"), TaskRef("B"))
+    assert _make_blockwise_graph(
+        t,
+        "z",
+        "ij",
+        "x",
+        "ij",
+        "y",
+        "ij",
+        numblocks={"x": (1, 2), "y": (2, 1)},
+        keys=("A", "B"),
     ) == {
-        ("z", 0, 0): (add, ("x", 0, 0), ("y", 0, 0)),
-        ("z", 0, 1): (add, ("x", 0, 1), ("y", 0, 0)),
-        ("z", 1, 0): (add, ("x", 0, 0), ("y", 1, 0)),
-        ("z", 1, 1): (add, ("x", 0, 1), ("y", 1, 0)),
+        ("z", 0, 0): Task(("z", 0, 0), add, TaskRef(("x", 0, 0)), TaskRef(("y", 0, 0))),
+        ("z", 0, 1): Task(("z", 0, 1), add, TaskRef(("x", 0, 1)), TaskRef(("y", 0, 0))),
+        ("z", 1, 0): Task(("z", 1, 0), add, TaskRef(("x", 0, 0)), TaskRef(("y", 1, 0))),
+        ("z", 1, 1): Task(("z", 1, 1), add, TaskRef(("x", 0, 1)), TaskRef(("y", 1, 0))),
     }
 
 
 def test_top_literals():
-    assert make_blockwise_graph(
-        add, "z", "ij", "x", "ij", 123, None, numblocks={"x": (2, 2)}
+    t = Task("add", add, TaskRef("A"), TaskRef("B"))
+    assert _make_blockwise_graph(
+        t, "z", "ij", "x", "ij", 123, None, numblocks={"x": (2, 2)}, keys=("A", "B")
     ) == {
-        ("z", 0, 0): (add, ("x", 0, 0), 123),
-        ("z", 0, 1): (add, ("x", 0, 1), 123),
-        ("z", 1, 0): (add, ("x", 1, 0), 123),
-        ("z", 1, 1): (add, ("x", 1, 1), 123),
+        ("z", 0, 0): Task(("z", 0, 0), add, TaskRef(("x", 0, 0)), DataNode(None, 123)),
+        ("z", 0, 1): Task(("z", 0, 1), add, TaskRef(("x", 0, 1)), DataNode(None, 123)),
+        ("z", 1, 0): Task(("z", 1, 0), add, TaskRef(("x", 1, 0)), DataNode(None, 123)),
+        ("z", 1, 1): Task(("z", 1, 1), add, TaskRef(("x", 1, 1)), DataNode(None, 123)),
     }
 
 
@@ -165,7 +236,7 @@ def test_blockwise_1_in_shape_I():
         assert 1 in b.shape
 
     p, k, N = 7, 2, 5
-    da.blockwise(
+    arr = da.blockwise(
         test_f,
         "x",
         da.zeros((2 * p, 9, k * N), chunks=(p, 3, k)),
@@ -174,7 +245,8 @@ def test_blockwise_1_in_shape_I():
         "xzt",
         concatenate=True,
         dtype=float,
-    ).compute()
+    )
+    arr.compute()
 
 
 def test_blockwise_1_in_shape_II():
@@ -222,8 +294,23 @@ def test_chunked_dot_product():
     getx = graph_from_arraylike(x, (5, 5), shape=(20, 20), name="x")
     geto = graph_from_arraylike(o, (5, 5), shape=(20, 20), name="o")
 
-    result = make_blockwise_graph(
-        dotmany, "out", "ik", "x", "ij", "o", "jk", numblocks={"x": (4, 4), "o": (4, 4)}
+    keys = list(map(str, range(2)))
+    t = Task(
+        "dotmany",
+        dotmany,
+        TaskRef(keys[0]),
+        TaskRef(keys[1]),
+    )
+    result = _make_blockwise_graph(
+        t,
+        "out",
+        "ik",
+        "x",
+        "ij",
+        "o",
+        "jk",
+        numblocks={"x": (4, 4), "o": (4, 4)},
+        keys=keys,
     )
 
     dsk = merge(getx, geto, result)
@@ -237,8 +324,10 @@ def test_chunked_transpose_plus_one():
 
     getx = graph_from_arraylike(x, (5, 5), shape=(20, 20), name="x")
 
-    f = lambda x: x.T + 1
-    comp = make_blockwise_graph(f, "out", "ij", "x", "ji", numblocks={"x": (4, 4)})
+    f = Task("f", lambda x: x.T + 1, TaskRef("x"))
+    comp = _make_blockwise_graph(
+        f, "out", "ij", "x", "ji", numblocks={"x": (4, 4)}, keys=("x",)
+    )
 
     dsk = merge(getx, comp)
     out = dask.get(dsk, [[("out", i, j) for j in range(4)] for i in range(4)])
@@ -3757,9 +3846,13 @@ def test_map_blocks_delayed():
     y = np.ones((5, 5))
 
     z = x.map_blocks(add, y, dtype=x.dtype)
+    z.dask.validate()
+    dask.optimize(z)[0].dask.validate()
 
     yy = delayed(y)
     zz = x.map_blocks(add, yy, dtype=x.dtype)
+    zz.dask.validate()
+    dask.optimize(zz)[0].dask.validate()
 
     assert_eq(z, zz)
 
@@ -5023,16 +5116,21 @@ def test_map_blocks_large_inputs_delayed():
 
 
 def test_blockwise_large_inputs_delayed():
+    def func(a, b):
+        return a
+
     a = da.ones(10, chunks=(5,))
     b = np.ones(1000000)
 
-    c = da.blockwise(add, "i", a, "i", b, None, dtype=a.dtype)
+    c = da.blockwise(func, "i", a, "i", b, None, dtype=a.dtype)
     assert any(b is v for v in c.dask.values())
     assert repr(dict(c.dask)).count(repr(b)[:10]) == 1  # only one occurrence
+    assert_eq(c, c)
 
-    d = da.blockwise(lambda x, y: x + y, "i", a, "i", y=b, dtype=a.dtype)
+    d = da.blockwise(lambda x, y: x, "i", a, "i", y=b, dtype=a.dtype)
     assert any(b is v for v in d.dask.values())
     assert repr(dict(c.dask)).count(repr(b)[:10]) == 1  # only one occurrence
+    assert_eq(d, d)
 
 
 def test_slice_reversed():
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
index bc77a1b3f..20fa0e86b 100644
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -8,6 +8,7 @@ import pytest
 
 import dask
 import dask.array as da
+from dask._task_spec import Task, TaskRef
 from dask.array.utils import assert_eq
 from dask.blockwise import (
     _BLOCKWISE_DEFAULT_PREFIX,
@@ -21,8 +22,8 @@ from dask.highlevelgraph import HighLevelGraph
 from dask.utils_test import dec, hlg_layer_topological, inc
 
 a, b, c, d, e, f, g = "a", "b", "c", "d", "e", "f", "g"
-_0, _1, _2, _3, _4, _5, _6, _7, _8, _9 = (
-    f"{_BLOCKWISE_DEFAULT_PREFIX}{i}" for i in range(10)
+_0, _1, _2, _3, _4, _5, _6, _7, _8, _9 = tuple(
+    map(TaskRef, (f"{_BLOCKWISE_DEFAULT_PREFIX}{i}" for i in range(10)))
 )
 i, j, k = "i", "j", "k"
 
@@ -31,258 +32,345 @@ i, j, k = "i", "j", "k"
     "inputs,expected",
     [
         # output name, output index, task, input indices
-        [[(b, "i", {b: (inc, _0)}, [(a, "i")])], (b, "i", {b: (inc, _0)}, [(a, "i")])],
+        # Case 0
+        [
+            [(b, "i", Task(b, inc, _0), [(a, "i")])],
+            (
+                b,
+                "i",
+                Task(b, inc, _0),
+                [
+                    (a, "i"),
+                ],
+            ),
+        ],
+        # Case 1
         [
             [
-                (b, "i", {b: (inc, _0)}, [(a, "i")]),
-                (c, "i", {c: (dec, _0)}, [(a, "i")]),
-                (d, "i", {d: (add, _0, _1, _2)}, [(a, "i"), (b, "i"), (c, "i")]),
+                (b, "i", Task(b, inc, _0), [(a, "i")]),
+                (c, "i", Task(c, dec, _0), [(a, "i")]),
+                (d, "i", Task(d, add, _0, _1, _2), [(a, "i"), (b, "i"), (c, "i")]),
             ],
             (
                 d,
                 "i",
-                {
-                    _unique_dep(b, "i"): (inc, _0),
-                    _unique_dep(c, "i"): (dec, _0),
-                    d: (add, _0, _unique_dep(b, "i"), _unique_dep(c, "i")),
-                },
+                Task.fuse(
+                    Task(
+                        d,
+                        add,
+                        _0,
+                        TaskRef(_unique_dep(b, "i")),
+                        TaskRef(_unique_dep(c, "i")),
+                    ),
+                    Task(_unique_dep(c, "i"), dec, _0),
+                    Task(_unique_dep(b, "i"), inc, _0),
+                ),
                 [(a, "i")],
             ),
         ],
+        # Case 2
         [
             [
-                (b, "i", {b: (inc, _0)}, [(a, "i")]),
-                (c, "j", {c: (inc, _0)}, [(b, "j")]),
+                (b, "i", Task(b, inc, _0), [(a, "i")]),
+                (c, "j", Task(c, inc, _0), [(b, "j")]),
             ],
             (
                 c,
                 "j",
-                {_unique_dep(b, "j"): (inc, _0), c: (inc, _unique_dep(b, "j"))},
+                Task.fuse(
+                    Task(c, inc, TaskRef(_unique_dep(b, "j"))),
+                    Task(_unique_dep(b, "j"), inc, _0),
+                ),
                 [(a, "j")],
             ),
         ],
+        # Case 3
         [
             [
-                (b, "i", {b: (sum, _0)}, [(a, "ij")]),
-                (c, "k", {c: (inc, _0)}, [(b, "k")]),
+                (b, "i", Task(b, sum, _0), [(a, "ij")]),
+                (c, "k", Task(c, inc, _0), [(b, "k")]),
             ],
             (
                 c,
                 "k",
-                {_unique_dep(b, "k"): (sum, _0), c: (inc, _unique_dep(b, "k"))},
+                Task.fuse(
+                    Task(c, inc, TaskRef(_unique_dep(b, "k"))),
+                    Task(_unique_dep(b, "k"), sum, _0),
+                ),
                 [(a, "kA")],
             ),
         ],
+        # Case 4
         [
             [
-                (c, "i", {c: (inc, _0)}, [(a, "i")]),
-                (d, "i", {d: (inc, _0)}, [(b, "i")]),
-                (g, "ij", {g: (add, _0, _1)}, [(c, "i"), (d, "j")]),
+                (c, "i", Task(c, inc, _0), [(a, "i")]),
+                (d, "i", Task(d, inc, _0), [(b, "i")]),
+                (g, "ij", Task(g, add, _0, _1), [(c, "i"), (d, "j")]),
             ],
             (
                 g,
                 "ij",
-                {
-                    g: (add, _unique_dep(c, "i"), _unique_dep(d, "j")),
-                    _unique_dep(c, "i"): (inc, _0),
-                    _unique_dep(d, "j"): (inc, _1),
-                },
+                Task.fuse(
+                    Task(
+                        g,
+                        add,
+                        TaskRef(_unique_dep(c, "i")),
+                        TaskRef(_unique_dep(d, "j")),
+                    ),
+                    Task(_unique_dep(c, "i"), inc, _0),
+                    Task(_unique_dep(d, "j"), inc, _1),
+                ),
                 [(a, "i"), (b, "j")],
             ),
         ],
+        # Case 5
         [
             [
-                (b, "ji", {b: (np.transpose, _0)}, [(a, "ij")]),
-                (c, "ij", {c: (add, _0, _1)}, [(a, "ij"), (b, "ij")]),
+                (b, "ji", Task(b, np.transpose, _0), [(a, "ij")]),
+                (c, "ij", Task(c, add, _0, _1), [(a, "ij"), (b, "ij")]),
             ],
             (
                 c,
                 "ij",
-                {
-                    c: (add, _0, _unique_dep(b, "ij")),
-                    _unique_dep(b, "ij"): (np.transpose, _1),
-                },
+                Task.fuse(
+                    Task(c, add, _0, TaskRef(_unique_dep(b, "ij"))),
+                    Task(_unique_dep(b, "ij"), np.transpose, _1),
+                ),
                 [(a, "ij"), (a, "ji")],
             ),
         ],
+        # Case 6
         [
             [
-                (c, "i", {c: (add, _0, _1)}, [(a, "i"), (b, "i")]),
-                (d, "i", {d: (inc, _0)}, [(c, "i")]),
+                (c, "i", Task(c, add, _0, _1), [(a, "i"), (b, "i")]),
+                (d, "i", Task(d, inc, _0), [(c, "i")]),
             ],
             (
                 d,
                 "i",
-                {d: (inc, _unique_dep(c, "i")), _unique_dep(c, "i"): (add, _0, _1)},
+                Task.fuse(
+                    Task(d, inc, TaskRef(_unique_dep(c, "i"))),
+                    Task(_unique_dep(c, "i"), add, _0, _1),
+                ),
                 [(a, "i"), (b, "i")],
             ),
         ],
+        # Case 7
         [
             [
-                (b, "ij", {b: (np.transpose, _0)}, [(a, "ji")]),
-                (d, "ij", {d: (np.dot, _0, _1)}, [(b, "ik"), (c, "kj")]),
+                (b, "ij", Task(b, np.transpose, _0), [(a, "ji")]),
+                (d, "ij", Task(d, np.dot, _0, _1), [(b, "ik"), (c, "kj")]),
             ],
             (
                 d,
                 "ij",
-                {
-                    d: (np.dot, _unique_dep(b, "ik"), _0),
-                    _unique_dep(b, "ik"): (np.transpose, _1),
-                },
+                Task.fuse(
+                    Task(
+                        d,
+                        np.dot,
+                        TaskRef(_unique_dep(b, "ik")),
+                        _0,
+                    ),
+                    Task(_unique_dep(b, "ik"), np.transpose, _1),
+                ),
                 [(c, "kj"), (a, "ki")],
             ),
         ],
+        # Case 8
         [
             [
-                (c, "i", {c: (add, _0, _1)}, [(a, "i"), (b, "i")]),
-                (f, "i", {f: (add, _0, _1)}, [(d, "i"), (e, "i")]),
-                (g, "i", {g: (add, _0, _1)}, [(c, "i"), (f, "i")]),
+                (c, "i", Task(c, add, _0, _1), [(a, "i"), (b, "i")]),
+                (f, "i", Task(f, add, _0, _1), [(d, "i"), (e, "i")]),
+                (g, "i", Task(g, add, _0, _1), [(c, "i"), (f, "i")]),
             ],
             (
                 g,
                 "i",
-                {
-                    g: (add, _unique_dep(c, "i"), _unique_dep(f, "i")),
-                    _unique_dep(f, "i"): (add, _2, _3),
-                    _unique_dep(c, "i"): (add, _0, _1),
-                },
+                Task.fuse(
+                    Task(
+                        g,
+                        add,
+                        TaskRef(_unique_dep(c, "i")),
+                        TaskRef(_unique_dep(f, "i")),
+                    ),
+                    Task(_unique_dep(f, "i"), add, _2, _3),
+                    Task(_unique_dep(c, "i"), add, _0, _1),
+                ),
                 [(a, i), (b, i), (d, i), (e, i)],
             ),
         ],
+        # Case 9
         [
             [
-                (c, "i", {c: (add, _0, _1)}, [(a, "i"), (b, "i")]),
-                (f, "i", {f: (add, _0, _1)}, [(a, "i"), (e, "i")]),
-                (g, "i", {g: (add, _0, _1)}, [(c, "i"), (f, "i")]),
+                (c, "i", Task(c, add, _0, _1), [(a, "i"), (b, "i")]),
+                (f, "i", Task(f, add, _0, _1), [(a, "i"), (e, "i")]),
+                (g, "i", Task(g, add, _0, _1), [(c, "i"), (f, "i")]),
             ],
             (
                 g,
                 "i",
-                {
-                    g: (add, _unique_dep(c, "i"), _unique_dep(f, "i")),
-                    _unique_dep(f, "i"): (add, _0, _2),
-                    _unique_dep(c, "i"): (add, _0, _1),
-                },
+                Task.fuse(
+                    Task(
+                        g,
+                        add,
+                        TaskRef(_unique_dep(c, "i")),
+                        TaskRef(_unique_dep(f, "i")),
+                    ),
+                    Task(_unique_dep(f, "i"), add, _0, _2),
+                    Task(_unique_dep(c, "i"), add, _0, _1),
+                ),
                 [(a, "i"), (b, "i"), (e, "i")],
             ),
         ],
+        # Case 10
         [
             [
-                (b, "i", {b: (sum, _0)}, [(a, "ij")]),
-                (c, "i", {c: (inc, _0)}, [(b, "i")]),
+                (b, "i", Task(b, sum, _0), [(a, "ij")]),
+                (c, "i", Task(c, inc, _0), [(b, "i")]),
             ],
             (
                 c,
                 "i",
-                {c: (inc, _unique_dep(b, "i")), _unique_dep(b, "i"): (sum, _0)},
+                Task.fuse(
+                    Task(c, inc, TaskRef(_unique_dep(b, "i"))),
+                    Task(_unique_dep(b, "i"), sum, _0),
+                ),
                 [(a, "iA")],
             ),
         ],
+        # Case 11
         [
             [
-                (c, "i", {c: (inc, _0)}, [(b, "i")]),
-                (d, "i", {d: (add, _0, _1, _2)}, [(a, "i"), (b, "i"), (c, "i")]),
+                (c, "i", Task(c, inc, _0), [(b, "i")]),
+                (d, "i", Task(d, add, _0, _1, _2), [(a, "i"), (b, "i"), (c, "i")]),
             ],
             (
                 d,
                 "i",
-                {d: (add, _0, _1, _unique_dep(c, "i")), _unique_dep(c, "i"): (inc, _1)},
+                Task.fuse(
+                    Task(
+                        d,
+                        add,
+                        _0,
+                        _1,
+                        TaskRef(_unique_dep(c, "i")),
+                    ),
+                    Task(_unique_dep(c, "i"), inc, _1),
+                ),
                 [(a, "i"), (b, "i")],
             ),
         ],
+        # Case 12
         # Include literals
         [
-            [(b, "i", {b: (add, _0, _1)}, [(a, "i"), (123, None)])],
-            (b, "i", {b: (add, _0, _1)}, [(a, "i"), (123, None)]),
+            [(b, "i", Task(b, add, _0, _1), [(a, "i"), (123, None)])],
+            (b, "i", Task(b, add, _0, _1), [(a, "i"), (123, None)]),
         ],
+        # Case 13
         [
             [
-                (b, "i", {b: (add, _0, _1)}, [(a, "i"), (123, None)]),
-                (c, "j", {c: (add, _0, _1)}, [(b, "j"), (456, None)]),
+                (b, "i", Task(b, add, _0, _1), [(a, "i"), (123, None)]),
+                (c, "j", Task(c, add, _0, _1), [(b, "j"), (456, None)]),
             ],
             (
                 c,
                 "j",
-                {_unique_dep(b, "j"): (add, _1, _2), c: (add, _unique_dep(b, "j"), _0)},
+                Task.fuse(
+                    Task(c, add, TaskRef(_unique_dep(b, "j")), _0),
+                    Task(_unique_dep(b, "j"), add, _1, _2),
+                ),
                 [(456, None), (a, "j"), (123, None)],
             ),
         ],
+        # Case 14
         # Literals that compare equal (e.g. 0 and False) aren't deduplicated
         [
             [
-                (b, "i", {b: (add, _0, _1)}, [(a, "i"), (0, None)]),
-                (c, "j", {c: (add, _0, _1)}, [(b, "j"), (False, None)]),
+                (b, "i", Task(b, add, _0, _1), [(a, "i"), (0, None)]),
+                (c, "j", Task(c, add, _0, _1), [(b, "j"), (False, None)]),
             ],
             (
                 c,
                 "j",
-                {_unique_dep(b, "j"): (add, _1, _2), c: (add, _unique_dep(b, "j"), _0)},
+                Task.fuse(
+                    Task(c, add, TaskRef(_unique_dep(b, "j")), _0),
+                    Task(_unique_dep(b, "j"), add, _1, _2),
+                ),
                 [(False, None), (a, "j"), (0, None)],
             ),
         ],
+        # Case 15
         # Literals are deduplicated
         [
             [
-                (b, "i", {b: (add, _0, _1)}, [(a, "i"), (123, None)]),
-                (c, "j", {c: (add, _0, _1)}, [(b, "j"), (123, None)]),
+                (b, "i", Task(b, add, _0, _1), [(a, "i"), (123, None)]),
+                (c, "j", Task(c, add, _0, _1), [(b, "j"), (123, None)]),
             ],
             (
                 c,
                 "j",
-                {_unique_dep(b, "j"): (add, _1, _0), c: (add, _unique_dep(b, "j"), _0)},
+                Task.fuse(
+                    Task(c, add, TaskRef(_unique_dep(b, "j")), _0),
+                    Task(_unique_dep(b, "j"), add, _1, _0),
+                ),
                 [(123, None), (a, "j")],
             ),
         ],
+        # Case 16
         # Check cases where two distinct indices are used
         # for the same dependency name, and where the same
         # dependency-index combination is repeated
         # (See: https://github.com/dask/dask/issues/8535)
         [
             [
-                (b, "jk", {b: (add, _0, _1)}, [(a, "jk"), (2, None)]),
-                (c, "ijk", {c: (add, _0, _1)}, [(b, "ij"), (b, "jk")]),
-                (d, "ijk", {d: (inc, _0, _1)}, [(c, "ijk"), (123, None)]),
+                (b, "jk", Task(b, add, _0, _1), [(a, "jk"), (2, None)]),
+                (c, "ijk", Task(c, add, _0, _1), [(b, "ij"), (b, "jk")]),
+                (d, "ijk", Task(d, inc, _0, _1), [(c, "ijk"), (123, None)]),
             ],
             (
                 "d",
                 "ijk",
-                {
-                    "d": (inc, _unique_dep(c, "ijk"), _0),
-                    _unique_dep(c, "ijk"): (
+                Task.fuse(
+                    Task(d, inc, TaskRef(_unique_dep(c, "ijk")), _0),
+                    Task(
+                        _unique_dep(c, "ijk"),
                         add,
-                        _unique_dep(b, "ij"),
-                        _unique_dep(b, "jk"),
+                        TaskRef(_unique_dep(b, "ij")),
+                        TaskRef(_unique_dep(b, "jk")),
                     ),
-                    _unique_dep(b, "ij"): (add, _1, _2),
-                    _unique_dep(b, "jk"): (add, _3, _2),
-                },
+                    Task(_unique_dep(b, "ij"), add, _1, _2),
+                    Task(_unique_dep(b, "jk"), add, _3, _2),
+                ),
                 [(123, None), (a, "ij"), (2, None), (a, "jk")],
             ),
         ],
+        # Case 17
         [
             [
-                (b, "jk", {b: (add, _0, _1)}, [(a, "jk"), (2, None)]),
-                (c, "ijk", {c: (add, _0, _1)}, [(b, "ij"), (b, "jk")]),
-                (d, "ijk", {d: (add, _0, _1, _2)}, [(b, "ij"), (c, "ij"), (b, "ij")]),
+                (b, "jk", Task(b, add, _0, _1), [(a, "jk"), (2, None)]),
+                (c, "ijk", Task(c, add, _0, _1), [(b, "ij"), (b, "jk")]),
+                (d, "ijk", Task(d, add, _0, _1, _2), [(b, "ij"), (c, "ij"), (b, "ij")]),
             ],
             (
                 "d",
                 "ijk",
-                {
-                    "d": (
+                Task.fuse(
+                    Task(
+                        d,
                         add,
-                        _unique_dep(b, "ij"),
-                        _unique_dep(c, "ij"),
-                        _unique_dep(b, "ij"),
+                        TaskRef(_unique_dep(b, "ij")),
+                        TaskRef(_unique_dep(c, "ij")),
+                        TaskRef(_unique_dep(b, "ij")),
                     ),
-                    _unique_dep(c, "ij"): (
+                    Task(
+                        _unique_dep(c, "ij"),
                         add,
-                        _unique_dep(b, "ij"),
-                        _unique_dep(b, "jk"),
+                        TaskRef(_unique_dep(b, "ij")),
+                        TaskRef(_unique_dep(b, "jk")),
                     ),
-                    _unique_dep(b, "ij"): (add, _0, _1),
-                    _unique_dep(b, "jk"): (add, _2, _1),
-                },
+                    Task(_unique_dep(b, "ij"), add, _0, _1),
+                    Task(_unique_dep(b, "jk"), add, _2, _1),
+                ),
                 [(a, "ij"), (2, None), (a, "jk")],
             ),
         ],
@@ -296,15 +384,18 @@ def test_rewrite(inputs, expected):
         for inp in inputs
     ]
     result = rewrite_blockwise(inputs)
+
     result2 = (
         result.output,
         "".join(result.output_indices),
-        result.dsk,
+        result.task,
         [
             (name, "".join(ind) if ind is not None else ind)
             for name, ind in result.indices
         ],
     )
+    # Assert on the task first to get a more informative error message
+    assert result2[2] == expected[2]
     assert result2 == expected
 
 
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
index 20f43b89f..c7ecc5115 100644
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -41,7 +41,6 @@ from dask.utils_test import add, hlg_layer, hlg_layer_topological, inc
 dsk: Graph = {("x", 0): (range, 5), ("x", 1): (range, 5), ("x", 2): (range, 5)}
 
 L = list(range(5)) * 3
-
 b = Bag(dsk, "x", 3)
 
 
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index 0a58b90d1..cab7cf422 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -2518,7 +2518,7 @@ def test_getitem_optimization(tmpdir, engine, preserve_index, index):
     subgraph_rd = hlg_layer(dsk, "read-parquet")
     assert isinstance(subgraph_rd, DataFrameIOLayer)
     assert subgraph_rd.columns == ["B"]
-    assert next(iter(subgraph_rd.dsk.values()))[0].columns == ["B"]
+    assert subgraph_rd.task.func.columns == ["B"]
 
     subgraph_wt = hlg_layer(dsk, "to-parquet")
     assert isinstance(subgraph_wt, Blockwise)
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index 66a06c28a..b08718b8e 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -5729,8 +5729,7 @@ def test_simple_map_partitions():
     ddf = dd.from_pandas(df, npartitions=2)
     ddf = ddf.clip(-4, 6)
     task = ddf.__dask_graph__()[ddf.__dask_keys__()[0]]
-    [v] = task[0].dsk.values()
-    assert v[0] == M.clip or v[1] == M.clip
+    assert task.func == M.clip
 
 
 def test_iter():
diff --git a/dask/tests/test_highgraph.py b/dask/tests/test_highgraph.py
index 1021a2284..409ef7a6a 100644
--- a/dask/tests/test_highgraph.py
+++ b/dask/tests/test_highgraph.py
@@ -10,7 +10,7 @@ import pytest
 
 import dask
 from dask.base import tokenize
-from dask.blockwise import Blockwise, blockwise_token
+from dask.blockwise import Blockwise
 from dask.highlevelgraph import HighLevelGraph, Layer, MaterializedLayer, to_graphviz
 from dask.utils_test import inc
 
@@ -264,11 +264,13 @@ def test_blockwise_cull(flat):
 
 
 def test_len_does_not_materialize():
+    from dask._task_spec import Task
+
     a = {"x": 1}
     b = Blockwise(
         output="b",
         output_indices=tuple("ij"),
-        dsk={"b": [[blockwise_token(0)]]},
+        task=Task("b", lambda: "1"),
         indices=(),
         numblocks={},
         new_axes={"i": (1, 1, 1), "j": (1, 1)},
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index 1b6cecd5a..60885ec1f 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -28,6 +28,7 @@ from dask.base import tokenize
 from dask.core import keys_in_tasks, reverse_dict
 from dask.optimization import SubgraphCallable
 from dask.sizeof import sizeof
+from dask.tokenize import tokenize
 from dask.utils import funcname
 
 
@@ -71,9 +72,7 @@ def test_repr():
     assert repr(t) == "<Task 'key' func('a', 'b')>"
 
     t = Task("nested", func2, t, t.ref())
-    assert (
-        repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias(key->key))>"
-    )
+    assert repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias('key'))>"
 
     def long_function_name_longer_even_longer(a, b):
         return a + b
@@ -88,6 +87,33 @@ def test_repr():
     assert repr(t) == "<Task 'kwarg' use_kwargs('foo', kwarg='kwarg_value')>"
 
 
+def test_task_eq():
+    assert Alias(("x", 0, 0)) == Alias(("x", 0, 0))
+    assert Alias(("x", 0, 0)) != Alias(("x", 0, 1))
+
+    assert List(Alias(("x", 0, 0))) == List(Alias(("x", 0, 0)))
+    assert List(Alias(("x", 0, 0))) == List(Alias("a", "a")).substitute(
+        {"a": ("x", 0, 0)}
+    )
+    assert List(Alias(("x", 0, 0)), Alias(("x", 0, 1))) != List(
+        Alias(("x", 0, 0)), Alias(("x", 0, 2))
+    )
+
+    obj = Task(
+        ("z", 0, 0),
+        func,
+        List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),
+        List(TaskRef(("y", 0, 0)), TaskRef(("y", 1, 0))),
+    )
+    assert obj == obj
+    assert obj == Task(
+        ("z", 0, 0),
+        func,
+        List(TaskRef(("x", 0, 0)), TaskRef(("x", 0, 1))),
+        List(TaskRef(("y", 0, 0)), TaskRef(("y", 1, 0))),
+    )
+
+
 def long_function_name_longer_even_longer(a, b):
     return a + b
 
@@ -101,9 +127,7 @@ def test_unpickled_repr():
     assert repr(t) == "<Task 'key' func('a', 'b')>"
 
     t = pickle.loads(pickle.dumps(Task("nested", func2, t, t.ref())))
-    assert (
-        repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias(key->key))>"
-    )
+    assert repr(t) == "<Task 'nested' func2(<Task 'key' func('a', 'b')>, Alias('key'))>"
 
     t = pickle.loads(
         pickle.dumps(Task("long", long_function_name_longer_even_longer, t, t.ref()))
@@ -984,6 +1008,9 @@ def test_nested_containers():
     t = Dict(k=Task("key-1", func, "a", "b"), v=Task("key-2", func, "c", "d"))
     assert not t.dependencies
     assert t() == {"k": "a-b", "v": "c-d"}
+    t2 = Dict({"k": Task("key-1", func, "a", "b"), "v": Task("key-2", func, "c", "d")})
+    assert t == t2
+    assert t() == t2()
     t = Dict(k=Task("key-1", func, "a", TaskRef("b")), v=Task("key-2", func, "c", "d"))
     assert t.dependencies == {"b"}
     assert t({"b": "b"}) == {"k": "a-b", "v": "c-d"}
@@ -999,6 +1026,17 @@ def test_nested_containers():
     assert t.dependencies == {"b"}
     assert t({"b": "b"}) == {"k": "a-b", ("v", 1): "c-d"}
 
+    t = Dict(
+        k=Task("key-1", func, "a", "b"),
+        v=Task("key-2", func, "c", "d"),
+    )
+    t2 = Dict(
+        v=Task("key-2", func, "c", "d"),
+        k=Task("key-1", func, "a", "b"),
+    )
+    assert t == t2
+    assert tokenize(t) == tokenize(t2)
+
 
 def test_block_io_fusion():
 
@@ -1035,6 +1073,39 @@ def test_nested_containers_different_types(task_type, python_type):
     assert t({"b": "b"}) == python_type(("a-b", "c-d"))
 
 
+def test_substitute():
+    t1 = Task("key-1", func, TaskRef("a"), "b")
+    assert t1.substitute({"a": "a"}) is t1
+    assert t1.substitute({"a": "a"}, key="foo") is not t1
+    t2 = t1.substitute({"a": "c"})
+    assert t2 is not t1
+    assert t2.dependencies == {"c"}
+
+    assert t1({"a": "a"}) == t2({"c": "a"})
+
+
+def test_substitute_nested():
+    def func(alist):
+        return alist[0] + alist[1]["foo"]
+
+    t1 = Task(
+        "key-1",
+        func,
+        List(
+            TaskRef("a"),
+            Dict(
+                {
+                    "foo": TaskRef("b"),
+                }
+            ),
+        ),
+    )
+    assert t1.dependencies == {"a", "b"}
+    t2 = t1.substitute({"a": "c", "b": "d"})
+    assert t2.dependencies == {"c", "d"}
+    assert t1({"a": "a", "b": "b"}) == t2({"c": "a", "d": "b"})
+
+
 @pytest.mark.parametrize("Container", [Dict, List, Set, Tuple])
 def test_nested_containers_empty(Container):
     assert Container(Container.klass())() == Container.klass()