def cull(dsk, keys):
    """Return new task graph with only the tasks required to calculate keys.

    In other words, remove unnecessary tasks from task graph.
    ``keys`` may be a single key or list of keys.

    Examples
    --------
    >>> from flash.core.serve.dag.utils_test import add, inc
    >>> d = {'x': 1, 'y': (inc, 'x'), 'out': (add, 'x', 10)}
    >>> dsk, dependencies = cull(d, 'out')
    >>> dsk  # doctest: +ELLIPSIS
    {'out': (<function add at ...>, 'x', 10), 'x': 1}
    >>> dependencies
    {'out': ['x'], 'x': []}

    Returns
    -------
    dsk: culled graph
    dependencies: Dict mapping {key: [deps]}.  Useful side effect to accelerate
        other optimizations, notably fuse.

    """
    if not isinstance(keys, (list, set)):
        keys = [keys]

    seen = set()
    dependencies = {}
    out = {}
    work = list(set(flatten(keys)))

    while work:
        new_work = []
        for k in work:
            dependencies_k = get_dependencies(dsk, k, as_list=True)  # fuse needs lists
            out[k] = dsk[k]
            dependencies[k] = dependencies_k
            for d in dependencies_k:
                if d not in seen:
                    seen.add(d)
                    new_work.append(d)

        work = new_work

    return out, dependencies
Return new task graph with only the tasks required to calculate keys.

    In other words, remove unnecessary tasks from task graph.
    ``keys`` may be a single key or list of keys.

    Examples
    --------
    >>> from flash.core.serve.dag.utils_test import add, inc
    >>> d = {'x': 1, 'y': (inc, 'x'), 'out': (add, 'x', 10)}
    >>> dsk, dependencies = cull(d, 'out')
    >>> dsk  # doctest: +ELLIPSIS
    {'out': (<function add at ...>, 'x', 10), 'x': 1}
    >>> dependencies
    {'out': ['x'], 'x': []}

    Returns
    -------
    dsk: culled graph
    dependencies: Dict mapping {key: [deps]}.  Useful side effect to accelerate
        other optimizations, notably fuse.

    
src/flash/core/serve/dag/optimization.py