AI-generated review comment - Finding 1: Guard newly MRO-sourced functions

This new MRO-derived fn_source is passed to UserFunctionVariable directly. That source lets bind_args reconstruct closure cells, but it does not guard that the function at the source is still the one Dynamo traced. Rebinding Helper.value therefore reuses the cached graph: the old return value is retained and the old mutation is replayed into the replacement function's closure. The same unguarded construction occurs in UserDefinedObjectVariable.call_method (including __init__) and both context-manager entry/exit paths.

Existing-work / severity: PARTIALLY_TRACKED. #197860 and #198124 cover the ordinary resolve_type_attr bound-method case, so that subcase is existing and its severity is lowered. This finding is retained for the untracked PR-specific delta: #198055 newly makes MRO-derived sources reachable in call_method/__init__ and context-manager __enter__/__exit__ without installing the function guard.

Reproduction:

import torch
from torch._dynamo.testing import CompileCounterWithBackend
from torch.testing._internal.common_utils import run_tests, TestCase


class TestMethodRebind(TestCase):
    def test_locally_instantiated_method_rebind(self):
        old_log = []
        new_log = []

        class Helper:
            def value(self):
                old_log.append("old")
                return 1

        def new_value(self):
            new_log.append("new")
            return 2

        def fn(x):
            return x + Helper().value()

        backend = CompileCounterWithBackend("eager")
        compiled = torch.compile(fn, backend=backend, fullgraph=True)
        out1 = compiled(torch.ones(1))
        Helper.value = new_value
        out2 = compiled(torch.ones(1))

        self.assertEqual(out1, torch.tensor([2.0]))
        self.assertEqual(out2, torch.tensor([3.0]))
        self.assertEqual(old_log, ["old"])
        self.assertEqual(new_log, ["new"])
        self.assertEqual(backend.frame_count, 2)


if __name__ == "__main__":
    run_tests()

Observed on the exact current base and head:

base 4eab5fba: out1=[2] out2=[2] old_log=[]      new_log=[]      frames=1
head fc1f8d74: out1=[2] out2=[2] old_log=['old'] new_log=['old'] frames=1
guarded fix:  out1=[2] out2=[3] old_log=['old'] new_log=['new'] frames=2

The head result is particularly problematic: it silently applies the traced function's mutation to the replacement function's closure. I also reproduced this for __init__ and __enter__ rebinding.

Suggested fix:

Use UserFunctionVariable.create_with_source whenever the function source came from the class MRO, retaining the existing fallback when no MRO source is available:

# UserDefinedObjectVariable.call_method
if source_fn is not None:
    fn_vt = variables.UserFunctionVariable.create_with_source(method, source_fn)
else:
    fn_vt = variables.UserFunctionVariable(
        method, source=source and AttrSource(source, "__func__")
    )
return UserMethodVariable(fn_vt, self, source=source).call_function(
    tx, args, kwargs
)

# UserDefinedObjectVariable.resolve_type_attr
if var_source is not None:
    fn_vt = variables.UserFunctionVariable.create_with_source(
        type_attr, var_source
    )
else:
    fn_vt = variables.UserFunctionVariable(type_attr, source=fn_source)
return variables.UserMethodVariable(fn_vt, self, source=source)

# GenericContextWrappingVariable.enter/exit
if fn_source is not None:
    fn_vt = variables.UserFunctionVariable.create_with_source(fn, fn_source)
else:
    fn_vt = variables.UserFunctionVariable(fn, source=fn_source)

create_with_source installs the required CLOSURE_MATCH guard. PR #198124 covers the ordinary resolve_type_attr site when composed with this change, but not call_method or context-manager entry/exit, so those sites still need to be handled here or coordinated explicitly.

Verification used source-spliced worktrees at exact base 4eab5fba (then-current origin/main, gap 0) and head fc1f8d74; the only environment adaptation was a compatibility-only HealthCheckHook import shim. All five tests added by this PR still passed with the guarded patch.

AI-generated review comment - Finding 2: Recover sources for lru_cache and __getattr__ methods

Widening can_use_mro_source fixes the plain-function branch, but two adjacent method paths still ignore it. The lru_cache branch passes the incoming instance source, which is None for a locally instantiated object, and call_getattr_fallback derives __getattr__ only from self.source, which is also None. Their closure cells therefore remain sourceless and their mutations are silently dropped even though both functions compile with fullgraph=True.

Existing-work / severity: EXISTING-IN-SCOPE; lowered. Both reproductions fail identically on the base and head, so this is not a regression introduced by #198055. It is an incomplete-fix gap within #197896: that issue tracks dropped closure mutations broadly, but no live issue or PR separately covers the lru_cache and custom __getattr__ variants.

Reproduction (lru_cache(0) disables caching so this only tests closure handling):

import functools

import torch
from torch._dynamo.testing import CompileCounter
from torch.testing._internal.common_utils import run_tests, TestCase


