Class OpStep

backends/mlx/pattern_utils.py:105–193

One step in a backward walk through the graph. Used with walk_back() to define pattern chains. Supports both exact op matching and predicate-based matching. Attributes: op: Specific op to match (e.g., torch.ops.aten.rsqrt.default) predicate: Alternative to op - a f

@dataclass
class OpStep:
    """
    One step in a backward walk through the graph.

    Used with walk_back() to define pattern chains. Supports both exact op
    matching and predicate-based matching.

    Attributes:
        op: Specific op to match (e.g., torch.ops.aten.rsqrt.default)
        predicate: Alternative to op - a function that returns True for matching nodes
        optional: If True, skip this step if it doesn't match
        repeat: If True, match this step 0 or more times (like regex *)
        require_single_user: If True (default), only match nodes with exactly one user
        nargs: Number of args required. Can be:
               - int: minimum number of args (default 1, since we advance via args[0])
               - tuple (min, max): range of args required (inclusive)
        kwargs: Set of kwargs we handle (node's kwargs must be subset of this)
        arg_index: Which arg to follow when advancing (default 0)

    Examples:
        # Match specific op
        OpStep(op=torch.ops.aten.rsqrt.default)

        # Match with predicate (for matching families of ops)
        OpStep(predicate=lambda n: match_target(n, torch.ops.aten.select.int))

        # Match chain of same op type (0 or more)
        OpStep(op=torch.ops.aten.select.int, repeat=True)

        # Optional dtype conversion
        OpStep(op=torch.ops.aten._to_copy.default, optional=True)

        # Require between 2 and 4 args
        OpStep(op=torch.ops.aten.some_op.default, nargs=(2, 4))

        # Declare that we handle 'dtype' kwarg
        OpStep(op=torch.ops.aten._to_copy.default, kwargs={"dtype"})

        # Follow second arg (e.g., mul(x, rsqrt(y)) -> follow rsqrt in args[1])
        OpStep(op=torch.ops.aten.mul.Tensor, arg_index=1)
    """

    op: Any = None
    predicate: Optional[Callable[[Node], bool]] = None
    optional: bool = False
    repeat: bool = False
    require_single_user: bool = True
    nargs: Union[int, Tuple[int, int]] = 1
    kwargs: Set[str] = field(default_factory=set)  # Empty = no kwargs allowed
    arg_index: int = 0

    def matches(self, node: Node) -> bool:
        """Check if this step fully matches the given node."""
        # Check op or predicate
        if self.op is not None:
            if not match_target(node, self.op):
                return False
        elif self.predicate is not None: