Function inline_functions

dask/optimization.py:307–369

Inline cheap functions into larger operations Examples -------- >>> inc = lambda x: x + 1 >>> add = lambda x, y: x + y >>> double = lambda x: x * 2 >>> dsk = {'out': (add, 'i', 'd'), # doctest: +SKIP ... 'i': (inc, 'x'), ... 'd': (double, 'y'), ...

(
    dsk, output, fast_functions=None, inline_constants=False, dependencies=None
)

def inline_functions(
    dsk, output, fast_functions=None, inline_constants=False, dependencies=None
):
    """Inline cheap functions into larger operations

    Examples
    --------
    >>> inc = lambda x: x + 1
    >>> add = lambda x, y: x + y
    >>> double = lambda x: x * 2
    >>> dsk = {'out': (add, 'i', 'd'),  # doctest: +SKIP
    ...        'i': (inc, 'x'),
    ...        'd': (double, 'y'),
    ...        'x': 1, 'y': 1}
    >>> inline_functions(dsk, [], [inc])  # doctest: +SKIP
    {'out': (add, (inc, 'x'), 'd'),
     'd': (double, 'y'),
     'x': 1, 'y': 1}

    Protect output keys.  In the example below ``i`` is not inlined because it
    is marked as an output key.

    >>> inline_functions(dsk, ['i', 'out'], [inc, double])  # doctest: +SKIP
    {'out': (add, 'i', (double, 'y')),
     'i': (inc, 'x'),
     'x': 1, 'y': 1}
    """
    if not fast_functions:
        return dsk

    output = set(output)

    fast_functions = set(fast_functions)

    if dependencies is None:
        dependencies = {k: get_dependencies(dsk, k) for k in dsk}
    dependents = reverse_dict(dependencies)

    def inlinable(key, task):
        if (
            not isinstance(task, GraphNode)
            and istask(task)
            and key not in output
            and dependents[key]
        ):
            try:
                if functions_of(task).issubset(fast_functions) and not any(
                    isinstance(dsk[d], GraphNode) for d in dependents[key]
                ):
                    return True
            except TypeError:
                pass
        return False

    keys = [k for k, v in dsk.items() if inlinable(k, v)]

    if keys:
        dsk = inline(