class TestMissingMethodSources(TestCase):
    def test_lru_cache_method_closure_mutation(self):
        log = []

        class Helper:
            @functools.lru_cache(0)
            def touch(self):
                log.append("touch")

        def fn(x):
            Helper().touch()
            return x + 1

        counter = CompileCounter()
        compiled = torch.compile(fn, backend=counter, fullgraph=True)
        self.assertEqual(compiled(torch.ones(1)), torch.tensor([2.0]))
        self.assertEqual(compiled(torch.ones(1)), torch.tensor([2.0]))
        self.assertGreater(counter.frame_count, 0)
        self.assertEqual(log, ["touch", "touch"])

    def test_getattr_closure_mutation(self):
        log = []

        class Helper:
            def __getattr__(self, name):
                log.append(name)
                return 1

        def fn(x):
            return x + Helper().missing

        counter = CompileCounter()
        compiled = torch.compile(fn, backend=counter, fullgraph=True)
        self.assertEqual(compiled(torch.ones(1)), torch.tensor([2.0]))
        self.assertEqual(compiled(torch.ones(1)), torch.tensor([2.0]))
        self.assertGreater(counter.frame_count, 0)
        self.assertEqual(log, ["missing", "missing"])


if __name__ == "__main__":
    run_tests()

Exact base and head both produced:

GETATTR frames=1 log=[]
LRU     frames=1 log=[]
FAILED (failures=2)

Suggested fix:

Give the wrapper's underlying function the MRO source while preserving the bound-method source. For __getattr__, keep the bound-method and function sources separate and use create_with_source for the MRO-derived function so this does not introduce Finding 1's rebinding bug:

elif is_lru_cache_wrapped_function(type_attr):
    fn_source = source
    if can_use_mro_source:
        fn_source = self.get_source_by_walking_mro(tx, name)
    fn_vt = variables.WrapperUserFunctionVariable(
        type_attr, "__wrapped__", source=fn_source
    )
    return variables.WrapperUserMethodVariable(fn_vt, self, source=source)

method_source = None
if self.source:
    method_source = AttrSource(self.source, "__getattr__")

if self.cls_source is not None:
    fn_source = self.get_source_by_walking_mro(tx, "__getattr__")
    fn_vt = variables.UserFunctionVariable.create_with_source(
        getattr_fn, fn_source
    )
else:
    fn_source = method_source and AttrSource(method_source, "__func__")
    fn_vt = variables.UserFunctionVariable(getattr_fn, source=fn_source)

out = variables.UserMethodVariable(
    fn_vt, self, source=method_source
).call_function(tx, [variables.ConstantVariable.create(name)], {})

With this patch, the same test produced:

GETATTR frames=2 log=['missing', 'missing']
LRU     frames=2 log=['touch', 'touch']
OK

Verification used source-spliced worktrees at exact base 4eab5fba and head fc1f8d74; the only environment adaptation was the compatibility-only import shim noted in Finding 1. The PR's five new tests also passed with this patch.

AI-generated review comment - Finding 3: Separate redundant build-class coverage from missing supported-path coverage

There are two independent coverage concerns here:

1. test_build_class_in_compiled_fn_nonlocal_rebind exercises behavior already covered more precisely by test_misc.py::MiscTests::test_build_class_closure_rebinding, which asserts that rebinding a materialized sourceless build-class cell remains unsupported under fullgraph=True. Because this new test omits fullgraph=True, it succeeds through graph-break/eager fallback and does not exercise this PR's production change. It can be dropped as redundant.
2. Separately, the supported locally defined class path still lacks coverage for nonlocal cell reassignment. The existing tests and Findings 1/2 cover container mutation, function-identity guarding, and the lru_cache/__getattr__ branches; none covers STORE_DEREF through an ordinary method whose class is defined outside the compiled function.

The distinction is observable on the exact current base and head:

added build-class test:                 base PASS, head PASS
added build-class test + fullgraph=True: base Write to immutable cell,
                                          head Write to immutable cell
supported-path test below:              base Write to immutable cell,
                                          head PASS (n == 2)

Suggested fix:

Drop the newly added fallback-only build-class test because that unsupported path is already pinned elsewhere. Independently add the following supported-path regression test alongside the other new closure-mutation tests.

Suggested test:

def test_locally_defined_method_nonlocal_rebind(self):
    n = 0

    class Helper:
        def touch(self):
            nonlocal n
            n += 1

    def fn(x):
        Helper().touch()
        return x + 1

    compiled = torch.compile(fn, backend="eager", fullgraph=True)
    self.assertEqual(compiled(torch.ones(1)), torch.tensor([2.0]))
    self.assertEqual(compiled(torch.ones(1)), torch.tensor([2.0]))
    self.assertEqual(n, 2)

This is not a semantic replacement for the existing build-class test: one confirms the unsupported build-class boundary, while the new test fills missing coverage for the supported ordinary-method path.