Method call

backends/mlx/passes.py:172–209

(self, graph_module: GraphModule)

    """

    def call(self, graph_module: GraphModule) -> PassResult:
        graph = graph_module.graph
        modified = False

        for node in list(graph.nodes):
            match = RMSNormMatch.maybe_create(node)
            if match is None:
                continue

            # Get input shape for normalized_shape
            input_meta = match.input_node.meta.get("val")
            if input_meta is None:
                continue

            # Create fused rms_norm node
            with graph.inserting_before(node):
                normalized_shape = [input_meta.shape[-1]]
                rms_norm_node = graph.call_function(
                    torch.ops.aten.rms_norm.default,
                    args=(
                        match.input_node,
                        normalized_shape,
                        match.weight_node,
                        match.eps,
                    ),
                )
                rms_norm_node.meta = node.meta.copy()

            node.replace_all_uses_with(rms_norm_node)
            match.remove_body_nodes(graph)
            graph.erase_node(node)
            modified = True

        if modified:
            graph.eliminate_dead_code()
            graph.lint()

        return PassResult(graph_module, modified)


class CanonicalizePermutePass(ExportPass):