import torch
import torch.fx
from torch._inductor.pattern_matcher import Arg, Match, MatchContext


def test_nested_boundary_stopping_logic():
    class M(torch.nn.Module):
        def forward(self, x, y):
            s1 = torch.sin(x)
            s2 = torch.sin(y)
            c = torch.cat([s1, s2], dim=1)
            return torch.cos(c)

    traced = torch.fx.symbolic_trace(M())
    graph = traced.graph
    nodes_before = set(graph.nodes)

    call_nodes = [n for n in graph.nodes if n.op == "call_function"]
    sin_nodes = [n for n in call_nodes if n.target == torch.sin]
    assert len(sin_nodes) == 2
    node_s1, node_s2 = sin_nodes
    node_cat = next(n for n in call_nodes if n.target == torch.cat)

    node_cat.meta["recompute"] = "1"
    node_cat.meta["ac_graph_id"] = 1

    def rep_fn(nested):
        x = nested[0]
        y = nested[1]
        return torch.add(x, y)

    replacement_gm = torch.fx.symbolic_trace(rep_fn)

    pattern = Arg()
    ctx = MatchContext([pattern], graph=graph)
    ctx.pattern_to_node[pattern] = node_cat
    match = Match(ctx, pattern)

    match.replace_with_graph(replacement_gm, args=[(node_s1, node_s2)])

    nodes_after = [n for n in graph.nodes if n not in nodes_before and n.op != "output"]
    new_add_nodes = [
        n for n in nodes_after if n.op == "call_function" and n.target == torch.add
    ]
    assert len(new_add_nodes) == 1
    new_add_node = new_add_nodes[0]

    assert new_add_node.meta.get("recompute") == "1"
    assert new_add_node.meta.get("ac_graph_id") == 1

    assert "recompute" not in node_s1.meta
    assert "recompute" not in node_s2.meta
    assert "ac_graph_id" not in node_s1.meta
    assert "ac_graph_id" not in node_s2.meta

    placeholders = [n for n in graph.nodes if n.op == "placeholder"]
    for n in placeholders:
        assert "recompute" not in n.meta
        assert "ac_graph_id" not in n.meta


if __name__ == "__main__":
    test_nested_boundary_stopping_logic()

### Error logs

Running the minimal reproducer fails with:

Traceback (most recent call last):
  File "/home/tiger/pattern-matcher-issue-repro/repro.py", line 70, in <module>
    test_nested_boundary_stopping_logic()
  File "/home/tiger/pattern-matcher-issue-repro/repro.py", line 58, in test_nested_boundary_stopping_logic
    assert "recompute" not in node_s1.meta
AssertionError

### Versions

PyTorch version: 2.7.1
Is debug build: False
CUDA used to build PyTorch: 12.6
Is CUDA available: False
Python version: 3.11.2 (main, Apr 28 2025, 14:11:48) [GCC 12.2.0]
Python platform: Linux-5.4.143.bsk.7-amd64-x86_64-with-glibc2.36

cc @soulitzer @ezyang @EikanWang @jgong5 @wenzhe-nrv @chauhang @penguinwu @voznesenskym @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo