Function walk_back

backends/mlx/pattern_utils.py:196–295

Walk backwards through a chain of ops, matching against a pattern. Starting from *node*, try to match each step against the current node. At every matched step the walk advances to ``cur.args[step.arg_index]``. Optional steps are silently skipped when they don't match. Repeat steps

(  # noqa: C901
    node: Node,
    steps: List[OpStep],
    debug: bool = False,
)

def walk_back(  # noqa: C901
    node: Node,
    steps: List[OpStep],
    debug: bool = False,
) -> Optional[Tuple[Node, List[WalkBackEntry]]]:
    """
    Walk backwards through a chain of ops, matching against a pattern.

    Starting from *node*, try to match each step against the current node.
    At every matched step the walk advances to ``cur.args[step.arg_index]``.
    Optional steps are silently skipped when they don't match. Repeat steps
    match 0 or more times.

    Args:
        node: Starting node
        steps: List of OpStep to match in order

    Returns:
        ``(base_node, entries)`` if the full chain matches, else ``None``.
        *base_node* is the input to the first (deepest) op in the chain.
        *entries* is a list with one entry per OpStep:
            - Node: matched node (for regular steps)
            - None: optional step that didn't match
            - List[Node]: repeat step (0 or more matches)

    Examples:
        # Match: rsqrt(add(mean(pow(x, 2)), eps))
        result = walk_back(rsqrt_node, [
            OpStep(op=torch.ops.aten.rsqrt.default),
            OpStep(op=torch.ops.aten.add.Tensor),
            OpStep(op=torch.ops.aten.mean.dim),
            OpStep(op=torch.ops.aten.pow.Tensor_Scalar),
        ])
        if result:
            base, entries = result
            rsqrt, add, mean, pow = entries  # Each is a Node

        # Match chain of select ops (like tensor[0][0])
        result = walk_back(node, [
            OpStep(op=torch.ops.aten.select.int, repeat=True),
        ])
        if result:
            base, entries = result
            select_nodes = entries[0]  # List[Node], may be empty

        # Skip optional _to_copy, then match rsqrt
        result = walk_back(node, [
            OpStep(op=torch.ops.aten._to_copy.default, optional=True),
            OpStep(op=torch.ops.aten.rsqrt.default),
        ])
        if result:
            base, entries = result
            to_copy, rsqrt = entries  # to_copy may be None
    """
    entries: List[WalkBackEntry] = []
    cur = node

    for i, step in enumerate(steps):