Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/__init__.py

@@ -0,0 +1,263 @@
+"""
+This package implements variable tracking and symbolic execution capabilities for Dynamo,
+which are essential for converting Python code into FX graphs. It provides a comprehensive
+set of variable types that handle different Python constructs during tracing.
+
+Each variable type (like BuiltinVariable, TensorVariable, NNModuleVariable, etc.) is responsible
+for tracking and symbolically executing operations on specific Python objects. This enables
+Dynamo to:
+- Track the flow of values through Python code
+- Maintain correct semantics during graph conversion
+- Handle complex Python features like context managers, iterators, and custom objects
+- Support both eager and symbolic execution modes
+
+The VariableTracker base class provides the foundation for all variable types, with each
+subclass implementing specific behavior for different Python constructs. This modular design
+allows Dynamo to accurately trace and optimize Python code while preserving its semantics.
+"""
+
+from .base import VariableTracker
+from .builtin import (
+    BaseBuiltinVariable,
+    BuiltinVariable,
+    DictBuiltinVariable,
+    IterBuiltinVariable,
+    ListBuiltinVariable,
+)
+from .constant import ConstantVariable
+from .ctx_manager import (
+    CatchWarningsCtxManagerVariable,
+    ContextWrappingVariable,
+    CUDADeviceVariable,
+    CudagraphOverrideVariable,
+    DisabledSavedTensorsHooksVariable,
+    DualLevelContextManager,
+    DynamoConfigPatchVariable,
+    ErrorOnGraphBreakVariable,
+    FSDPParamGroupUseTrainingStateVariable,
+    FxTracebackAnnotateVariable,
+    GenericContextWrappingVariable,
+    GradIncrementNestingCtxManagerVariable,
+    GradInplaceRequiresGradCtxManagerVariable,
+    GradModeVariable,
+    InferenceModeVariable,
+    JvpIncrementNestingCtxManagerVariable,
+    SDPAKernelVariable,
+    SetFwdGradEnabledContextManager,
+    TemporarilyPopInterpreterStackCtxManagerVariable,
+    VmapIncrementNestingCtxManagerVariable,
+    WithEnterFunctionVariable,
+    WithExitFunctionVariable,
+)
+from .dicts import (
+    ConstDictVariable,
+    DefaultDictVariable,
+    DictItemsVariable,
+    DunderDictVariable,
+    MappingProxyVariable,
+    NNModuleHooksDictVariable,
+)
+from .distributed import BackwardHookVariable, DistributedVariable
+from .functions import (
+    BaseUserFunctionVariable,
+    BuiltinMethodVariable,
+    CollectionsNamedTupleFunction,
+    CreateTMADescriptorExperimentalVariable,
+    CreateTMADescriptorStableVariable,
+    FunctionDecoratedByContextlibContextManagerVariable,
+    FunctoolsPartialVariable,
+    InspectSignatureVariable,
+    LocalGeneratorFunctionVariable,
+    LocalGeneratorObjectVariable,
+    NestedUserFunctionVariable,
+    PolyfilledFunctionVariable,
+    PyTreeGetNodeTypeFunctionVariable,
+    PyTreeTreeIsLeafFunctionVariable,
+    SkipFunctionVariable,
+    SparseTensorCreationSkipVariable,
+    TMADescriptorExperimentalVariable,
+    TMADescriptorStableVariable,
+    TritonSetAllocatorVariable,
+    UserFunctionVariable,
+    UserMethodVariable,
+    WrapperUserFunctionVariable,
+    WrapperUserMethodVariable,
+)
+from .higher_order_ops import (
+    FunctionalCallVariable,
+    FunctorchHigherOrderVariable,
+    ReparametrizeModuleCallVariable,
+    TorchHigherOrderOperatorVariable,
+)
+from .iter import (
+    CountIteratorVariable,
+    FilterVariable,
+    IteratorVariable,
+    ItertoolsVariable,
+    MapVariable,
+    ObjectIteratorVariable,
+    RepeatIteratorVariable,
+    ZipVariable,
+)
+from .lazy import LazyConstantVariable, LazyVariableTracker
+from .lists import (
+    BaseListVariable,
+    ListIteratorVariable,
+    ListVariable,
+    RangeVariable,
+    SliceVariable,
+    TupleIteratorVariable,
+    TupleVariable,
+)
+from .misc import (
+    AutogradFunctionContextVariable,
+    AutogradFunctionVariable,
+    CellVariable,
+    DeletedVariable,
+    ExceptionVariable,
+    GetAttrVariable,
+    LambdaVariable,
+    MethodWrapperVariable,
+    NewGlobalVariable,
+    NumpyVariable,
+    ObjectVariable,
+    PythonModuleVariable,
+    RandomClassVariable,
+    RandomVariable,
+    StringFormatVariable,
+    SuperVariable,
+    TorchVersionVariable,
+    TracebackVariable,
+    TypingVariable,
+    UnknownVariable,
+    WeakRefVariable,
+)
+from .nn_module import (
+    FSDPManagedNNModuleVariable,
+    NNModuleVariable,
+    UnspecializedBuiltinNNModuleVariable,
+    UnspecializedNNModuleVariable,
+)
+from .optimizer import OptimizerVariable
+from .sdpa import SDPAParamsVariable
+from .sets import (
+    DictKeySetVariable,
+    FrozensetVariable,
+    OrderedSetClassVariable,
+    OrderedSetVariable,
+    SetVariable,
+)
+from .streams import (
+    CudaStreamVariable,
+    EventVariable,
+    StreamContextVariable,
+    StreamVariable,
+)
+from .tensor import (
+    DataPtrVariable,
+    FakeItemVariable,
+    NumpyNdarrayVariable,
+    SymNodeVariable,
+    TensorVariable,
+    UnspecializedPythonVariable,
+    UntypedStorageVariable,
+)
+from .torch import TorchCtxManagerClassVariable, TorchInGraphFunctionVariable
+from .user_defined import (
+    FrozenDataClassVariable,
+    InspectVariable,
+    MutableMappingVariable,
+    NamedTupleVariable,
+    RemovableHandleVariable,
+    StructSequenceVariable,
+    UserDefinedClassVariable,
+    UserDefinedConstantVariable,
+    UserDefinedDictVariable,
+    UserDefinedExceptionClassVariable,
+    UserDefinedExceptionObjectVariable,
+    UserDefinedListVariable,
+    UserDefinedObjectVariable,
+    UserDefinedSetVariable,
+    UserDefinedTupleVariable,
+    UserDefinedVariable,
+)
+
+
+__all__ = [
+    "AutogradFunctionContextVariable",
+    "AutogradFunctionVariable",
+    "BackwardHookVariable",
+    "BaseBuiltinVariable",
+    "BaseListVariable",
+    "BuiltinVariable",
+    "CatchWarningsCtxManagerVariable",
+    "ConstantVariable",
+    "ConstDictVariable",
+    "DictBuiltinVariable",
+    "ContextWrappingVariable",
+    "CountIteratorVariable",
+    "CreateTMADescriptorExperimentalVariable",
+    "CreateTMADescriptorStableVariable",
+    "CUDADeviceVariable",
+    "CudagraphOverrideVariable",
+    "DataPtrVariable",
+    "DefaultDictVariable",
+    "DeletedVariable",
+    "DictKeySetVariable",
+    "DynamoConfigPatchVariable",
+    "FakeItemVariable",
+    "GetAttrVariable",
+    "GradModeVariable",
+    "InspectSignatureVariable",
+    "InspectVariable",
+    "IterBuiltinVariable",
+    "IteratorVariable",
+    "ItertoolsVariable",
+    "LambdaVariable",
+    "LazyConstantVariable",
+    "LazyVariableTracker",
+    "ListBuiltinVariable",
+    "ListIteratorVariable",
+    "ListVariable",
+    "NestedUserFunctionVariable",
+    "CellVariable",
+    "NewGlobalVariable",
+    "NNModuleVariable",
+    "NumpyNdarrayVariable",
+    "NumpyVariable",
+    "OptimizerVariable",
+    "PolyfilledFunctionVariable",
+    "PythonModuleVariable",
+    "RangeVariable",
+    "RemovableHandleVariable",
+    "RepeatIteratorVariable",
+    "SDPAParamsVariable",
+    "ErrorOnGraphBreakVariable",
+    "SkipFunctionVariable",
+    "SliceVariable",
+    "StringFormatVariable",
+    "SuperVariable",
+    "TemporarilyPopInterpreterStackCtxManagerVariable",
+    "TensorVariable",
+    "TMADescriptorExperimentalVariable",
+    "TMADescriptorStableVariable",
+    "TorchCtxManagerClassVariable",
+    "TorchInGraphFunctionVariable",
+    "TorchVersionVariable",
+    "TupleVariable",
+    "UnknownVariable",
+    "UnspecializedNNModuleVariable",
+    "UnspecializedPythonVariable",
+    "UntypedStorageVariable",
+    "UserDefinedClassVariable",
+    "UserDefinedTupleVariable",
+    "NamedTupleVariable",
+    "StructSequenceVariable",
+    "UserDefinedObjectVariable",
+    "UserFunctionVariable",
+    "UserMethodVariable",
+    "VariableTracker",
+    "WithEnterFunctionVariable",
+    "WithExitFunctionVariable",
+    "MappingProxyVariable",
+]

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/constant.py

@@ -0,0 +1,482 @@
+"""
+Constant variable tracking in Dynamo.
+
+This module is fundamental to Dynamo's ability to track and propagate constant
+values during compilation, ensuring proper handling of Python literals and
+maintaining type safety through the compilation process.
+"""
+
+from __future__ import annotations
+
+import operator
+from typing import Any, Literal, overload, TYPE_CHECKING
+from typing_extensions import override
+
+import torch
+from torch._dynamo.source import GetItemSource
+
+from .. import variables
+from ..exc import raise_observed_exception, unimplemented
+from ..utils import common_constant_types, istype, np, raise_args_mismatch
+from .base import ValueMutationNew, VariableTracker
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+    from .functions import UserFunctionVariable
+
+
+class ConstantVariable(VariableTracker):
+    """
+    Variable tracker for Python literals and basic immutable types, with automatic
+    routing support for collection types (lists, tuples, sets, etc.).
+
+    The create() method intelligently constructs appropriate variable types for
+    nested collections.
+    """
+
+    # PyLong_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/longobject.c#L6585
+    # PyFloat_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/floatobject.c#L1880
+    # PyBool_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/boolobject.c#L171
+    # PyUnicode_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/unicodeobject.c#L14931
+    # PyBytes_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/bytesobject.c#L3017
+    # PyComplex_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/complexobject.c#L1099
+    # _PyNone_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/object.c#L2022
+    _cpython_type = (int, float, str, bytes, bool, type(None), complex, type(...))
+
+    @overload
+    @staticmethod
+    def create(value: None) -> ConstantVariable: ...
+
+    @overload
+    @staticmethod
+    def create(value: bool) -> ConstantVariable: ...
+
+    @overload
+    @staticmethod
+    def create(value: Any, **kwargs: Any) -> VariableTracker: ...
+
+    @staticmethod
+    def create(value: Any, **kwargs: Any) -> VariableTracker:
+        """
+        Create a `ConstantVariable` based on the given value, and supports
+        automatic routing for collection types like `tuple` (in which case we'd
+        create `ConstantVariable` for the leaf items).
+
+        NOTE: the caller must install the proper guards if needed; most often
+        the guard will be `CONSTANT_MATCH`.
+        """
+        # Return pre-allocated sentinels for None/True/False when there are
+        # no extra kwargs (source, etc.) that would differentiate the instance.
+        if not kwargs:
+            match value:
+                case None:
+                    return CONSTANT_VARIABLE_NONE
+                case True:
+                    return CONSTANT_VARIABLE_TRUE
+                case False:
+                    return CONSTANT_VARIABLE_FALSE
+
+        source = kwargs.get("source")
+
+        # Routing for supported collection literals.
+        if isinstance(value, set):
+            items = [ConstantVariable.create(x) for x in value]
+            return variables.SetVariable(items, **kwargs)  # type: ignore[arg-type]
+        elif isinstance(value, frozenset):
+            items = [ConstantVariable.create(x) for x in value]
+            return variables.FrozensetVariable(items, **kwargs)  # type: ignore[arg-type]
+        elif isinstance(value, slice):
+            slice_args = (value.start, value.stop, value.step)
+            slice_args_vars = tuple(ConstantVariable.create(arg) for arg in slice_args)
+            return variables.SliceVariable(slice_args_vars, **kwargs)
+        elif isinstance(value, (list, tuple)):
+            items = []
+            for i, x in enumerate(value):
+                item_source = GetItemSource(source, i) if source else None
+                items.append(
+                    ConstantVariable.create(
+                        x,
+                        source=item_source,
+                    )
+                )
+            return variables.BaseListVariable.cls_for(type(value))(items, **kwargs)
+
+        return ConstantVariable(value, **kwargs)
+
+    def __init__(self, value: Any, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        assert ConstantVariable.is_base_literal(value), f"""
+Cannot construct `ConstantVariable` for value of type {type(value)}.
+
+This failure likely due to PyTorch-internal use of `ConstantVariable` on
+non-literal python values, please try using `VariableTracker.build` instead. If
+you believe it's a necessary and legitimate use case (the value is immutable and
+can't easily be represented with another `VariableTracker` class), please add
+its type to `common_constant_types`.
+"""
+        if np is not None and isinstance(value, np.number):
+            self.value = value.item()
+        else:
+            self.value = value
+
+    def as_proxy(self) -> Any:
+        return self.value
+
+    def __repr__(self) -> str:
+        return f"ConstantVariable({type(self.value).__name__}: {repr(self.value)})"
+
+    def as_python_constant(self) -> Any:
+        return self.value
+
+    def is_python_constant(self) -> Literal[True]:
+        return True
+
+    def is_symnode_like(self) -> bool:
+        return isinstance(self.value, (int, bool))
+
+    def is_constant_match(self, *values: Any) -> bool:
+        return self.value in values
+
+    def is_constant_none(self) -> bool:
+        return self.value is None
+
+    @property
+    def items(self) -> list[VariableTracker]:
+        """
+        Need this when adding a BaseListVariable and a ConstantVariable together.
+        Happens in detectron2.
+        """
+        return self.unpack_var_sequence(tx=None)
+
+    def getitem_const(
+        self, tx: InstructionTranslator, arg: VariableTracker
+    ) -> VariableTracker:
+        return ConstantVariable.create(
+            self.value[arg.as_python_constant()],
+        )
+
+    @staticmethod
+    def is_base_literal(obj: object) -> bool:
+        return type(obj) in common_constant_types
+
+    @staticmethod
+    def is_literal(obj: object, cache: dict[int, object] | None = None) -> bool:
+        if cache is None:
+            cache = {}
+        if id(obj) in cache:
+            # no-op if there is a cyclical reference
+            return True
+        if type(obj) in (list, tuple, set, frozenset, torch.Size):
+            cache[id(obj)] = obj
+            return all(ConstantVariable.is_literal(x, cache) for x in obj)  # type: ignore[attr-defined]
+        return ConstantVariable.is_base_literal(obj)
+
+    def unpack_var_sequence(
+        self, tx: InstructionTranslator | None
+    ) -> list[VariableTracker]:
+        try:
+            return [ConstantVariable.create(x) for x in self.as_python_constant()]
+        except TypeError as e:
+            raise NotImplementedError from e
+
+    def len_impl(self, tx: InstructionTranslator) -> VariableTracker:
+        """Generic len for any constant value (sequence or mapping)."""
+        try:
+            return ConstantVariable.create(len(self.value))
+        except TypeError as e:
+            raise_observed_exception(type(e), tx, args=list(e.args))
+
+    def sq_length(self, tx: InstructionTranslator) -> VariableTracker:
+        """Sequence length - delegates to len_impl for constants."""
+        return self.len_impl(tx)
+
+    def mp_length(self, tx: InstructionTranslator) -> VariableTracker:
+        """Mapping length - delegates to len_impl for constants."""
+        return self.len_impl(tx)
+
+    def const_getattr(self, tx: InstructionTranslator, name: str) -> VariableTracker:
+        if not hasattr(self.value, name):
+            raise_observed_exception(AttributeError, tx, args=[name])
+        member = getattr(self.value, name)
+        if callable(member):
+            raise NotImplementedError
+        return member
+
+    def call_method(
+        self,
+        tx: InstructionTranslator,
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .tensor import SymNodeVariable
+
+        if name == "format" and istype(self.value, str):
+            return variables.BuiltinVariable(str.format).call_function(
+                tx, [self, *args], kwargs
+            )
+        elif name == "join" and istype(self.value, str):
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            arg_unpacked = args[0].force_unpack_var_sequence(tx)
+            try:
+                arg_const = [x.as_python_constant() for x in arg_unpacked]
+                return ConstantVariable.create(self.value.join(arg_const))
+            except NotImplementedError:
+                return super().call_method(tx, name, args, kwargs)
+        elif name == "__iter__" and istype(self.value, str):
+            # this could be some generic iterator to avoid the circular import,
+            # but ListIterator does what we want
+            from .lists import ListIteratorVariable
+
+            return ListIteratorVariable(
+                self.unpack_var_sequence(tx), mutation_type=ValueMutationNew()
+            )
+
+        if any(isinstance(x, SymNodeVariable) for x in args):
+            # Promote to SymNodeVariable for operations involving dynamic shapes.
+            return variables.SymNodeVariable.create(
+                tx, self.as_proxy(), self.value
+            ).call_method(tx, name, args, kwargs)
+
+        try:
+            const_args = [a.as_python_constant() for a in args]
+            const_kwargs = {k: v.as_python_constant() for k, v in kwargs.items()}
+        except NotImplementedError:
+            return super().call_method(tx, name, args, kwargs)
+
+        if isinstance(self.value, str) and name in str.__dict__:
+            method = getattr(self.value, name)
+            try:
+                return ConstantVariable.create(method(*const_args, **const_kwargs))
+            except Exception as e:
+                raise_observed_exception(type(e), tx)
+        elif isinstance(self.value, (float, int)) and hasattr(self.value, name):
+            if not (args or kwargs):
+                try:
+                    return ConstantVariable.create(getattr(self.value, name)())
+                except (OverflowError, ValueError) as exc:
+                    raise_observed_exception(
+                        type(exc),
+                        tx,
+                        args=list(exc.args),
+                    )
+            if (
+                hasattr(operator, name)
+                and len(args) == 1
+                and args[0].is_python_constant()
+            ):
+                add_target = const_args[0]
+                op = getattr(operator, name)
+                if isinstance(
+                    add_target, (torch.SymBool, torch.SymFloat, torch.SymInt)
+                ):
+                    # Addition between a non sym and sym makes a sym
+                    proxy = tx.output.create_proxy(
+                        "call_function", op, (self.value, add_target), {}
+                    )
+                    return SymNodeVariable.create(tx, proxy, add_target)
+                else:
+                    try:
+                        return ConstantVariable.create(op(self.value, add_target))
+                    except Exception as e:
+                        raise_observed_exception(type(e), tx, args=list(e.args))
+        elif isinstance(self.value, bytes) and name == "decode":
+            method = getattr(self.value, name)
+            return ConstantVariable.create(method(*const_args, **const_kwargs))
+        elif type(self.value) is complex and name in complex.__dict__:
+            method = getattr(self.value, name)
+            try:
+                return ConstantVariable.create(method(*const_args, **const_kwargs))
+            except Exception as e:
+                raise_observed_exception(type(e), tx)
+
+        if name == "__round__" and len(args) == 1 and args[0].is_python_constant():
+            try:
+                return ConstantVariable.create(
+                    round(self.value, args[0].as_python_constant())
+                )
+            except Exception as e:
+                raise_observed_exception(type(e), tx, args=list(e.args))
+        elif name == "__contains__" and len(args) == 1 and args[0].is_python_constant():
+            assert not kwargs
+            search = args[0].as_python_constant()
+            try:
+                result = search in self.value
+                return ConstantVariable.create(result)
+            except TypeError as e:
+                raise_observed_exception(type(e), tx, args=list(e.args))
+        return super().call_method(tx, name, args, kwargs)
+
+    def call_tree_map(
+        self,
+        tx: InstructionTranslator,
+        tree_map_fn: UserFunctionVariable,
+        map_fn: VariableTracker,
+        rest: Sequence[VariableTracker],
+        tree_map_kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if self.value is None:
+            none_is_leaf_var = tree_map_kwargs.get("none_is_leaf")
+            if none_is_leaf_var is not None:
+                try:
+                    none_is_leaf = bool(none_is_leaf_var.as_python_constant())
+                except NotImplementedError:
+                    return self._tree_map_fallback(
+                        tx,
+                        tree_map_fn,
+                        map_fn,
+                        rest,
+                        tree_map_kwargs,
+                    )
+            else:
+                tree_map_module = getattr(
+                    getattr(tree_map_fn, "fn", None), "__module__", ""
+                )
+                # torch.utils._pytree and torch.utils._cxx_pytree treat None as a leaf
+                # by default, while optree keeps it as an internal node unless
+                # none_is_leaf=True is provided.
+                none_is_leaf = not tree_map_module.startswith("optree")
+            if none_is_leaf:
+                return map_fn.call_function(tx, [self, *rest], {})
+            else:
+                for other in rest:
+                    if not other.is_constant_none():
+                        return self._tree_map_fallback(
+                            tx,
+                            tree_map_fn,
+                            map_fn,
+                            rest,
+                            tree_map_kwargs,
+                        )
+                return self.clone()
+        if isinstance(self.value, (int, float, bool, complex, str, bytes, torch.dtype)):
+            return map_fn.call_function(tx, [self, *rest], {})
+        return super().call_tree_map(
+            tx,
+            tree_map_fn,
+            map_fn,
+            rest,
+            tree_map_kwargs,
+        )
+
+    @override
+    def call_obj_hasattr(
+        self, tx: InstructionTranslator, name: str
+    ) -> ConstantVariable:
+        result = hasattr(self.value, name)
+        return variables.ConstantVariable.create(result)
+
+    def is_python_hashable(self) -> Literal[True]:
+        return True
+
+    def get_python_hash(self) -> int:
+        return hash(self.value)
+
+    def is_python_equal(self, other: object) -> bool:
+        from .tensor import SymNodeVariable
+
+        if isinstance(other, SymNodeVariable):
+            return self.as_python_constant() == other.evaluate_expr()
+        return (
+            isinstance(other, VariableTracker)
+            and self.as_python_constant() == other.as_python_constant()
+        )
+
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
+    def nb_index_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        # CPython: int and bool define nb_index (returns self for int,
+        # int(self) for bool). All other constant types do not.
+        if isinstance(self.value, (int, bool)):
+            return ConstantVariable.create(operator.index(self.value))
+        return super().nb_index_impl(tx)
+
+    def nb_int_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        # CPython: int defines nb_int (long_long, returns copy).
+        # bool inherits nb_int from int via slot inheritance.
+        # float defines nb_int (truncates toward zero via PyLong_FromDouble).
+        return ConstantVariable.create(int(self.value))
+
+    def nb_float_impl(
+        self,
+        tx: Any,
+    ) -> VariableTracker:
+        # CPython: float defines nb_float (float_float, returns copy).
+        # int defines nb_float (long_float, converts to float).
+        # bool inherits nb_float from int via slot inheritance.
+        return ConstantVariable.create(float(self.value))
+
+
+CONSTANT_VARIABLE_NONE = ConstantVariable(None)
+CONSTANT_VARIABLE_TRUE = ConstantVariable(True)
+CONSTANT_VARIABLE_FALSE = ConstantVariable(False)
+
+
+class FakeIdVariable(VariableTracker):
+    """A compile-time-only id value that can be used as a dict key but cannot
+    be reconstructed across graph breaks.
+
+    When dynamo evaluates ``id(x)`` on a variable tracker that has no
+    corresponding runtime object (e.g. a ``ConstDictVariable`` created during
+    tracing), we mint a fake integer id.  This variable holds that id and
+    supports the minimal interface needed to participate as a dict key
+    (hashing and equality).  It intentionally blocks reconstruction so that a
+    graph break does not silently bake a stale id into the resumed bytecode.
+    """
+
+    # PyLong_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/longobject.c#L6585
+    _cpython_type = int
+
+    def __init__(self, value: int, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.value = value
+
+    def as_python_constant(self) -> int:
+        return self.value
+
+    def is_python_constant(self) -> bool:
+        return False
+
+    def python_type(self) -> type:
+        return int
+
+    def is_python_hashable(self) -> bool:
+        return True
+
+    def get_python_hash(self) -> int:
+        return hash(self.value)
+
+    def is_python_equal(self, other: object) -> bool:
+        if isinstance(other, (FakeIdVariable, ConstantVariable)):
+            return self.value == other.as_python_constant()
+        return False
+
+    def reconstruct(self, codegen: Any) -> None:
+        unimplemented(
+            gb_type="Reconstruction of FakeIdVariable",
+            context=str(self.value),
+            explanation=(
+                "A fake id produced by id() on a compile-time container "
+                "cannot be reconstructed across a graph break."
+            ),
+            hints=[
+                "Avoid using id() on containers in code that may graph-break.",
+            ],
+        )

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/distributed.py

@@ -0,0 +1,258 @@
+"""
+Distributed computing variable tracking classes for PyTorch Dynamo.
+
+This module implements variable tracking for distributed computing components:
+- Process Groups (for collective communication)
+- Device Meshes (for distributed tensor sharding)
+- Placement Types (for specifying distribution strategies)
+- Distributed Tensors and their operations
+- Backward hooks for distributed module operations
+
+These classes are responsible for tracking distributed operations during graph
+compilation while maintaining proper guards and handling distributed-specific
+behaviors. They ensure correct handling of distributed components like process
+groups, device meshes, and placement strategies while preserving proper semantics
+for distributed tensor operations in the compiled code.
+
+The implementation provides special handling for distributed package availability
+checks and proper tracking of distributed state and operations across processes.
+"""
+
+import functools
+import inspect
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+from torch.fx.experimental._backward_state import BackwardState
+
+from .. import compiled_autograd
+from .._trace_wrapped_higher_order_op import trace_wrapped
+from ..exc import unimplemented
+from ..external_utils import call_module_hooks_from_backward_state
+from ..guards import GuardBuilder, install_guard
+from ..source import AttrSource
+from .base import VariableTracker
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+
+class DistributedVariable(VariableTracker):
+    """
+    The base distributed variable that encapsulates common methods
+    for the distributed objects (i.e. ProcessGroup, DeviceMesh, etc.).
+    Concrete distributed objects could inherit this class and add object
+    specific logic.
+
+    i.e. It provides the check on the distributed package existence
+    and hold the tracking value for the corresponding distributed object.
+    """
+
+    def __init__(self, value: Any, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        if not DistributedVariable.is_available():
+            unimplemented(
+                gb_type="torch.distributed package is not available!",
+                context="",
+                explanation="The PyTorch package doesn't include torch.distributed when building from source.",
+                hints=[
+                    "Set USE_DISTRIBUTED=1 to enable it when building PyTorch from source."
+                ],
+            )
+        self.value = value
+
+    def python_type(self) -> type:
+        return type(self.value)
+
+    @staticmethod
+    def is_available() -> bool:
+        # check if the distributed package is available or not
+        return torch.distributed.is_available()
+
+    def is_python_hashable(self) -> Literal[True]:
+        return True
+
+    def get_python_hash(self) -> int:
+        return hash(self.value)
+
+    def is_python_equal(self, other: object) -> bool:
+        return (
+            isinstance(other, VariableTracker)
+            and self.as_python_constant() == other.as_python_constant()
+        )
+
+
+def is_from_local(value: object) -> bool:
+    if not DistributedVariable.is_available():
+        return False
+    from torch.distributed.tensor import DTensor
+
+    return inspect.isfunction(value) and value is DTensor.from_local
+
+
+def is_constant_pg_functions(value: object) -> bool:
+    if not DistributedVariable.is_available():
+        return False
+
+    from torch.distributed.distributed_c10d import (
+        _get_group_size_by_name,
+        _get_group_tag,
+        _rank_not_in_group,
+        _resolve_group_name_by_ranks_and_tag,
+        get_process_group_ranks,
+    )
+
+    constant_processgroup_functions = [
+        _get_group_size_by_name,
+        _get_group_tag,
+        _rank_not_in_group,
+        get_process_group_ranks,
+        _resolve_group_name_by_ranks_and_tag,
+    ]
+
+    return inspect.isfunction(value) and value in constant_processgroup_functions
+
+
+class WorldMetaClassVariable(DistributedVariable):
+    """
+    Tracks torch.distributed.GroupMember and torch.distributed.group, which are
+    instances of the metaclass _WorldMeta.
+    """
+
+    @classmethod
+    def is_group_member_type(cls, value: object) -> bool:
+        if not cls.is_available():
+            return False
+
+        from torch.distributed.distributed_c10d import _WorldMeta
+
+        return type(value) is _WorldMeta
+
+    def python_type(self) -> type:
+        return type(self.value)
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        if name == "WORLD":
+            from .builder import SourcelessBuilder
+
+            assert self.source
+            source = AttrSource(base=self.source, member="WORLD")
+            install_guard(source.make_guard(GuardBuilder.ID_MATCH))
+            return SourcelessBuilder.create(tx, self.value.WORLD)
+        elif name == "NON_GROUP_MEMBER":
+            assert self.source
+            source = AttrSource(base=self.source, member="NON_GROUP_MEMBER")
+            install_guard(source.make_guard(GuardBuilder.ID_MATCH))
+            return VariableTracker.build(tx, self.value.NON_GROUP_MEMBER)
+        return super().var_getattr(tx, name)
+
+
+class BackwardHookVariable(VariableTracker):
+    """
+    Handles torch.utils.hooks.BackwardHook for module-level backward
+    hooks.
+    """
+
+    @staticmethod
+    def create(
+        tx: "InstructionTranslator",
+        module: VariableTracker,
+        user_hooks: VariableTracker,
+        user_pre_hooks: VariableTracker,
+    ) -> "BackwardHookVariable":
+        if not compiled_autograd.compiled_autograd_enabled:
+            unimplemented(
+                gb_type="Module-level backwards hooks require compiled autograd.",
+                context="",
+                explanation="",
+                hints=[
+                    "Enable compiled autograd by setting torch._dynamo.config.compiled_autograd = True."
+                ],
+            )
+
+        def _in_graph_bw_hooks(
+            bw_state: BackwardState,
+        ) -> torch.utils.hooks.BackwardHook:
+            """
+            Rather than installing the user hooks in the graph (which
+            don't survive AotAutograd), we install hooks that will call
+            trace_wrapped in the backward pass that CompiledAutograd
+            can turn into actual hook calls.
+            """
+            return torch.utils.hooks.BackwardHook(
+                None,
+                (
+                    functools.partial(
+                        trace_wrapped,
+                        fn=call_module_hooks_from_backward_state,
+                        bw_state=bw_state,
+                        hooks_name=user_hooks_name,
+                        module_name=module_name,
+                    ),
+                ),
+                (
+                    functools.partial(
+                        trace_wrapped,
+                        fn=call_module_hooks_from_backward_state,
+                        bw_state=bw_state,
+                        hooks_name=user_pre_hooks_name,
+                        module_name=module_name,
+                    ),
+                ),
+            )
+
+        module_name, bw_state_proxy = tx.output.add_backward_state_hook(module, "mod")
+        user_pre_hooks_name, _ = tx.output.add_backward_state_hook(user_pre_hooks)
+        user_hooks_name, _ = tx.output.add_backward_state_hook(user_hooks)
+        proxy = tx.output.create_proxy(
+            "call_function",
+            _in_graph_bw_hooks,
+            (bw_state_proxy,),
+            {},
+        )
+        proxy.node.meta["example_value"] = torch.utils.hooks.BackwardHook(None, (), ())
+        return BackwardHookVariable(proxy, module, user_hooks, user_pre_hooks)
+
+    def __init__(
+        self,
+        proxy: torch.fx.Proxy,
+        module: VariableTracker,
+        user_hooks: VariableTracker,
+        user_pre_hooks: VariableTracker,
+        **options: Any,
+    ) -> None:
+        super().__init__(**options)
+        self.proxy = proxy
+        self.module = module
+        self.user_hooks = user_hooks
+        self.user_pre_hooks = user_pre_hooks
+
+    def as_proxy(self) -> torch.fx.Proxy:
+        return self.proxy
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name in ("setup_input_hook", "setup_output_hook"):
+            return self._setup_hook(tx, name, *args, **kwargs)
+        return super().call_method(tx, name, args, kwargs)
+
+    def _setup_hook(
+        self, tx: "InstructionTranslator", hook_method_name: str, args: VariableTracker
+    ) -> VariableTracker:
+        from .builder import wrap_fx_proxy
+
+        return wrap_fx_proxy(
+            tx,
+            tx.output.create_proxy(
+                "call_method",
+                hook_method_name,
+                (self.as_proxy(), args.as_proxy()),
+                {},
+            ),
+        )

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/hashable.py

@@ -0,0 +1,163 @@
+"""
+Hashability utilities for PyTorch Dynamo variable tracking.
+
+This module provides the HashableTracker wrapper class and associated utilities
+for making VariableTracker instances usable as dictionary keys and set elements
+during symbolic execution. Used by both ConstDictVariable and SetVariable.
+"""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from .. import variables
+from ..exc import raise_observed_exception
+from ..utils import specialize_symnode
+from .base import VariableTracker
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+
+def raise_unhashable(
+    arg: VariableTracker, tx: "InstructionTranslator | None" = None
+) -> None:
+    if tx is None:
+        from torch._dynamo.symbolic_convert import InstructionTranslator
+
+        tx = InstructionTranslator.current_tx()
+    try:
+        arg_type = arg.python_type()
+    except Exception:
+        arg_type = type(arg)
+
+    raise_observed_exception(
+        TypeError,
+        tx,
+        args=[
+            f"unhashable type: {arg_type!r} and variable tracker = {type(arg.realize())}",
+        ],
+    )
+
+
+def is_hashable(x: VariableTracker) -> bool:
+    # NB - performing isinstance check on a LazVT realizes the VT, accidentally
+    # inserting the guard. To avoid this, lazyVT `is_hashable` methods looks at
+    # the underlying value without realizing the VT. Consider updating the
+    # lazyVT `is_hashable` method if you see unnecessary guarding for a key VT.
+    if (
+        isinstance(x, variables.LazyVariableTracker)
+        and not x.is_realized()
+        and x.is_hashable()
+    ):
+        return True
+    return x.is_python_hashable()
+
+
+class HashableTracker:
+    """
+    Class that wraps a VariableTracker and makes it hashable.
+    Note that it's fine to put VTs into dictionaries and sets, but doing so
+    does not take into account aliasing.
+    """
+
+    _MISSING = object()
+
+    def __init__(self, vt: VariableTracker) -> None:
+        # We specialize SymNodes
+        vt = specialize_symnode(vt)
+
+        # If Dynamo does not know the hashability of the vt, it will raise unsupported here
+        # TODO(follow-up): check tp_hash via C-level slot detection — unhashable keys
+        # (e.g. list) should raise TypeError, not graph break via is_python_hashable/unimplemented.
+        if not is_hashable(vt):
+            raise_unhashable(vt)
+        self.vt = vt
+
+    @classmethod
+    def _maybe_constant_torch_size(cls, vt: VariableTracker) -> object:
+        from .lists import SizeVariable
+        from .tensor import TensorVariable
+
+        if (
+            isinstance(vt, variables.LazyVariableTracker)
+            and not vt.is_realized()
+            and isinstance(vt.original_value(), torch.Size)
+        ):
+            return vt.original_value()
+
+        if not isinstance(vt, SizeVariable):
+            return cls._MISSING
+
+        items = []
+        for item in vt.items:
+            if item.is_python_constant():
+                items.append(item.as_python_constant())
+                continue
+
+            if isinstance(item, TensorVariable):
+                proxy = getattr(item, "proxy", None)
+                node = getattr(proxy, "node", None)
+                meta = getattr(node, "meta", None) if node is not None else None
+                example_value = (
+                    meta.get("example_value") if isinstance(meta, dict) else None
+                )
+                constant = getattr(example_value, "constant", None)
+
+                if isinstance(constant, torch.Tensor) and constant.numel() == 1:
+                    items.append(constant.item())
+                    continue
+
+            return cls._MISSING
+
+        return torch.Size(items)
+
+    def __hash__(self) -> int:
+        """
+        Computes the hash value for the wrapped VariableTracker.
+
+        For unrealized LazyVariableTrackers, uses the hash of the original value
+        to avoid realizing the tracker and inserting unnecessary guards.
+        For all other cases, delegates to the VariableTracker's get_python_hash method.
+
+        Returns:
+            The hash value of the underlying variable tracker
+        """
+        if (
+            isinstance(self.vt, variables.LazyVariableTracker)
+            and not self.vt.is_realized()
+            and self.vt.is_hashable()
+        ):
+            return hash(self.vt.original_value())
+
+        maybe_constant = self._maybe_constant_torch_size(self.vt)
+        if maybe_constant is not self._MISSING:
+            return hash(maybe_constant)
+
+        return self.vt.get_python_hash()
+
+    def __eq__(self, other: object) -> bool:
+        """
+        Checks equality between two HashableTracker instances.
+
+        Delegates to the VariableTracker's is_python_equal method to compare
+        the underlying variable trackers for Python-level equality.
+
+        Args:
+            other: Another HashableTracker instance to compare with
+
+        Returns:
+            True if the underlying variable trackers are Python-equal, False otherwise
+        """
+        if not isinstance(other, HashableTracker):
+            return False
+        if self.vt is other.vt:
+            return True
+
+        self_constant = self._maybe_constant_torch_size(self.vt)
+        other_constant = self._maybe_constant_torch_size(other.vt)
+        if self_constant is not self._MISSING and other_constant is not self._MISSING:
+            return self_constant == other_constant
+
+        return self.vt.is_python_equal(other.vt)

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/iter.py

@@ -0,0 +1,657 @@
+"""
+This module provides iterator-related variable tracking functionality for Dynamo.
+It implements variable classes for handling Python iterators and itertools functions
+during symbolic execution and tracing.
+
+The module includes:
+- Base iterator variable classes for tracking iterator state
+- Implementations of built-in iterators (zip, map, filter)
+- Support for itertools functions (product, accumulate, combinations, etc.)
+- Mutation tracking and reconstruction capabilities for iterator operations
+
+These classes integrate with Dynamo's variable tracking system to enable proper
+handling of iterator operations during code transformation and optimization.
+"""
+
+import itertools
+import sys
+from collections.abc import Callable, Sequence
+from typing import Any, TYPE_CHECKING
+
+from .. import graph_break_hints, polyfills, variables
+from ..bytecode_transformation import (
+    create_build_tuple,
+    create_call_function,
+    create_call_function_ex,
+    create_instruction,
+)
+from ..exc import (
+    handle_observed_exception,
+    ObservedUserStopIteration,
+    raise_observed_exception,
+    unimplemented,
+    UserError,
+)
+from .base import ValueMutationNew, VariableTracker
+from .constant import ConstantVariable
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.codegen import PyCodegen
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+
+MAX_ITERATOR_LIMIT = 100 * 1024  # 100k
+
+
+class ItertoolsVariable(VariableTracker):
+    def __init__(self, value: Any, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.value = value
+
+    def __repr__(self) -> str:
+        return f"ItertoolsVariable({self.value})"
+
+    def as_python_constant(self) -> Any:
+        return self.value
+
+    def get_real_python_backed_value(self) -> Any:
+        return self.value
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence["VariableTracker"],
+        kwargs: "dict[str, VariableTracker]",
+    ) -> "VariableTracker":
+        # See also: module `torch._dynamo.polyfills.itertools`
+
+        if self.value is itertools.product:
+            if any(kw != "repeat" for kw in kwargs):
+                unimplemented(
+                    gb_type="Unsupported kwargs for itertools.product",
+                    context=f"call_function {self} {args} {kwargs}",
+                    explanation=f"Expected kwargs: 'repeat', but got "
+                    f"{','.join(set(kwargs.keys()) - {'repeat'})}",
+                    hints=[*graph_break_hints.USER_ERROR],
+                )
+
+            if "repeat" in kwargs:
+                r = kwargs["repeat"].as_python_constant()
+            else:
+                r = 1
+            seqs = [arg.force_unpack_var_sequence(tx) for arg in args]
+            items = [
+                variables.TupleVariable(list(item))
+                for item in itertools.product(*seqs, repeat=r)
+            ]
+            return variables.ListIteratorVariable(
+                items,  # type: ignore[arg-type]
+                mutation_type=ValueMutationNew(),
+            )
+        elif (
+            self.value is itertools.combinations
+            and not kwargs
+            and len(args) == 2
+            and args[0].has_unpack_var_sequence(tx)
+            and args[1].is_python_constant()
+        ):
+            iterable = args[0].unpack_var_sequence(tx)
+            r = args[1].as_python_constant()
+
+            items = []
+            for item in itertools.combinations(iterable, r):
+                items.append(variables.TupleVariable(list(item)))
+            return variables.ListIteratorVariable(
+                items,  # type: ignore[arg-type]
+                mutation_type=ValueMutationNew(),
+            )
+        elif self.value is itertools.groupby:
+            if any(kw != "key" for kw in kwargs):
+                unimplemented(
+                    gb_type="Unsupported kwargs for itertools.groupby",
+                    context=f"call_function {self} {args} {kwargs}",
+                    explanation=f"Expected kwargs: 'key', but got "
+                    f"{','.join(set(kwargs.keys()) - {'key'})}",
+                    hints=[*graph_break_hints.USER_ERROR],
+                )
+
+            def retrieve_const_key(key: VariableTracker) -> Any:
+                if isinstance(key, variables.SymNodeVariable):
+                    return key.evaluate_expr()
+                elif key.is_python_constant():
+                    return key.as_python_constant()
+                else:
+                    unimplemented(
+                        gb_type="Unsupported key type for itertools.groupby",
+                        context=f"call_function {self} {args} {kwargs}",
+                        explanation="Dynamo does not know how to trace "
+                        f"itertools.groupby with key type: {str(type(key))}. "
+                        "We only support grouping keys that are constants (int, float, str, etc.)",
+                        hints=[*graph_break_hints.SUPPORTABLE],
+                    )
+
+            if len(args) == 1 and args[0].has_unpack_var_sequence(tx):
+                seq = args[0].unpack_var_sequence(tx)
+            else:
+                unimplemented(
+                    gb_type="Unsupported arguments for itertools.groupby",
+                    context=f"call_function {self} {args} {kwargs}",
+                    explanation="Dynamo does not know how to trace "
+                    f"itertools.groupby with args: {args} and kwargs: {kwargs}. "
+                    "itertools.groupby expects an iterable to group and an "
+                    "optional key function to determine groupings.",
+                    hints=[
+                        "Make sure the arguments to itertools.groupby are correct.",
+                        *graph_break_hints.SUPPORTABLE,
+                    ],
+                )
+
+            if "key" in kwargs:
+
+                def keyfunc(x: VariableTracker) -> Any:
+                    return retrieve_const_key(
+                        kwargs.get("key").call_function(tx, [x], {})  # type: ignore[union-attr]
+                    )
+
+            else:
+
+                def keyfunc(x: VariableTracker) -> Any:
+                    return retrieve_const_key(x)
+
+            result = []
+            try:
+                for k, v in itertools.groupby(seq, key=keyfunc):
+                    result.append(
+                        variables.TupleVariable(
+                            [
+                                (
+                                    variables.ConstantVariable.create(k)
+                                    if variables.ConstantVariable.is_literal(k)
+                                    else k
+                                ),
+                                variables.ListIteratorVariable(
+                                    list(v), mutation_type=ValueMutationNew()
+                                ),
+                            ],
+                            mutation_type=ValueMutationNew(),
+                        )
+                    )
+            except Exception as e:
+                unimplemented(
+                    gb_type="Unexpected failure during itertools.groupby() iteration",
+                    context=f"call_function {self} {args} {kwargs}",
+                    explanation="Unexpected failure in invoking function during groupby",
+                    hints=[*graph_break_hints.SUPPORTABLE],
+                    from_exc=e,
+                )
+            return variables.ListIteratorVariable(
+                result,  # type: ignore[arg-type]
+                mutation_type=ValueMutationNew(),
+            )
+        elif self.value is itertools.repeat:
+            if len(args) < 2:
+                return variables.RepeatIteratorVariable(
+                    *args, mutation_type=ValueMutationNew()
+                )
+
+            return tx.inline_user_function_return(
+                VariableTracker.build(tx, polyfills.repeat), args, kwargs
+            )
+        elif self.value is itertools.count and not kwargs:
+            if len(args) == 0:
+                return variables.CountIteratorVariable(mutation_type=ValueMutationNew())
+            if len(args) == 1:
+                return variables.CountIteratorVariable(
+                    item=args[0], mutation_type=ValueMutationNew()
+                )
+            if len(args) == 2:
+                return variables.CountIteratorVariable(
+                    item=args[0],
+                    step=args[1],
+                    mutation_type=ValueMutationNew(),
+                )
+            return super().call_function(tx, args, kwargs)
+        elif (
+            self.value is itertools.permutations
+            and (len(args) == 1 or (len(args) == 2 and args[1].is_python_constant()))
+            and not kwargs
+        ):
+            if len(args) == 2:
+                r = args[1].as_python_constant()
+            else:
+                r = None
+            items = [
+                variables.TupleVariable(list(item))
+                for item in itertools.permutations(
+                    args[0].force_unpack_var_sequence(tx), r
+                )
+            ]
+            return variables.ListIteratorVariable(
+                items,  # type: ignore[arg-type]
+                mutation_type=ValueMutationNew(),
+            )
+        else:
+            return super().call_function(tx, args, kwargs)
+
+
+class IteratorVariable(VariableTracker):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        unimplemented(
+            gb_type="Unimplemented next() call",
+            context=f"next({self})",
+            explanation="This abstract method must be implemented",
+            hints=[*graph_break_hints.DYNAMO_BUG],
+        )
+
+    # NOTE: only call when unpacking this iterator safely done eagerly!
+    # Normally, iterators are accessed lazily.
+    # Example of safe eager unpacking: list(map(f, seq))
+    # Example of unsafe eager unpacking: list(islice(map(f, seq), 5))
+    def force_unpack_var_sequence(
+        self, tx: "InstructionTranslator"
+    ) -> list[VariableTracker]:
+        result: list[VariableTracker] = []
+        self.force_apply_to_var_sequence(tx, result.append)
+        return result
+
+    def force_apply_to_var_sequence(
+        self, tx: "InstructionTranslator", fn: Callable[[Any], Any]
+    ) -> None:
+        while True:
+            try:
+                fn(self.next_variable(tx))
+            except ObservedUserStopIteration:
+                handle_observed_exception(tx)
+                break
+
+    # don't call force_unpack_var_sequence since it can mutate
+    # IteratorVariable state!
+    def has_force_unpack_var_sequence(self, tx: "InstructionTranslator") -> bool:
+        return True
+
+    def call_obj_hasattr(
+        self, tx: "InstructionTranslator", name: str
+    ) -> "ConstantVariable":
+        if name == "__iter__" or name == "__next__":
+            return variables.ConstantVariable.create(True)
+        return super().call_obj_hasattr(tx, name)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name == "__iter__":
+            return self
+        elif name == "__next__":
+            return self.next_variable(tx)
+        return super().call_method(tx, name, args, kwargs)
+
+
+class ObjectIteratorVariable(IteratorVariable):
+    """
+    VariableTracker for iter(obj) that implements the iterator protocol (i.e.,
+    has a `__next__` method).
+
+    We use this class to track the state of the iterator and handle the case
+    when the iterator is exhausted:
+
+    Example usage:
+        > b = iter(obj)
+        > list(b)  # exhaust the iterator
+        > list(b)  # empty list
+    """
+
+    def __init__(self, obj: VariableTracker, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.obj = obj
+        self.generator_exhausted = False
+
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        if self.generator_exhausted:
+            raise_observed_exception(StopIteration, tx)
+
+        try:
+            return self.obj.next_variable(tx)
+        except ObservedUserStopIteration:
+            # Do not rely on the object to always return StopIteration once it
+            # is exhausted.
+            self.generator_exhausted = True
+            raise
+
+
+class RepeatIteratorVariable(IteratorVariable):
+    def __init__(self, item: VariableTracker, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self.item = item
+
+    def python_type(self) -> type:
+        return itertools.repeat
+
+    # Repeat needs no mutation, clone self
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        return self.item
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.extend_output(
+                [
+                    codegen.create_load_python_module(itertools),
+                    codegen.create_load_attr("repeat"),
+                ]
+            )
+        )
+        codegen(self.item)
+        codegen.extend_output(create_call_function(1, False))
+
+
+class CountIteratorVariable(IteratorVariable):
+    # advance_count tracks how many next() calls were made during tracing,
+    # used by side_effects.py to replay them on the real iterator post-execution.
+    _nonvar_fields = {
+        "advance_count",
+        *IteratorVariable._nonvar_fields,
+    }
+
+    def python_type(self) -> type:
+        return itertools.count
+
+    def __init__(
+        self,
+        item: int | VariableTracker = 0,
+        step: int | VariableTracker = 1,
+        advance_count: int = 0,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        if not isinstance(item, VariableTracker):
+            item = ConstantVariable.create(item)
+        if not isinstance(step, VariableTracker):
+            step = ConstantVariable.create(step)
+        self.item = item
+        self.step = step
+        self.advance_count = advance_count
+
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        assert self.is_mutable()
+        old_item = self.item
+        tx.output.side_effects.mutation(self)
+        self.item = self.item.call_method(tx, "__add__", [self.step], {})
+        self.advance_count += 1
+        return old_item
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.extend_output(
+                [
+                    codegen.create_load_python_module(itertools),
+                    codegen.create_load_attr("count"),
+                ]
+            )
+        )
+        codegen(self.item)
+        codegen(self.step)
+        codegen.extend_output(create_call_function(2, False))
+
+
+class ZipVariable(IteratorVariable):
+    """
+    Represents zip(*iterables)
+    """
+
+    # PyZip_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L3011
+    _cpython_type = zip
+
+    _nonvar_fields = {
+        "index",
+        "strict",
+        *IteratorVariable._nonvar_fields,
+    }
+
+    def __init__(
+        self,
+        iterables: list[VariableTracker],
+        strict: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        assert isinstance(iterables, list)
+        # can be list[Variable] or VariableTracker (with next_variable implemented)
+        self.iterables = iterables
+        self.index = 0
+        self.strict = strict
+
+    def python_type(self) -> type[zip]:  # type: ignore[type-arg]
+        return zip
+
+    def has_unpack_var_sequence(self, tx: "InstructionTranslator") -> bool:
+        return all(
+            isinstance(it, list) or it.has_unpack_var_sequence(tx)
+            for it in self.iterables
+        )
+
+    def unpack_var_sequence(
+        self, tx: "InstructionTranslator"
+    ) -> list["VariableTracker"]:
+        assert self.has_unpack_var_sequence(tx)
+        iterables = []
+        for it in self.iterables:
+            if isinstance(it, list):
+                iterables.append(it[self.index :])
+            else:
+                iterables.append(it.unpack_var_sequence(tx))
+        kwargs = {"strict": self.strict} if self.strict else {}
+        zipped = zip(*iterables, **kwargs)
+        return [variables.TupleVariable(list(var)) for var in zipped]
+
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        assert self.is_mutable()
+
+        if len(self.iterables) == 0:
+            raise_observed_exception(StopIteration, tx)
+
+        old_index = self.index
+        args = []
+
+        def get_item(
+            it: list[VariableTracker] | VariableTracker,
+        ) -> VariableTracker:
+            if isinstance(it, list):
+                if old_index >= len(it):
+                    raise_observed_exception(StopIteration, tx)
+                return it[old_index]
+            else:
+                return it.next_variable(tx)
+
+        idx: int | None = None
+        try:
+            for idx, it in enumerate(self.iterables):  # noqa:B007
+                args.append(get_item(it))
+        except ObservedUserStopIteration:
+            if self.strict:
+                if idx == 0:
+                    # all other iterables should be exhausted
+                    for it in self.iterables:
+                        try:
+                            get_item(it)
+                        except ObservedUserStopIteration:
+                            handle_observed_exception(tx)
+                            continue
+                        # no ObservedUserStopIteration - fall through to UserError
+                        break
+                    else:
+                        # all iterables exhausted, raise original error
+                        raise
+                handle_observed_exception(tx)
+                raise UserError(
+                    ValueError,  # type: ignore[arg-type]
+                    "zip() has one argument of len differing from others",
+                ) from None
+            raise
+
+        tx.output.side_effects.mutation(self)
+        self.index += 1
+        return variables.TupleVariable(args)
+
+    def reconstruct_items(self, codegen: "PyCodegen") -> None:
+        for it in self.iterables:
+            if isinstance(it, list):
+                remaining_items = it[self.index :]
+                codegen.foreach(remaining_items)
+                codegen.append_output(create_build_tuple(len(remaining_items)))
+            else:
+                codegen(it)
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.load_import_from("builtins", "zip"), call_function_ex=True
+        )
+        self.reconstruct_items(codegen)
+        codegen.append_output(create_build_tuple(len(self.iterables)))
+        codegen.extend_output(
+            [
+                codegen.create_load_const("strict"),
+                codegen.create_load_const(self.strict),
+                create_instruction("BUILD_MAP", arg=1),
+                *create_call_function_ex(True, False),
+            ]
+        )
+
+
+class MapVariable(ZipVariable):
+    """
+    Represents map(fn, *iterables)
+    """
+
+    # PyMap_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L1484
+    _cpython_type = map
+
+    def __init__(
+        self,
+        fn: VariableTracker,
+        iterables: list[VariableTracker],
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(iterables, **kwargs)
+        self.fn = fn
+
+    def python_type(self) -> type:
+        return map
+
+    def has_unpack_var_sequence(self, tx: "InstructionTranslator") -> bool:
+        return False
+
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        args = super().next_variable(tx)
+        return self.fn.call_function(tx, args.items, {})  # type: ignore[attr-defined]
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.load_import_from("builtins", "map"), call_function_ex=True
+        )
+        codegen(self.fn)
+        self.reconstruct_items(codegen)
+        codegen.append_output(create_build_tuple(len(self.iterables) + 1))
+        if self.strict:
+            assert sys.version_info >= (3, 14), (
+                "Unexpected bug: map(strict=True) requires Python 3.14+"
+            )
+            codegen.extend_output(
+                [
+                    codegen.create_load_const("strict"),
+                    codegen.create_load_const(self.strict),
+                    create_instruction("BUILD_MAP", arg=1),
+                    *create_call_function_ex(True, False),
+                ]
+            )
+        else:
+            codegen.extend_output(create_call_function_ex(False, False))
+
+
+class FilterVariable(IteratorVariable):
+    """
+    Represents filter(fn, iterable)
+    """
+
+    # PyFilter_Type: https://github.com/python/cpython/blob/v3.13.0/Python/bltinmodule.c#L630
+    _cpython_type = filter
+
+    _nonvar_fields = {
+        "index",
+        *IteratorVariable._nonvar_fields,
+    }
+
+    def __init__(
+        self,
+        fn: VariableTracker,
+        iterable: list[VariableTracker],
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.fn = fn
+        self.iterable = iterable
+        self.index = 0
+
+    def python_type(self) -> type:
+        return filter
+
+    def has_unpack_var_sequence(self, tx: "InstructionTranslator") -> bool:
+        return isinstance(self.iterable, list) or self.iterable.has_unpack_var_sequence(
+            tx
+        )
+
+    def unpack_var_sequence(
+        self, tx: "InstructionTranslator"
+    ) -> list["VariableTracker"]:
+        assert self.has_unpack_var_sequence(tx)
+        it = None
+        if isinstance(self.iterable, list):
+            it = self.iterable[self.index :]
+        else:
+            it = self.iterable.unpack_var_sequence(tx)
+        filtered = self.fn.call_function(tx, it, {})
+        return [variables.TupleVariable([filtered])]
+
+    def next_variable(self, tx: "InstructionTranslator") -> VariableTracker:
+        def _next() -> VariableTracker:
+            old_index = self.index
+            if isinstance(self.iterable, list):
+                if old_index >= len(self.iterable):
+                    raise_observed_exception(StopIteration, tx)
+                return self.iterable[old_index]
+            else:
+                return self.iterable.next_variable(tx)
+
+        # A do-while loop to find elements that make fn return true
+        while True:
+            item = _next()
+            self.index += 1
+            if self.fn.is_constant_none():
+                res = item
+            else:
+                res = self.fn.call_function(tx, [item], {})
+            pred_res = variables.UserFunctionVariable(
+                polyfills.predicate  # type: ignore[arg-type]
+            ).call_function(tx, [res], {})
+            if pred_res.as_python_constant():
+                return item
+
+    def reconstruct_items(self, codegen: "PyCodegen") -> None:
+        if isinstance(self.iterable, list):
+            remaining_items = self.iterable[self.index :]
+            codegen.foreach(remaining_items)
+            codegen.append_output(create_build_tuple(len(remaining_items)))
+        else:
+            codegen(self.iterable)
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(lambda: codegen.load_import_from("builtins", "filter"))
+        codegen(self.fn)
+        self.reconstruct_items(codegen)
+        codegen.extend_output(create_call_function(2, False))

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/lazy.py

@@ -0,0 +1,306 @@
+from __future__ import annotations
+
+import collections
+import functools
+import inspect
+from typing import Any, TYPE_CHECKING
+
+from ..utils import is_function_or_wrapper
+from .base import VariableTracker, VariableTrackerMeta
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing_extensions import Self
+
+    from .tensor import SymNodeVariable
+
+
+class LazyCache:
+    """Container to cache the real VariableTracker"""
+
+    def __init__(self, value: Any, source: Any) -> None:
+        if not isinstance(value, LazySymNodeFormatString):
+            assert source
+        self.value = value
+        self.source = source
+        self.name_hint: str | None = None
+        self.vt: VariableTracker | None = None
+
+    def realize(self) -> None:
+        assert self.vt is None
+        from ..symbolic_convert import InstructionTranslator
+        from . import builder
+
+        tx = InstructionTranslator.current_tx()
+
+        if isinstance(self.value, LazySymNodeFormatString):
+            self.vt = builder.SourcelessBuilder.create(tx, self.value)
+        else:
+            # Pass allow_lazy_constant=False to prevent VariableBuilder from
+            # returning LazyConstantVariable, which would cause infinite recursion
+            # when LazyVariableTracker.realize() returns LazyConstantVariable.
+            self.vt = builder.VariableBuilder(
+                tx, self.source, allow_lazy_constant=False
+            )(self.value)
+
+        if self.name_hint is not None:
+            self.vt.set_name_hint(self.name_hint)
+
+        del self.value
+        del self.source
+        del self.name_hint
+
+
+class LazyVariableTracker(VariableTracker, metaclass=VariableTrackerMeta):
+    """
+    A structure that defers the creation of the actual VariableTracker
+    for a given underlying value until it is accessed.
+
+    The `realize` function invokes VariableTracker.build() to produce the real object.
+    Once a LazyVariableTracker has been realized, internal bookkeeping will
+    prevent double realization.
+
+    This object should be utilized for processing containers, or objects that
+    reference other objects where we may not want to take on creating all the
+    VariableTrackers right away.
+    """
+
+    # Flag to prevent implicit realization in isinstance checks (inherited by subclasses)
+    _no_implicit_realize = True
+    _nonvar_fields = {"_cache", *VariableTracker._nonvar_fields}
+
+    @staticmethod
+    def create(value: Any, source: Any, **options: Any) -> VariableTracker:
+        if type(value) in LazyConstantVariable.supported_types:
+            return LazyConstantVariable.create(value, source, **options)
+
+        # Cache based on source when no extra options are passed
+        if source is not None and not options:
+            from ..symbolic_convert import InstructionTranslator
+
+            tx = InstructionTranslator.current_tx()
+            if tx is not None:
+                cache = tx.output.variable_tracker_cache
+                cached = cache.get(source)
+                if cached is not None:
+                    return cached
+                vt = LazyVariableTracker(LazyCache(value, source), source=source)
+                cache[source] = vt
+                return vt
+
+        return LazyVariableTracker(LazyCache(value, source), source=source, **options)
+
+    def __init__(self, _cache: LazyCache, **kwargs: Any) -> None:
+        assert isinstance(_cache, LazyCache)
+        super().__init__(**kwargs)
+        self._cache = _cache
+
+    def realize(self) -> VariableTracker:
+        """Force construction of the real VariableTracker"""
+        if self._cache.vt is None:
+            self._cache.realize()
+            assert self._cache.vt is not None
+        return self._cache.vt
+
+    def lazy_isinstance(self, cls: type) -> bool:
+        """Check isinstance after realizing, used by ImplicitRealizingVariableTrackerMeta"""
+        return type.__instancecheck__(cls, self.realize())
+
+    def unwrap(self) -> VariableTracker | Self:
+        """Return the real VariableTracker if it already exists"""
+        if self.is_realized():
+            assert self._cache.vt is not None
+            return self._cache.vt
+        return self
+
+    def is_realized(self) -> bool:
+        return self._cache.vt is not None
+
+    def clone(self, **kwargs: Any) -> VariableTracker:
+        assert kwargs.get("_cache", self._cache) is self._cache
+        if kwargs.get("source", self.source) is not self.source:
+            self.realize()
+        return VariableTracker.clone(self.unwrap(), **kwargs)
+
+    def peek_type(self) -> type[Any]:
+        assert not self.is_realized()
+        return type(self._cache.value)
+
+    def peek_value(self) -> Any:
+        assert not self.is_realized()
+        return self._cache.value
+
+    def set_name_hint(self, name: str) -> None:
+        if self.is_realized():
+            self._cache.vt.set_name_hint(name)  # type: ignore[union-attr]
+        else:
+            self._cache.name_hint = name
+
+    def __str__(self) -> str:
+        variable_info = "LazyVariableTracker("
+        if self.is_realized():
+            variable_info += f"realized: {repr(self.unwrap())})"
+        else:
+            variable_info += f"unrealized: {self.peek_type()})"
+
+        return variable_info
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self.realize(), item)
+
+    # most methods are auto-generated below, these are the ones we want to exclude
+    visit = VariableTracker.visit  # type: ignore[assignment]
+    __repr__ = __str__
+
+    @classmethod
+    def realize_all(
+        cls,
+        value: Any,
+        cache: dict[int, tuple[Any, Any]] | None = None,
+        *,
+        allow_lazy_constant: bool = False,
+    ) -> Any:
+        """
+        Walk an object and realize all LazyVariableTrackers inside it.
+        """
+        if cache is None:
+            cache = {}
+
+        idx = id(value)
+        if idx in cache:
+            return cache[idx][0]
+
+        value_cls = type(value)
+        if issubclass(value_cls, LazyVariableTracker):
+            # Allow LazyConstantVariable to stay lazy when returning from a frame
+            keep_lazy = allow_lazy_constant and isinstance(value, LazyConstantVariable)
+            if keep_lazy:
+                result = value
+            else:
+                result = cls.realize_all(
+                    value.realize(), cache, allow_lazy_constant=allow_lazy_constant
+                )
+        elif issubclass(value_cls, VariableTracker):
+            # update value in-place
+            result = value
+            # update cache now to prevent infinite recursion
+            cache[idx] = (result, value)
+            value_dict = value.__dict__
+            nonvars = value._nonvar_fields
+            for key in value_dict:
+                if key not in nonvars:
+                    value_dict[key] = cls.realize_all(
+                        value_dict[key], cache, allow_lazy_constant=allow_lazy_constant
+                    )
+        elif value_cls is list:
+            result = [
+                cls.realize_all(v, cache, allow_lazy_constant=allow_lazy_constant)
+                for v in value
+            ]
+        elif value_cls is tuple:
+            result = tuple(
+                cls.realize_all(v, cache, allow_lazy_constant=allow_lazy_constant)
+                for v in value
+            )
+        elif value_cls in (dict, collections.OrderedDict):
+            result = {
+                k: cls.realize_all(v, cache, allow_lazy_constant=allow_lazy_constant)
+                for k, v in list(value.items())
+            }
+        else:
+            result = value
+
+        # save `value` to keep it alive and ensure id() isn't reused
+        cache[idx] = (result, value)
+        return result
+
+    def is_hashable(self) -> bool:
+        # Checks that the underlying value is hashable without realizing the VT.
+        # This is used by ConstDictVariable tracker to find if the key LazyVT
+        # can be hashed.
+        def _helper(value: Any) -> bool:
+            # TODO: Add support for more types
+            return (
+                inspect.isbuiltin(value)
+                or issubclass(type(value), type)
+                or is_function_or_wrapper(value)
+            )
+
+        assert not self.is_realized()
+        value = self._cache.value
+        if isinstance(value, tuple):
+            return all(_helper(v) for v in value)
+        return _helper(value)
+
+    def original_value(self) -> Any:
+        # Returns the value without realizing the VT.
+        assert not self.is_realized()
+        return self._cache.value
+
+    def original_source(self) -> Any:
+        # Returns the source without realizing the VT.
+        assert not self.is_realized()
+        return self._cache.source
+
+
+class LazyConstantVariable(LazyVariableTracker):
+    """
+    A lazy variable tracker for constants (int, float, bool, str) that defers
+    guarding until the value is actually used in a way that requires it.
+
+    This allows constants that are just passed through (e.g., returned without
+    being used in control flow or math) to avoid unnecessary recompilation when
+    their values change.
+    """
+
+    supported_types = (int, float, bool, str)
+
+    @staticmethod
+    def create(
+        value: Any,
+        source: Any,
+        **options: Any,
+    ) -> LazyConstantVariable:
+        assert type(value) in LazyConstantVariable.supported_types
+        return LazyConstantVariable(LazyCache(value, source), source=source, **options)
+
+
+class LazySymNodeFormatString:
+    def __init__(
+        self, sym_node_variable: SymNodeVariable, fmt_spec_var: VariableTracker
+    ) -> None:
+        from .constant import ConstantVariable
+
+        self.sym_node_var = sym_node_variable
+        self.fmt_var = ConstantVariable.create(
+            "{:" + fmt_spec_var.as_python_constant() + "}"
+        )
+
+    def __repr__(self) -> str:
+        return str.format(
+            self.fmt_var.as_python_constant(),
+            str(self.sym_node_var.evaluate_expr()),
+        )
+
+
+def _create_realize_and_forward(
+    name: str,
+) -> Callable[[LazyVariableTracker, Any, Any], Any]:
+    @functools.wraps(getattr(VariableTracker, name))
+    def realize_and_forward(
+        self: LazyVariableTracker, *args: Any, **kwargs: Any
+    ) -> Any:
+        return getattr(self.realize(), name)(*args, **kwargs)
+
+    return realize_and_forward
+
+
+def _populate() -> None:
+    for name, value in VariableTracker.__dict__.items():
+        if name not in LazyVariableTracker.__dict__:
+            if callable(value):
+                setattr(LazyVariableTracker, name, _create_realize_and_forward(name))
+
+
+_populate()

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/object_protocol.py

@@ -0,0 +1,319 @@
+"""
+Dynamo implementations of CPython's PyObject_* default slot algorithms.
+
+Analogous to CPython's Objects/object.c, this module holds the general
+dispatch machinery that is independent of any specific type.
+Per-type hook implementations (bool_impl, richcompare_impl, etc.)
+live in their respective VT files.
+"""
+
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+from torch._C._dynamo import (
+    get_type_slots,
+    has_slot,
+    PyMappingSlots,
+    PyNumberSlots,
+    PySequenceSlots,
+)
+
+from .. import graph_break_hints
+from ..exc import (
+    handle_observed_exception,
+    ObservedTypeError,
+    raise_observed_exception,
+    raise_type_error,
+    unimplemented,
+)
+from ..utils import istype
+from .base import NO_SUCH_SUBOBJ, VariableTracker
+from .constant import ConstantVariable
+
+
+if TYPE_CHECKING:
+    from ..symbolic_convert import InstructionTranslator
+
+
+def vt_identity_compare(
+    left: VariableTracker,
+    right: VariableTracker,
+) -> "VariableTracker | None":
+    """Try to determine Python identity (left is right) at trace time.
+
+    Returns ConstantVariable(True/False) if determinable, else None.
+    Mirrors the logic in BuiltinVariable's handle_is handler.
+    """
+    if left is right:
+        return ConstantVariable.create(True)
+
+    left_val = left.get_real_python_backed_value()
+    right_val = right.get_real_python_backed_value()
+    left_known = left_val is not NO_SUCH_SUBOBJ
+    right_known = right_val is not NO_SUCH_SUBOBJ
+
+    if left_known and right_known:
+        return (
+            ConstantVariable.create(True)
+            if left_val is right_val
+            else ConstantVariable.create(False)
+        )
+
+    # One side has a concrete backing object, the other doesn't — they can't
+    # be the same object.
+    if left_known != right_known:
+        return ConstantVariable.create(False)
+
+    # Mutable containers created during tracing: VT identity = Python identity.
+    from .dicts import ConstDictVariable
+    from .lists import ListVariable
+    from .sets import SetVariable
+
+    if isinstance(left, (ConstDictVariable, ListVariable, SetVariable)):
+        return ConstantVariable.create(False)
+
+    # Different Python types can never be the same object.
+    try:
+        if left.python_type() is not right.python_type():
+            return ConstantVariable.create(False)
+    except NotImplementedError:
+        pass
+
+    # Different exception types are never identical.
+    from .. import variables
+
+    if (
+        istype(left, variables.ExceptionVariable)
+        and istype(right, variables.ExceptionVariable)
+        and left.exc_type is not right.exc_type  # type: ignore[attr-defined]
+    ):
+        return ConstantVariable.create(False)
+
+    return None
+
+
+@lru_cache(maxsize=256)
+def _get_cached_slots(obj_type: type) -> tuple[int, int, int, int]:
+    """Get all type slots for a type (cached)."""
+    return get_type_slots(obj_type)
+
+
+def type_implements_sq_length(obj_type: type) -> bool:
+    """Check whether obj_type implements __len__ as sequence protocol"""
+    seq_slots, _, _, _ = _get_cached_slots(obj_type)
+    return has_slot(seq_slots, PySequenceSlots.SQ_LENGTH)
+
+
+def type_implements_mp_length(obj_type: type) -> bool:
+    """Check whether obj_type implements __len__ as mapping protocol"""
+    _, map_slots, _, _ = _get_cached_slots(obj_type)
+    return has_slot(map_slots, PyMappingSlots.MP_LENGTH)
+
+
+def type_implements_nb_bool(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_bool slot (i.e. has __bool__ or __len__)."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_BOOL)
+
+
+def type_implements_nb_int(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_int slot."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_INT)
+
+
+def type_implements_nb_index(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_index slot."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_INDEX)
+
+
+def type_implements_nb_float(obj_type: type) -> bool:
+    """Check whether obj_type implements the nb_float slot."""
+    _, _, number_slots, _ = _get_cached_slots(obj_type)
+    return has_slot(number_slots, PyNumberSlots.NB_FLOAT)
+
+
+def maybe_get_python_type(obj: VariableTracker) -> type:
+    try:
+        return obj.python_type()
+    except NotImplementedError:
+        unimplemented(
+            gb_type="Unsupported python_type() call",
+            context=f"{obj} does not implement python_type()",
+            explanation="This VariableTracker does not implement python_type(), "
+            "which is required for object protocol operations.",
+            hints=[
+                *graph_break_hints.DYNAMO_BUG,
+            ],
+        )
+
+
+def vt_mapping_size(
+    tx: "InstructionTranslator", obj: "VariableTracker"
+) -> "VariableTracker":
+    # ref: https://github.com/python/cpython/blob/v3.13.3/Objects/abstract.c#L2308-L2330
+    T = maybe_get_python_type(obj)
+    if type_implements_mp_length(T):
+        return obj.mp_length(tx)
+
+    if type_implements_sq_length(T):
+        raise_type_error(tx, f"{obj.python_type_name()} is not a mapping")
+
+    raise_type_error(tx, f"object of type {obj.python_type_name()} has no len()")
+
+
+def generic_len(
+    tx: "InstructionTranslator", obj: "VariableTracker"
+) -> "VariableTracker":
+    # ref: https://github.com/python/cpython/blob/v3.13.3/Objects/abstract.c#L53-L69
+    """
+    Implements PyObject_Size/PyObject_Length semantics for VariableTracker objects.
+    Dispatches to sq_length (sequences) or mp_length (mappings) depending on the VT type.
+    """
+
+    T = maybe_get_python_type(obj)
+    if type_implements_sq_length(T):
+        return obj.sq_length(tx)
+    return vt_mapping_size(tx, obj)
+
+
+def generic_bool(tx: "InstructionTranslator", obj: VariableTracker) -> VariableTracker:
+    """Mirrors PyObject_IsTrue.
+
+    https://github.com/python/cpython/blob/c09ccd9c429/Objects/object.c#L2135-L2158
+
+    Resolution order: constants → nb_bool → mp_length/sq_length → truthy.
+    """
+    from .constant import ConstantVariable
+
+    if obj.is_python_constant():
+        return ConstantVariable.create(bool(obj.as_python_constant()))
+
+    obj_type = maybe_get_python_type(obj)
+
+    if type_implements_nb_bool(obj_type):
+        result = obj.bool_impl(tx)
+        if result is not None:
+            return result
+
+    try:
+        length = generic_len(tx, obj)
+        from .tensor import SymNodeVariable
+
+        if isinstance(length, SymNodeVariable):
+            return SymNodeVariable.create(tx, length.as_proxy() > 0)
+        return ConstantVariable.create(length.as_python_constant() > 0)
+    except ObservedTypeError:
+        handle_observed_exception(tx)
+
+    return ConstantVariable.create(True)
+
+
+def vt_getitem(
+    tx: "InstructionTranslator",
+    obj: VariableTracker,
+    key: VariableTracker,
+) -> VariableTracker:
+    """CPython's PyObject_GetItem — dispatch to the type's mp_subscript/sq_item.
+
+    PyObject_GetItem: https://github.com/python/cpython/blob/62a6e898e01/Objects/abstract.c#L155-L206
+
+    CPython checks three branches in order:
+      1. tp_as_mapping->mp_subscript  (L161-166)
+      2. tp_as_sequence->sq_item      (L168-181) — only if key passes _PyIndex_Check
+      3. PyType_Check(o)              (L183-203) — type[int] → GenericAlias/__class_getitem__
+
+    Branch 1 is the common path (list, tuple, dict, range all have mp_subscript).
+    TODO(follow-up): use has_slot(map_slots, PyMappingSlots.MP_SUBSCRIPT) to gate
+    Branch 1 and has_slot(seq_slots, PySequenceSlots.SQ_ITEM) to gate Branch 2,
+    matching CPython's dispatch order.
+    TODO(follow-up): Branch 2 (sq_item) for C extension types that only have
+    tp_as_sequence (e.g. deque — Modules/_collectionsmodule.c:1888).
+    Branch 3 is handled by TypingVariable.mp_subscript_impl for typing module types
+    and by BuiltinVariable for builtin types like list[int].
+
+    Types that work via constant fold fallback (no dedicated mp_subscript_impl):
+    TODO(follow-up): str (unicode_subscript, Objects/unicodeobject.c:13809)
+    TODO(follow-up): bytes (bytes_subscript, Objects/bytesobject.c)
+    """
+    return obj.mp_subscript_impl(tx, key)
+
+
+def generic_int(tx: "InstructionTranslator", obj: VariableTracker) -> VariableTracker:
+    """Mirrors PyNumber_Long (int(x) dispatch).
+
+    https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1520-L1632
+
+    Resolution: nb_int → nb_index → str/bytes/bytearray parsing → TypeError.
+    """
+    from .constant import ConstantVariable
+
+    # Fast path for int (sub)class instances — mirrors PyLong_Check at the
+    # top of PyNumber_Long (abstract.c:1531). Avoids infinite recursion for
+    # int subclasses like IntEnum whose __int__ calls int() again.
+    if obj.is_python_constant() and isinstance(obj.as_python_constant(), int):
+        return ConstantVariable.create(int(obj.as_python_constant()))
+
+    obj_type = maybe_get_python_type(obj)
+
+    if type_implements_nb_int(obj_type):
+        return obj.nb_int_impl(tx)
+
+    if type_implements_nb_index(obj_type):
+        return obj.nb_index_impl(tx)
+
+    # String/bytes/bytearray parsing fallback.
+    # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1598-L1612
+    if obj.is_python_constant() and isinstance(
+        obj.as_python_constant(), (str, bytes, bytearray)
+    ):
+        try:
+            return ConstantVariable.create(int(obj.as_python_constant()))
+        except ValueError as e:
+            raise_observed_exception(ValueError, tx, args=[str(e)])
+
+    raise_type_error(
+        tx,
+        f"int() argument must be a string, a bytes-like object "
+        f"or a real number, not '{obj.python_type_name()}'",
+    )
+
+
+def generic_float(tx: "InstructionTranslator", obj: VariableTracker) -> VariableTracker:
+    """Mirrors PyNumber_Float (float(x) dispatch).
+
+    https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1635-L1692
+
+    Resolution: nb_float → nb_index → str parsing → TypeError.
+    """
+    from .constant import ConstantVariable
+
+    # Fast path: if the value is already a float constant, return it directly.
+    # Mirrors PyFloat_CheckExact fast path at the top of PyNumber_Float
+    # (abstract.c:1641-1643).
+    if obj.is_python_constant() and isinstance(obj.as_python_constant(), float):
+        return ConstantVariable.create(float(obj.as_python_constant()))
+
+    obj_type = maybe_get_python_type(obj)
+
+    if type_implements_nb_float(obj_type):
+        return obj.nb_float_impl(tx)
+
+    # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1674-L1685
+    if type_implements_nb_index(obj_type):
+        return obj.nb_index_impl(tx)
+
+    # PyFloat_FromString fallback — handles str and bytes.
+    # https://github.com/python/cpython/blob/v3.13.0/Objects/abstract.c#L1691
+    if obj.is_python_constant() and isinstance(obj.as_python_constant(), (str, bytes)):
+        try:
+            return ConstantVariable.create(float(obj.as_python_constant()))
+        except ValueError as e:
+            raise_observed_exception(ValueError, tx, args=[str(e)])
+
+    raise_type_error(
+        tx,
+        f"float() argument must be a string or a real number, "
+        f"not '{obj.python_type_name()}'",
+    )

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/optimizer.py

@@ -0,0 +1,425 @@
+"""
+This module implements variable tracking for PyTorch optimizers during Dynamo tracing.
+
+The OptimizerVariable class provides specialized handling for optimizer instances by:
+- Optimizing the tracing of expensive optimizer initialization
+- Managing optimizer state and parameter group tracking
+- Handling tensor sources and guards for optimizer state tensors
+- Supporting CUDA graph execution through static tensor address management
+- Providing special handling for parameter gradients and optimizer state tensors
+
+Key features include:
+- Efficient initialization tracing via _init_group optimization
+- Automatic marking of optimizer state tensors as static for CUDA graphs
+- Proper source tracking for parameter groups, gradients, and state tensors
+- Guard installation for optimizer state structure
+- Support for both CPU and GPU tensor handling
+- Cleanup of static tensor references via finalizers
+
+The module integrates with Dynamo's broader tracing system while providing
+optimizer-specific optimizations and safety guarantees.
+"""
+
+import logging
+import weakref
+from collections.abc import Iterable
+from typing import Any, TYPE_CHECKING
+
+import torch
+from torch._dynamo.variables.tensor import TensorVariable
+from torch._guards import Source
+from torch._logging import getArtifactLogger
+from torch.utils._pytree import tree_map_only
+
+from ..guards import GuardBuilder, install_guard
+from ..source import (
+    AttrSource,
+    ConstDictKeySource,
+    DictGetItemSource,
+    GetItemSource,
+    GlobalWeakRefSource,
+    GradSource,
+)
+from ..utils import GLOBAL_KEY_PREFIX
+from .base import VariableTracker
+from .constant import ConstantVariable
+from .dicts import ConstDictVariable
+from .hashable import HashableTracker
+from .lists import ListVariable
+from .misc import GetAttrVariable
+from .user_defined import UserDefinedObjectVariable
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+
+class ArgMappingException(Exception):
+    pass
+
+
+class GuardInstallException(Exception):
+    pass
+
+
+perf_hint_log = getArtifactLogger(__name__, "perf_hints")
+
+
+def _is_static_for_cudagraphs(x: torch.Tensor) -> bool:
+    from torch._inductor.cudagraph_trees import get_manager
+
+    if x.is_cuda:
+        manager = get_manager(x.device.index, False)
+        is_static_address = torch._dynamo.utils.get_static_address_type(x) is not None
+        if manager:
+            assert manager.current_node is not None
+            return (
+                is_static_address
+                or manager.current_node._is_cuda_graph_recorded_tensor(x)
+            )
+        else:
+            return is_static_address
+    else:
+        # Don't print a warning for non-cuda tensors
+        return True
+
+
+class OptimizerVariable(UserDefinedObjectVariable):
+    _nonvar_fields = {
+        "grad_to_source",
+        "tensor_to_source",
+        "static_tensor_names",
+        *UserDefinedObjectVariable._nonvar_fields,
+    }
+
+    def __init__(
+        self,
+        value: torch.optim.Optimizer,
+        grad_to_source: dict[Any, GradSource] | None = None,
+        static_tensor_names: set[str] | None = None,
+        tensor_to_source: dict[torch.Tensor, Source] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(value, **kwargs)
+        # pyrefly: ignore [bad-override]
+        self.value: torch.optim.Optimizer = value
+        self.grad_to_source = grad_to_source or {}
+        self.tensor_to_source = tensor_to_source or {}
+        self.static_tensor_names = static_tensor_names or set()
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> "VariableTracker":
+        """This is an optimization to avoid tracing the very slow initialization of the optimizer"""
+        if name == "_init_group":
+            if not hasattr(self.value, "_init_group"):
+                # Fallback: if the optimizer does not have _init_group, trace normally
+                return super().call_method(tx, name, args, kwargs)
+            try:
+                self.graph_break_if_pending_mutation(tx)
+                self.move_step_if_cpu()
+                py_args, py_kwargs = self.get_python_args(*args, **kwargs)
+                ret_val = self.value._init_group(*py_args, **py_kwargs)
+                self.map_sources_and_install_guards(tx)
+                self.update_list_args(tx, args, kwargs, py_args, py_kwargs)
+                # stash a weak_ptr to optimizer to invalidate code
+                # if the optimizer object dies
+                mangled_name = f"__optimizer_{id(self.value)}"
+                tx.store_global_weakref_by_id(mangled_name, self.value)
+                self.create_finalizer(tx)
+
+                # This is currently safe only because the only actual `ret_val`s returned
+                # by the `_init_group` of existing optimizers are properties that are invariant
+                # to the input tensors (e.g. dtype, layout). Changing these would trigger a
+                # recompilation and hence never result in the wrong specialization of `ret_val`.
+                return ConstantVariable.create(ret_val)
+            except (ArgMappingException, GuardInstallException) as _:
+                # trace normally if we can't map args or install guards correctly
+                pass
+
+        return super().call_method(tx, name, args, kwargs)
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        # Note: this allows us to intercept the call in call_method
+        # in the typical case, we return a UserMethodVariable
+        # which will directly inline
+        if name in ("_init_group"):
+            assert self.source
+            return GetAttrVariable(
+                self,
+                name,
+                py_type=type(getattr(self.value, name)),
+                source=AttrSource(self.source, name),
+            )
+
+        if name == "param_groups":
+            from ..decorators import mark_static_address
+
+            for group in self.value.param_groups:
+                for p in group["params"]:
+                    mark_static_address(p, guard=True)
+
+            self._set_capturable(tx)
+
+        return super().var_getattr(tx, name)
+
+    def graph_break_if_pending_mutation(self, tx: "InstructionTranslator") -> None:
+        # If there are pending mutations on a parameter (due to using closure)
+        # then we need to graph break to allow the python version of the parameter
+        # to update, so that running _init_group will initialize the states with
+        # the correct values
+        for g in self.value.param_groups:
+            for p in g["params"]:
+                side_effects = tx.output.side_effects
+                variable = side_effects.id_to_variable.get(id(p), None)
+                if variable and side_effects.has_pending_mutation(variable):
+                    from ..exc import unimplemented
+
+                    unimplemented(
+                        gb_type="optimizer: pending mutation on parameter",
+                        context=f"variable: {variable}, parameter: {p}",
+                        explanation="Pending mutations on a parameter (e.g. due to using closure) require a graph break.",
+                        hints=[],
+                    )
+
+    def _set_capturable(self, tx: "InstructionTranslator") -> None:
+        from . import LazyVariableTracker
+
+        # We only set capturable if params are on cuda
+        # and the state is not initialized
+        def safe_to_set_capturable(group: dict[str, Any]) -> bool:
+            all_uninitialized = True
+            all_gpu = True
+
+            for p in group.get("params", []):
+                all_gpu &= p.is_cuda or p.is_xpu
+                all_uninitialized &= p not in self.value.state
+
+            return "capturable" in group and all_uninitialized and all_gpu
+
+        # track indices to not set so we don't need to
+        # in the variable tracker realize the whole state
+        # we handle guarding the state specially
+        for group in self.value.param_groups:
+            if safe_to_set_capturable(group):
+                group["capturable"] = True
+
+        source = self.source and AttrSource(self.source, "param_groups")
+        param_groups_vt = LazyVariableTracker.realize_all(
+            VariableTracker.build(tx, self.value.param_groups, source)
+        )
+        for param_group_vt in param_groups_vt.items:
+            key = HashableTracker(ConstantVariable.create("capturable"))
+            param_group_vt.items[key] = ConstantVariable.create(True)
+
+    def get_python_args(
+        self, *args: Any, **kwargs: Any
+    ) -> tuple[list[Any], dict[str, Any]]:
+        """Get python values equivalent to the variable tracker args"""
+
+        def map_arg(arg: Any) -> Any:
+            if isinstance(arg, VariableTracker) and arg.is_python_constant():
+                return arg.as_python_constant()
+            elif isinstance(arg, ListVariable) and not arg.items:
+                # pyrefly: ignore [implicit-any]
+                return []
+            elif (
+                isinstance(arg, ConstDictVariable)
+                and isinstance(arg.source, GetItemSource)
+                and isinstance(arg.source.base, AttrSource)
+                and arg.source.base.member == "param_groups"
+            ):
+                return self.value.param_groups[arg.source.index]
+
+            raise ArgMappingException
+
+        new_args = [map_arg(arg) for arg in args]
+        new_kwargs = {k: map_arg(v) for k, v in kwargs.items()}
+
+        return new_args, new_kwargs
+
+    # If users load an old state dictionary,
+    # it's possible that step could be on the cpu
+    # if this is the case, move it to the GPU
+    # corresponding to the parameter
+    # in most cases this is a no-op because the state is empty
+    def move_step_if_cpu(self) -> None:
+        for p, state in self.value.state.items():
+            if "step" in state and state["step"].is_cpu:
+                state["step"] = state["step"].to(p.device)
+
+    def map_sources_and_install_guards(self, tx: "InstructionTranslator") -> None:
+        from ..decorators import mark_static_address
+        from .lazy import LazyVariableTracker
+
+        self.grad_to_source = {}
+        self.tensor_to_source = {}
+
+        def mark_static(x: Any) -> None:
+            mark_static_address(x, guard=True)
+
+        tree_map_only(torch.Tensor, mark_static, self.value.state)
+
+        # Recursively realize the variable trackers for optim.state and
+        # optim.param_groups, which recursively install the necessary guards.
+        params_groups_source = self.source and AttrSource(self.source, "param_groups")
+        param_groups_vt = LazyVariableTracker.realize_all(
+            VariableTracker.build(tx, self.value.param_groups, params_groups_source)
+        )
+
+        state_source = self.source and AttrSource(self.source, "state")
+        state_vt = VariableTracker.build(tx, self.value.state, state_source)
+
+        # We need to realize the top level state dict to populate
+        # the guard locals
+        state_vt.realize()
+        assert state_source is not None
+        tx.output.guard_on_key_order.add(state_source)
+
+        # Populate self.grad_to_source and self.tensor_to_source so that we can
+        # manually update_list_args
+        for group, group_vt in zip(self.value.param_groups, param_groups_vt.items):
+            # we assume here that all params within a param group
+            # are initialized similarly
+            if len(group["params"]) > 0:
+                for param in group["params"]:
+                    if param.grad is not None:
+                        key_index = None
+                        for i, k in enumerate(self.value.state.keys()):
+                            if k is param:
+                                key_index = i
+                                break
+                        if key_index:
+                            LazyVariableTracker.realize_all(
+                                VariableTracker.build(
+                                    tx,
+                                    self.value.state[param],
+                                    DictGetItemSource(
+                                        state_source,
+                                        ConstDictKeySource(state_source, key_index),
+                                    ),
+                                )
+                            )
+                            break
+
+            params_vt = group_vt.getitem_const(tx, ConstantVariable.create("params"))
+            all_static = True
+            non_static_grads = []
+            for p, p_vt in zip(group["params"], params_vt.unpack_var_sequence(tx)):
+                param_source = p_vt.source
+                self.tensor_to_source[p] = param_source
+                grad_source = GradSource(
+                    param_source,
+                    "grad",
+                )
+
+                if p.grad is not None:
+                    self.grad_to_source[p.grad] = grad_source
+                    if not _is_static_for_cudagraphs(p.grad):
+                        all_static = False
+                        non_static_grads.append(grad_source)
+                else:
+                    install_guard(grad_source.make_guard(GuardBuilder.CONSTANT_MATCH))
+
+            # Note: to avoid spam logs only warn if perf hint artifact is enabled
+            # (NB: artifacts are only enabled at the debug or warning level)
+            if not all_static and perf_hint_log.isEnabledFor(logging.DEBUG):
+                non_static_grad_names = [src.name for src in non_static_grads]
+                perf_hint_log.warning(
+                    (
+                        "Grad tensors %s will be copied during cudagraphs execution."
+                        "If using cudagraphs and the grad tensor addresses will be the same across runs,"
+                        " use torch._dynamo.decorators.mark_static_address to elide this copy.",
+                    ),
+                    non_static_grad_names,
+                )
+
+        # We have to again iterate over the state dict to collect the
+        # tensor_to_source dict. This is used for the finalizer.
+        for idx, value in enumerate(self.value.state.values()):
+            p_state_source = DictGetItemSource(
+                state_source, ConstDictKeySource(state_source, idx)
+            )
+            tx.output.guard_on_key_order.add(p_state_source)
+            for inner_idx, v in enumerate(value.values()):
+                if (
+                    isinstance(v, torch.Tensor)
+                    and v not in self.grad_to_source
+                    and v not in self.tensor_to_source
+                ):
+                    self.tensor_to_source[v] = DictGetItemSource(
+                        p_state_source, ConstDictKeySource(p_state_source, inner_idx)
+                    )
+
+    def wrap_tensor(
+        self, tx: "InstructionTranslator", tensor_value: torch.Tensor
+    ) -> TensorVariable:
+        """Wrap state tensor in a TensorVariable"""
+        from ..decorators import mark_static_address
+
+        # If we have a source for a tensor already use it,
+        # if we have not seen a tensor before, stash and use a
+        # global weak ref source, since it must be an optimizer tensor
+        # that we have missed
+
+        if tensor_value in self.tensor_to_source:
+            # mark these tensors as static for cudagraphs
+            mark_static_address(tensor_value, guard=True)
+            source = self.tensor_to_source[tensor_value]
+            self.static_tensor_names.add(tx.output.module_key_name(source.name))
+        elif tensor_value in self.grad_to_source:
+            source = self.grad_to_source[tensor_value]
+        else:
+            # mark these tensors as static for cudagraphs
+            mark_static_address(tensor_value, guard=True)
+
+            global_name = tx.store_global_weakref_by_id(GLOBAL_KEY_PREFIX, tensor_value)
+            source = GlobalWeakRefSource(global_name)
+            self.static_tensor_names.add(tx.output.module_key_name(source.name))
+
+        return VariableTracker.build(tx, tensor_value, source)
+
+    def update_list_args(
+        self,
+        tx: "InstructionTranslator",
+        args: Iterable[VariableTracker],
+        kwargs: Any,
+        py_args: Iterable[Any],
+        py_kwargs: Any,
+    ) -> None:
+        """Update the args and kwargs to the traced optimizer call"""
+        for arg, py_arg in zip(args, py_args):
+            if isinstance(arg, ListVariable):
+                assert isinstance(py_arg, list), (
+                    "py_arg should be a list in optimizer variable"
+                )
+                for i, val in enumerate(py_arg):
+                    tx.output.side_effects.mutation(arg)
+                    if isinstance(val, torch.Tensor):
+                        arg.items.append(self.wrap_tensor(tx, val))
+                    else:
+                        source = arg.source and GetItemSource(arg.source, i)
+                        arg.items.append(VariableTracker.build(tx, val, source))
+
+    def create_finalizer(self, tx: "InstructionTranslator") -> None:
+        names_to_delete = self.static_tensor_names
+        value = self.value
+        tc = tx.output.tracing_context
+
+        def init_finalizer(gm: torch.fx.GraphModule) -> None:
+            def clear_static_tensor_refs() -> None:
+                for name in names_to_delete:
+                    gm._buffers.pop(name, None)
+                    gm._parameters.pop(name, None)
+                    if tc.params_flat:
+                        tc.params_flat.clear()
+                    if tc.params_flat_unwrap_subclasses:
+                        tc.params_flat_unwrap_subclasses.clear()
+
+            weakref.finalize(value, clear_static_tensor_refs)
+
+        tx.output.add_graph_finalizer(init_finalizer)

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/script_object.py

@@ -0,0 +1,532 @@
+"""
+This module implements variable tracking for TorchScript objects during Dynamo tracing.
+
+The TorchScriptObjectVariable class provides specialized handling for TorchScript
+objects with strong safety guarantees by:
+- Enforcing method-call-only access to prevent unsafe attribute manipulation
+- Converting graph breaks into hard errors via _raise_hard_error_if_graph_break
+- Proper proxy and source tracking for TorchScript method calls
+- Integration with higher-order operators for method call handling
+
+Key safety features:
+- Strict validation that only method calls are allowed (no direct attribute access)
+- Immediate error reporting for potentially unsafe operations
+- Proper source tracking for debugging and guard installation
+- Safe handling of TorchScript object method calls through torchbind
+
+The module ensures that TorchScript objects are handled safely during tracing
+by limiting operations to known-safe patterns and failing fast for unsafe usage.
+"""
+
+import enum
+import functools
+import inspect
+import types
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, TYPE_CHECKING, TypeVar
+from typing_extensions import ParamSpec
+
+import torch
+import torch.utils._pytree as pytree
+from torch._guards import Source
+from torch._library.fake_class_registry import FakeScriptObject
+from torch._library.opaque_object import (
+    get_member_type,
+    is_opaque_reference_type,
+    is_opaque_type,
+    is_opaque_value_type,
+    MemberType,
+    should_hoist,
+)
+from torch.fx.proxy import Proxy
+
+from .. import graph_break_hints
+from ..eval_frame import skip_code
+from ..exc import (
+    raise_observed_exception,
+    unimplemented,
+    UnsafeScriptObjectError,
+    Unsupported,
+)
+from ..source import AttrSource
+from ..utils import proxy_args_kwargs
+from .base import VariableTracker
+from .constant import ConstantVariable
+from .dicts import ConstDictVariable
+from .lists import TupleVariable
+from .misc import LambdaVariable
+from .user_defined import UserDefinedObjectVariable, UserDefinedVariable
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+
+
+def _raise_hard_error_if_graph_break(
+    reason: str,
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
+    def deco(fn: Callable[_P, _T]) -> Callable[_P, _T]:
+        @functools.wraps(fn)
+        def graph_break_as_hard_error(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+            try:
+                return fn(*args, **kwargs)
+            except Unsupported as e:
+                raise UnsafeScriptObjectError(e.msg) from e
+
+        return graph_break_as_hard_error
+
+    return deco
+
+
+class OpaqueObjectClassVariable(UserDefinedVariable):
+    """
+    A variable that represents an opaque object class (not instance).
+    Since UserDefinedClassVariable has some special handling for side effects,
+    we have a separate class here which will directly return the object when
+    __init__ is called.
+    """
+
+    def __init__(self, value: Any, **kwargs: Any) -> None:
+        assert not (isinstance(value, type) and issubclass(value, enum.Enum)), (
+            f"Enum class {value} should use UserDefinedClassVariable, "
+            "not OpaqueObjectClassVariable"
+        )
+        super().__init__(**kwargs)
+        self.value = value
+
+    def as_python_constant(self) -> Any:
+        return self.value
+
+    def is_python_constant(self) -> bool:
+        # prevents constant folding of attribute accesses on
+        # opaque classes. this ensures var_getattr is called,
+        # allowing for proper validation and error handling
+        return False
+
+    def is_python_hashable(self) -> bool:
+        return is_opaque_value_type(self.value)  # pyrefly: ignore[bad-argument-type]
+
+    def get_python_hash(self) -> int:
+        return hash(self.value)
+
+    def as_proxy(self) -> Any:
+        return self.value
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.value})"
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        obj = None
+        try:
+            obj = inspect.getattr_static(self.value, name)
+        except AttributeError:
+            unimplemented(
+                gb_type="Attribute not found on opaque class",
+                context=f"class={self.value}, attr={name}",
+                explanation=f"The attribute '{name}' does not exist on opaque class {self.value}.",
+                hints=[
+                    f"Ensure '{name}' is a valid attribute of {type(self.value)}.",
+                ],
+            )
+
+        if isinstance(obj, staticmethod):
+            obj = obj.__get__(self.value)
+        elif isinstance(obj, property):
+            obj = obj.__get__(None, self.value)  # pyrefly: ignore[no-matching-overload]
+        elif hasattr(obj, "__get__"):
+            if not isinstance(type(obj).__dict__.get("__get__"), types.FunctionType):
+                # C-level descriptors are safe to resolve dynamically.
+                obj = getattr(self.value, name)
+            else:
+                type_name = type(obj).__name__
+                unimplemented(
+                    gb_type="Unsupported descriptor on opaque class",
+                    context=f"class={self.value}, attr={name}, descriptor={type_name}",
+                    explanation=f"The attribute '{name}' is a descriptor of type '{type_name}' which is not supported.",
+                    hints=[
+                        "Only staticmethod, property, and pybind11_static_property are supported.",
+                        "Consider accessing this attribute outside of the compiled region.",
+                    ],
+                )
+
+        if ConstantVariable.is_literal(obj):
+            return VariableTracker.build(tx, obj)
+
+        source = AttrSource(self.source, name) if self.source else None
+        return VariableTracker.build(tx, obj, source)
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        # disallow creating reference-type opaque objects in the middle of the
+        # program
+        if is_opaque_reference_type(self.value):
+            # Skip __init__ to prevent dynamo from tracing it during resume
+            skip_code(self.value.__init__.__code__)
+
+            unimplemented(
+                gb_type="An opaque object was created in the middle of the program.",
+                context=f"Opaque object type: {self.value}.",
+                explanation=(
+                    "Opaque objects cannot be created inside the torch.compile region. "
+                    "They must be created before entering the compiled function."
+                ),
+                hints=[
+                    "Please create the opaque object before calling torch.compile "
+                    "and pass it in as an argument or as a global variable."
+                ],
+            )
+
+        var_args = TupleVariable(list(args))
+        var_kwargs = ConstDictVariable(
+            {VariableTracker.build(tx, k): v for k, v in kwargs.items()}
+        )
+        if should_hoist(self.value):
+            with tx.output.tracing_context.guards_context.skip_guard_install():
+                constant_args = var_args.as_python_constant()
+                constant_kwargs = var_kwargs.as_python_constant()
+        else:
+            constant_args = var_args.as_python_constant()
+            constant_kwargs = var_kwargs.as_python_constant()
+        opaque_obj = self.value(  # pyrefly: ignore[not-callable]
+            *constant_args, **constant_kwargs
+        )
+
+        # Capture sources from the VT args so subgraph reuse can apply
+        # source replacement to resolve new ctor arg values on stamp-out.
+        ctor_arg_sources = tuple(getattr(a, "source", None) for a in args)
+
+        if is_opaque_value_type(type(opaque_obj)):
+            fake_script_obj = opaque_obj
+        else:
+            fake_script_obj = torch._library.fake_class_registry.maybe_to_fake_obj(
+                tx.output.fake_mode, opaque_obj
+            )
+
+        return TorchScriptObjectVariable.create(
+            opaque_obj,
+            fake_script_obj,
+            (constant_args, constant_kwargs),
+            ctor_arg_sources=ctor_arg_sources,
+        )
+
+
+class TorchScriptObjectVariable(UserDefinedObjectVariable):
+    _fake_script_object_cache: dict[int, "TorchScriptObjectVariable"] = {}
+
+    @classmethod
+    def is_matching_cls(cls, user_cls: type) -> bool:
+        return (
+            issubclass(user_cls, torch.ScriptObject)
+            or is_opaque_type(user_cls)
+            or issubclass(user_cls, FakeScriptObject)
+        )
+
+    @staticmethod
+    def create(
+        proxy: Proxy,
+        value: Any,
+        ctor_args_kwargs: Any = None,
+        ctor_arg_sources: tuple[Source | None, ...] | None = None,
+        **options: Any,
+    ) -> "TorchScriptObjectVariable":
+        assert not isinstance(value, enum.Enum), (
+            f"Enum {type(value)} should use UserDefinedObjectVariable, not TorchScriptObjectVariable"
+        )
+        out = TorchScriptObjectVariable(
+            proxy, value, ctor_args_kwargs, ctor_arg_sources=ctor_arg_sources, **options
+        )
+        if isinstance(proxy, torch.fx.Proxy) and proxy.node.op != "placeholder":
+            from torch._dynamo.symbolic_convert import InstructionTranslator
+
+            tx = InstructionTranslator.current_tx()
+            tx.output.current_tracer.record_proxyable_vt(out)
+        return out
+
+    def __init__(
+        self,
+        proxy: Proxy,
+        value: Any,
+        ctor_args_kwargs: Any = None,
+        source: Source | None = None,
+        ctor_arg_sources: tuple[Source | None, ...] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(value, **kwargs)
+        self.proxy = proxy
+        if isinstance(self.proxy, torch.fx.Proxy):
+            self.proxy.node.meta["example_value"] = value
+        self.source = source
+        # If the OpaqueObject is sourceless, then this is
+        # the constant (args, kwargs) that Dynamo used to construct it.
+        self.ctor_args_kwargs = ctor_args_kwargs
+        # Sources of the constructor args, used by subgraph reuse to
+        # resolve new values via source replacement on stamp-out.
+        self.ctor_arg_sources = ctor_arg_sources
+
+    def as_proxy(self) -> Proxy:
+        if not isinstance(self.proxy, torch.fx.Proxy):
+            # If we have a hoisted value type, then lazily lift it to be a graph
+            # input when as_proxy() is called.
+            assert is_opaque_value_type(type(self.proxy))
+            if should_hoist(type(self.proxy)):
+                from torch._dynamo.symbolic_convert import InstructionTranslator
+
+                tx = InstructionTranslator.current_tx()
+                # if any kwargs (synthetic_graph_input doesn't support them yet)
+                # not a graph break because hard error more explicit here
+                # (and opaque objects are really just used for compile)
+                if self.ctor_args_kwargs[1]:
+                    raise RuntimeError(
+                        "NYI: hoisted opaque objects that accept kwargs, please pass as args"
+                    )
+                hoisted_vt = tx.output.synthetic_graph_input(
+                    type(self.proxy),
+                    self.ctor_args_kwargs[0],
+                    ctor_arg_sources=self.ctor_arg_sources,
+                )
+                self.proxy = hoisted_vt.as_proxy()
+
+        return self.proxy
+
+    def __str__(self) -> str:
+        value = (
+            self.value.real_obj
+            if isinstance(self.value, FakeScriptObject)
+            else self.value
+        )
+        return f"{self.__class__.__name__}({value})"
+
+    __repr__ = __str__
+
+    @_raise_hard_error_if_graph_break(
+        "Dynamo cannot safely trace script object due to graph break."
+    )
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        from torch._higher_order_ops.torchbind import call_torchbind
+
+        from .higher_order_ops import TorchHigherOrderOperatorVariable
+
+        real_obj = self.as_python_constant()
+        real_obj_type = type(real_obj)
+        if is_opaque_type(real_obj_type):
+            member_type = get_member_type(real_obj_type, name)
+
+            if member_type == MemberType.USE_REAL:
+                value = getattr(real_obj, name)
+                if inspect.ismethod(value) or isinstance(
+                    value, types.MethodWrapperType
+                ):
+                    return LambdaVariable(
+                        lambda *args, **kwargs: self.call_method(tx, name, args, kwargs)
+                    )
+                else:
+                    return super().var_getattr(tx, name)
+
+            elif member_type == MemberType.INLINED:
+                value = getattr(real_obj, name)
+                if (
+                    inspect.ismethod(value)
+                    or isinstance(value, types.MethodWrapperType)
+                ) and self.source is None:
+                    # When we don't have a source, fall back to call_method
+                    # which creates a proxy node.
+                    return LambdaVariable(
+                        lambda *args, **kwargs: self.call_method(tx, name, args, kwargs)
+                    )
+                return super().var_getattr(tx, name)
+
+            elif is_opaque_value_type(real_obj_type):
+                return super().var_getattr(tx, name)
+
+            elif name in ("__bool__", "__len__") and not hasattr(real_obj, name):
+                # Special case: __bool__ and __len__ are used for truthiness checks.
+                # If they're not registered and the real object doesn't have them,
+                # raise ObservedAttributeError so the caller can fall back to
+                # treating the object as truthy (Python default behavior
+                raise_observed_exception(AttributeError, tx)
+
+            else:
+                unimplemented(
+                    gb_type="Attempted to access unregistered member on an OpaqueObject",
+                    context=f"value={real_obj}, attr={name}",
+                    explanation=f"Member '{name}' is not registered for this opaque object type.",
+                    hints=[
+                        f"Register '{name}' with a MemberType in register_opaque_type(members=...).",
+                    ],
+                )
+
+        method = getattr(self.value, name, None)
+        if method is None:
+            unimplemented(
+                gb_type="FakeScriptObject missing method implementation",
+                context=f"value={self.value}, method={name}",
+                explanation=f"TorchScript object {self.value} doesn't define the method {name}.",
+                hints=[
+                    f"Ensure the method {name} is implemented in {self.value}.",
+                    *graph_break_hints.USER_ERROR,
+                ],
+            )
+
+        if not callable(method):
+            unimplemented(
+                gb_type="Attempted to access non-callable attribute of TorchScript object",
+                context=f"value={self.value}, method={name}",
+                explanation="Attribute accesses of TorchScript objects to non-callable attributes are not supported.",
+                hints=[
+                    "Use method calls instead of attribute access.",
+                ],
+            )
+
+        assert self.source is not None
+        return TorchHigherOrderOperatorVariable.make(
+            call_torchbind,
+            source=AttrSource(self.source, name),
+            script_obj_var=self,
+            method_name=name,
+        )
+
+    def mp_subscript_impl(
+        self,
+        tx: "InstructionTranslator",
+        key: "VariableTracker",
+    ) -> "VariableTracker":
+        # Call call_method directly on this class to avoid the __getitem__ →
+        # mp_subscript_impl loop in VariableTracker.call_method.
+        return TorchScriptObjectVariable.call_method(self, tx, "__getitem__", [key], {})
+
+    # We only support method calls on script objects. Interpreting the bytecodes
+    # should go through var_getattr then call_function instead of call_method.
+
+    # However, it's possible for call_method to be used directly e.g. for __setattr__.
+    @_raise_hard_error_if_graph_break(
+        "Dynamo cannot safely trace script object due to graph break."
+    )
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: Iterable[Any],
+        kwargs: dict[str, Any],
+    ) -> VariableTracker:
+        from .builder import wrap_fx_proxy
+
+        real_obj = self.as_python_constant()
+        real_obj_type = type(real_obj)
+        if is_opaque_type(real_obj_type):
+            member_type = get_member_type(real_obj_type, name)
+
+            if member_type == MemberType.USE_REAL:
+                if (
+                    inspect.getattr_static(real_obj_type, "__getattr__", None)
+                    is not None
+                ):
+                    unimplemented(
+                        gb_type="Opaque object with custom __getattr__ not supported",
+                        context=f"{real_obj_type.__name__} with custom __getattr__",
+                        explanation="Dynamo does not support opaque objects types with custom __getattr__ methods",
+                        hints=[],
+                    )
+
+                args_const = [x.as_python_constant() for x in args]
+                kwargs_const = {k: v.as_python_constant() for k, v in kwargs.items()}
+
+                method = getattr(real_obj, name)
+
+                if name == "__setattr__":
+                    method(*args_const, **kwargs_const)
+                    return real_obj  # pyrefly: ignore[bad-return]
+
+                constant_val = method(*args_const, **kwargs_const)
+
+                if any(
+                    is_opaque_reference_type(type(r))
+                    for r in pytree.tree_leaves(constant_val)
+                ):
+                    unimplemented(
+                        gb_type="Opaque object member with method-type USE_REAL returned a reference-type opaque object.",
+                        context=f"Opaque object type: {real_obj_type}. Method name: '{name}'",
+                        explanation=(
+                            "To properly guard reference-type opaque objects, "
+                            "we must lift them as inputs to the graph. In order "
+                            "to do this, they must all have a source, meaning they "
+                            "come from a global value or are an attribute of an input."
+                        ),
+                        hints=[
+                            f"Register member '{name}' with MemberType.INLINED in "
+                            f"register_opaque_type({real_obj_type}, members=...).",
+                        ],
+                    )
+
+                return VariableTracker.build(tx, constant_val)
+
+            elif member_type == MemberType.INLINED or is_opaque_value_type(
+                real_obj_type
+            ):
+                proxy_args, proxy_kwargs = proxy_args_kwargs(args, kwargs)
+
+                proxy = tx.output.create_proxy(
+                    "call_method",
+                    name,
+                    args=(self.proxy, *proxy_args),
+                    kwargs=proxy_kwargs,
+                )
+
+                return wrap_fx_proxy(tx=tx, proxy=proxy)
+
+            else:
+                unimplemented(
+                    gb_type="Attempted to access unregistered member on an OpaqueObject",
+                    context=f"value={real_obj}, attr={name}",
+                    explanation=f"Member '{name}' is not registered for this opaque object type.",
+                    hints=[
+                        f"Register '{name}' with a MemberType in register_opaque_type(members=...).",
+                    ],
+                )
+
+        unimplemented(
+            gb_type="Weird method call on TorchScript object",
+            context=f"value={self.value}, method={name}",
+            explanation=(
+                f"This particular method call ({name}) is not supported (e.g. calling `__setattr__`). "
+                "Most method calls to TorchScript objects should be supported."
+            ),
+            hints=[
+                "Avoid calling this method.",
+            ],
+        )
+
+    def as_python_constant(self) -> Any:
+        if isinstance(self.value, FakeScriptObject):
+            return self.value.real_obj
+        elif is_opaque_value_type(type(self.value)):
+            return self.value
+        elif isinstance(self.value, torch.ScriptObject):
+            return self.value
+        return super().as_python_constant()
+
+    def is_python_hashable(self) -> bool:
+        try:
+            self.get_python_hash()
+            return True
+        except TypeError:
+            return False
+
+    def get_python_hash(self) -> int:
+        real_obj = self.as_python_constant()
+        return hash(real_obj)
+
+    def is_python_equal(self, other: object) -> bool:
+        assert isinstance(other, VariableTracker)
+        real_self = self.as_python_constant()
+        real_other = other.as_python_constant()
+        return real_self == real_other
+
+    def get_real_value(self) -> Any:
+        return self.as_python_constant()

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/sdpa.py

@@ -0,0 +1,98 @@
+from collections.abc import Sequence
+from inspect import getattr_static
+from typing import Any, TYPE_CHECKING, TypeGuard
+
+from torch._guards import Source
+from torch.backends.cuda import SDPAParams
+from torch.fx.proxy import Proxy
+
+from ..bytecode_transformation import create_call_function
+from ..exc import unimplemented
+from ..source import AttrSource
+from .base import VariableTracker
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.codegen import PyCodegen
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+PARAM_NAMES = [
+    "query",
+    "key",
+    "value",
+    "attn_mask",
+    "dropout",
+    "is_causal",
+    "enable_gqa",
+]
+
+
+class SDPAParamsVariable(VariableTracker):
+    """Represents the c++ params struct for scaled dot product attention.
+    This is a read-only container."""
+
+    @staticmethod
+    def create(
+        tx: "InstructionTranslator", value: Any, source: Source
+    ) -> VariableTracker:
+        from .torch import TorchInGraphFunctionVariable
+
+        params = [
+            VariableTracker.build(tx, getattr(value, p), AttrSource(source, p))
+            for p in PARAM_NAMES
+        ]
+        return TorchInGraphFunctionVariable(SDPAParams).call_function(tx, params, {})
+
+    def __init__(
+        self, proxy: Proxy, param_vars: Sequence[VariableTracker], **kwargs: Any
+    ) -> None:
+        self.proxy = proxy
+        self.param_vars = param_vars
+        super().__init__(**kwargs)
+
+    def python_type(self) -> type:
+        return SDPAParams
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        assert self.source is None
+        assert self.param_vars is not None
+        codegen.add_push_null(
+            lambda: codegen.load_import_from("torch._C", "_SDPAParams")
+        )
+        codegen.foreach(self.param_vars)
+        codegen.extend_output(create_call_function(len(self.param_vars), False))
+
+    def as_proxy(self) -> Proxy:
+        return self.proxy
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        import torch._C
+
+        from .builder import wrap_fx_proxy
+        from .misc import GetAttrVariable
+
+        try:
+            getattr_static(torch._C._SDPAParams, name)
+        except AttributeError:
+            import torch._dynamo.graph_break_hints as graph_break_hints
+
+            unimplemented(
+                gb_type="unsupported torch._C._SDPAParams attribute",
+                context=f"name: {name}",
+                explanation=f"Unable to fetch attribute {name} from torch._C._SDPAParams.",
+                hints=[
+                    *graph_break_hints.USER_ERROR,
+                ],
+            )
+
+        proxy = GetAttrVariable.create_getattr_proxy(self.as_proxy(), name)
+        if self.source is not None:
+            return wrap_fx_proxy(
+                tx=tx, proxy=proxy, source=AttrSource(self.source, name)
+            )
+        else:
+            return wrap_fx_proxy(tx=tx, proxy=proxy)
+
+    @staticmethod
+    def is_sdpa_params(value: Any) -> TypeGuard["SDPAParams"]:
+        return value is SDPAParams

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/sets.py

@@ -0,0 +1,797 @@
+"""
+Set-related variable tracking classes for PyTorch Dynamo.
+
+This module implements variable tracking for different types of set-like objects:
+- Regular Python sets (set)
+- Frozen sets (frozenset)
+- Ordered sets (torch.utils._ordered_set.OrderedSet)
+- Dictionary key sets (dict_keys views used as sets)
+
+These classes are responsible for tracking set operations during graph compilation,
+maintaining proper guards for set mutations and element existence checks.
+
+The implementation uses a special HashableTracker wrapper to handle set elements
+while preserving proper aliasing semantics. Sets are modeled internally as
+dictionaries with None values.
+"""
+
+import functools
+import operator
+from collections.abc import Iterable, Sequence
+from typing import Any, Literal, TYPE_CHECKING
+
+from torch.utils._ordered_set import OrderedSet
+
+from .. import polyfills, variables
+from ..bytecode_transformation import create_call_function, create_instruction
+from ..exc import raise_observed_exception
+from ..guards import GuardBuilder, install_guard
+from ..source import AttrSource, is_constant_source, is_from_local_source
+from ..utils import cmp_name_to_op_mapping, istype, raise_args_mismatch
+from .base import ValueMutationNew, VariableTracker
+from .constant import ConstantVariable
+from .hashable import HashableTracker, is_hashable, raise_unhashable
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.codegen import PyCodegen
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+    from torch._dynamo.variables.builtin import BuiltinVariable
+
+
+# [Adding a new supported class within the keys of SetVariable]
+# see steps outlined for ConstDictVariable
+
+
+class SetVariable(VariableTracker):
+    """Represents a Python set during symbolic execution."""
+
+    # PySet_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/setobject.c#L2436
+    _cpython_type = set
+
+    CONTAINS_GUARD = GuardBuilder.SET_CONTAINS
+    NOT_CONTAINS_GUARD = GuardBuilder.SET_NOT_CONTAINS
+
+    def __init__(
+        self,
+        items: Iterable[VariableTracker | HashableTracker],
+        **kwargs: Any,
+    ) -> None:
+        # .clone() passes these arguments in kwargs but they're recreated below
+        if "original_items" in kwargs:
+            kwargs.pop("original_items")
+        if "should_reconstruct_all" in kwargs:
+            kwargs.pop("should_reconstruct_all")
+
+        super().__init__(**kwargs)
+
+        # Items can be either VariableTrackers or HashableTrackers (from set ops).
+        # For VariableTrackers, realize them to ensure aliasing guards are installed
+        # when the same object appears multiple times.
+        hashable_items = []
+        for item in items:
+            if isinstance(item, HashableTracker):
+                # Already a HashableTracker from a set operation
+                hashable_items.append(item)
+            else:
+                # VariableTracker - realize to install guards, then wrap
+                # pyrefly: ignore [bad-argument-type]
+                hashable_items.append(HashableTracker(item.realize()))
+        self.items = dict.fromkeys(hashable_items, SetVariable._default_value())
+        self.should_reconstruct_all = (
+            not is_from_local_source(self.source) if self.source else True
+        )
+        self.original_items = dict.fromkeys(
+            hashable_items, SetVariable._default_value()
+        )
+
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "set()"
+        else:
+            items: list[str] = []
+            for v in self.items:
+                vt = v.vt if isinstance(v, HashableTracker) else v
+                val_str = repr(vt.value) if hasattr(vt, "value") else vt.debug_repr()
+                items.append(val_str)
+            return "{" + ",".join(items) + "}"
+
+    @property
+    def set_items(self) -> set["HashableTracker"]:
+        return set(self.items.keys())
+
+    @staticmethod
+    def _default_value() -> VariableTracker:
+        # Variable to fill in the keys of the dictionary
+        return ConstantVariable.create(None)
+
+    def as_proxy(self) -> Any:
+        return {k.vt.as_proxy() for k in self.set_items}
+
+    def python_type(self) -> type:
+        return set
+
+    def as_python_constant(self) -> Any:
+        return {k.vt.as_python_constant() for k in self.set_items}
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.foreach([x.vt for x in self.set_items])
+        codegen.append_output(create_instruction("BUILD_SET", arg=len(self.set_items)))
+
+    def __contains__(self, vt: VariableTracker) -> bool:
+        assert isinstance(vt, VariableTracker)
+        if not is_hashable(vt):
+            return False
+        key = HashableTracker(vt)
+        return key in self.items and not isinstance(
+            self.items[key], variables.DeletedVariable
+        )
+
+    def len(self) -> int:
+        return sum(
+            not isinstance(x, variables.DeletedVariable) for x in self.items.values()
+        )
+
+    def has_new_items(self) -> bool:
+        return self.should_reconstruct_all or any(
+            self.is_new_item(self.original_items.get(key.vt), value)
+            for key, value in self.items.items()
+        )
+
+    def is_new_item(
+        self, value: VariableTracker | None, other: VariableTracker
+    ) -> bool:
+        if value and value.is_realized() and other.is_realized():
+            return id(value.realize()) != id(other.realize())
+        return id(value) != id(other)
+
+    def unpack_var_sequence(self, tx: "InstructionTranslator") -> list[VariableTracker]:
+        return [x.vt for x in self.items]
+
+    def clone(self, **kwargs: Any) -> VariableTracker:
+        return super().clone(**kwargs)
+
+    def is_python_hashable(self) -> bool:
+        return False
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str):
+        if name == "__class__":
+            return VariableTracker.build(tx, self.python_type())
+        return super().var_getattr(tx, name)
+
+    def call_obj_hasattr(
+        self, tx: "InstructionTranslator", name: str
+    ) -> ConstantVariable:
+        return VariableTracker.build(tx, hasattr(set, name))
+
+    def install_set_contains_guard(
+        self, tx: "InstructionTranslator", args: list[VariableTracker]
+    ) -> None:
+        if not self.source:
+            return
+
+        if tx.output.side_effects.is_modified(self):
+            return
+
+        contains = args[0] in self
+        if args[0].source is None and args[0].is_python_constant():
+            guard_fn = (
+                type(self).CONTAINS_GUARD if contains else type(self).NOT_CONTAINS_GUARD
+            )
+            install_guard(
+                self.make_guard(
+                    functools.partial(
+                        guard_fn,
+                        key=args[0].as_python_constant(),
+                    )
+                )
+            )
+
+    def _fast_set_method(
+        self,
+        tx: "InstructionTranslator",
+        fn: Any,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        try:
+            res = fn(
+                *[x.as_python_constant() for x in [self, *args]],
+                **{k: v.as_python_constant() for k, v in kwargs.items()},
+            )
+        except Exception as exc:
+            raise_observed_exception(type(exc), tx, args=list(exc.args))
+        return VariableTracker.build(tx, res)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from ..utils import check_constant_args
+        from .builder import SourcelessBuilder
+
+        if (
+            name
+            in (
+                "isdisjoint",
+                "union",
+                "intersection",
+                "difference",
+                "symmetric_difference",
+            )
+            and check_constant_args(args, kwargs)
+            and self.python_type() is set
+        ):
+            py_type = self.python_type()
+            return self._fast_set_method(tx, getattr(py_type, name), args, kwargs)
+
+        # Lazy imports to avoid circular dependencies
+        from .dicts import DictItemsVariable, DictKeysVariable
+
+        if name == "__init__":
+            temp_set_vt = SourcelessBuilder.create(tx, set).call_set(
+                tx, *args, **kwargs
+            )
+            tx.output.side_effects.mutation(self)
+            self.items.clear()
+            self.items.update(temp_set_vt.items)  # type: ignore[attr-defined]
+            return ConstantVariable.create(None)
+        elif name == "add":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            # Convert add to __setitem__ with None value
+            if not is_hashable(args[0]):
+                raise_unhashable(args[0], tx)
+            tx.output.side_effects.mutation(self)
+            self.items[HashableTracker(args[0])] = SetVariable._default_value()
+            return ConstantVariable.create(None)
+        elif name == "pop":
+            if kwargs or args:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            # Choose an item at random and pop it
+            try:
+                result: VariableTracker = self.set_items.pop().vt  # type: ignore[assignment]
+            except KeyError as e:
+                raise_observed_exception(KeyError, tx, args=list(e.args))
+            self.should_reconstruct_all = True
+            tx.output.side_effects.mutation(self)
+            self.items.pop(HashableTracker(result))
+            return result
+        elif name == "isdisjoint":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return SourcelessBuilder.create(tx, polyfills.set_isdisjoint).call_function(
+                tx, [self, args[0]], {}
+            )
+        elif name == "intersection":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(
+                tx, polyfills.set_intersection
+            ).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "intersection_update":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(
+                tx, polyfills.set_intersection_update
+            ).call_function(tx, [self, *args], {})
+        elif name == "union":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(tx, polyfills.set_union).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "difference":
+            if kwargs:
+                raise_args_mismatch(
+                    tx, name, f"Expect: 0 kwargs, Actual: {len(kwargs)} kwargs"
+                )
+            return SourcelessBuilder.create(tx, polyfills.set_difference).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "difference_update":
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(
+                tx, polyfills.set_difference_update
+            ).call_function(tx, [self, *args], {})
+        elif name == "symmetric_difference":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return SourcelessBuilder.create(
+                tx, polyfills.set_symmetric_difference
+            ).call_function(
+                tx,
+                [self, *args],
+                {"cls": self.python_type_var()},
+            )
+        elif name == "symmetric_difference_update":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return SourcelessBuilder.create(
+                tx, polyfills.set_symmetric_difference_update
+            ).call_function(tx, [self, *args], {})
+        elif name == "update" and self.is_mutable():
+            if kwargs:
+                raise_args_mismatch(tx, name, "0 kwargs", f"{len(kwargs)} kwargs")
+            return SourcelessBuilder.create(tx, polyfills.set_update).call_function(
+                tx, [self, *args], {}
+            )
+        elif name == "remove":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            if args[0] not in self:
+                raise_observed_exception(KeyError, tx, args=args)
+            self.should_reconstruct_all = True
+            tx.output.side_effects.mutation(self)
+            self.items.pop(HashableTracker(args[0]))
+            return ConstantVariable.create(None)
+        elif name == "discard":
+            if kwargs or len(args) != 1:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            if args[0] in self:
+                self.should_reconstruct_all = True
+                tx.output.side_effects.mutation(self)
+                self.items.pop(HashableTracker(args[0]))
+            return ConstantVariable.create(None)
+        elif name in ("issubset", "issuperset"):
+            if len(args) != 1:
+                raise_args_mismatch(tx, name, "1 args", f"{len(args)} args")
+
+            op = {
+                "issubset": operator.le,
+                "issuperset": operator.ge,
+            }
+            other = args[0].realize()
+            if not istype(other, SetVariable):
+                other = SourcelessBuilder.create(tx, set).call_function(tx, [other], {})
+            return SourcelessBuilder.create(tx, op.get(name)).call_function(
+                tx, [self, other], {}
+            )
+        elif name in ("__and__", "__or__", "__xor__", "__sub__"):
+            m = {
+                "__and__": "intersection",
+                "__or__": "union",
+                "__xor__": "symmetric_difference",
+                "__sub__": "difference",
+            }.get(name)
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for {name}: '{self.python_type_name()}' and '{args[0].python_type_name()}'"
+                    ],
+                )
+            assert m is not None
+            return self.call_method(tx, m, args, kwargs)
+        elif name in ("__rand__", "__ror__", "__rxor__", "__rsub__"):
+            m = {
+                "__rand__": "__and__",
+                "__ror__": "__or__",
+                "__rxor__": "__xor__",
+                "__rsub__": "__sub__",
+            }.get(name)
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for {name}: '{args[0].python_type_name()}' and '{self.python_type_name()}'"
+                    ],
+                )
+            assert m is not None
+            return args[0].call_method(tx, m, [self], kwargs)
+        elif name in ("__iand__", "__ior__", "__ixor__", "__isub__"):
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                raise_observed_exception(
+                    TypeError,
+                    tx,
+                    args=[
+                        f"unsupported operand type(s) for {name}: '{self.python_type_name()}' and '{args[0].python_type_name()}'"
+                    ],
+                )
+            m = {
+                "__iand__": "intersection_update",
+                "__ior__": "update",
+                "__ixor__": "symmetric_difference_update",
+                "__isub__": "difference_update",
+            }.get(name)
+            assert m is not None
+            self.call_method(tx, m, args, kwargs)
+            return self
+        elif name == "__eq__":
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return ConstantVariable.create(False)
+            r = self.call_method(tx, "symmetric_difference", args, kwargs)
+            return VariableTracker.build(tx, len(r.set_items) == 0)  # type: ignore[attr-defined]
+        elif name == "__ne__":
+            eq_result = self.call_method(tx, "__eq__", args, kwargs)
+            return VariableTracker.build(tx, not eq_result.value)  # type: ignore[attr-defined]
+        elif name in cmp_name_to_op_mapping:
+            if not isinstance(
+                args[0],
+                (
+                    SetVariable,
+                    variables.UserDefinedSetVariable,
+                    DictItemsVariable,
+                    DictKeysVariable,
+                ),
+            ):
+                return VariableTracker.build(tx, NotImplemented)
+            return VariableTracker.build(
+                tx,
+                cmp_name_to_op_mapping[name](self.set_items, args[0].set_items),  # type: ignore[attr-defined]
+            )
+        elif name == "__contains__":
+            if not len(args):
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "more than 1 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            if not (args and is_hashable(args[0])):
+                raise_unhashable(args[0], tx)
+            self.install_set_contains_guard(tx, args)
+            contains = args[0] in self
+            return VariableTracker.build(tx, contains)
+        elif name == "__len__":
+            if args or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return VariableTracker.build(tx, len(self.items))
+        elif name == "copy":
+            if args or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            return self.clone(
+                items=self.items.copy(), mutation_type=ValueMutationNew(), source=None
+            )
+        elif name == "clear":
+            if args or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "0 args and 0 kwargs",
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+            self.should_reconstruct_all = True
+            tx.output.side_effects.mutation(self)
+            self.items.clear()
+            return ConstantVariable.create(None)
+        elif name == "__iter__":
+            from .lists import ListIteratorVariable
+
+            if self.source and not is_constant_source(self.source):
+                tx.output.guard_on_key_order.add(self.source)
+            return ListIteratorVariable(
+                self.unpack_var_sequence(tx), mutation_type=ValueMutationNew()
+            )
+        return super().call_method(tx, name, args, kwargs)
+
+    def python_type_var(self) -> "BuiltinVariable":
+        return variables.BuiltinVariable(set)
+
+    def getitem_const(
+        self, tx: "InstructionTranslator", arg: VariableTracker
+    ) -> VariableTracker:
+        raise RuntimeError("Illegal to getitem on a set")
+
+    def sq_length(self, tx: "InstructionTranslator") -> VariableTracker:
+        return VariableTracker.build(tx, len(self.set_items))
+
+
+class OrderedSetClassVariable(VariableTracker):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+    def as_python_constant(self) -> type[OrderedSet[Any]]:
+        return OrderedSet
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        if name == "__new__":
+            from .misc import GetAttrVariable
+
+            if self.source:
+                attr_source = AttrSource(self.source, name)
+            else:
+                attr_source = None
+            return GetAttrVariable(
+                self, name, py_type=type(getattr(OrderedSet, name)), source=attr_source
+            )
+        else:
+            return super().var_getattr(tx, name)
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from .builtin import set_methods
+
+        if name == "__new__":
+            if len(args) != 2 or kwargs:
+                raise_args_mismatch(
+                    tx,
+                    name,
+                    "OrderedSet.__new__ only accepts one arg"
+                    f"{len(args)} args and {len(kwargs)} kwargs",
+                )
+
+            return variables.OrderedSetVariable([], mutation_type=ValueMutationNew())
+
+        resolved_fn = getattr(set, name)
+        if resolved_fn in set_methods and isinstance(args[0], variables.SetVariable):
+            return args[0].call_method(tx, name, args[1:], kwargs)
+
+        return super().call_method(tx, name, args, kwargs)
+
+    def call_function(
+        self,
+        tx: "InstructionTranslator",
+        args: Sequence[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> "OrderedSetVariable":
+        if len(args) > 1 or kwargs:
+            raise_args_mismatch(
+                tx,
+                "OrderedSet",
+                "OrderedSet only accepts one arg"
+                f"{len(args)} args and {len(kwargs)} kwargs",
+            )
+
+        if len(args) == 0:
+            # pyrefly: ignore [implicit-any]
+            items = []
+        else:
+            items = args[0].force_unpack_var_sequence(tx)
+        return variables.OrderedSetVariable(items, mutation_type=ValueMutationNew())
+
+
+class OrderedSetVariable(SetVariable):
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "OrderedSet([])"
+        else:
+            items: list[str] = []
+            for k in self.items:
+                key_str = (
+                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
+                )
+                items.append(key_str)
+            return "OrderedSet([" + ",".join(items) + "])"
+
+    def as_python_constant(self) -> OrderedSet[Any]:
+        return OrderedSet([k.vt.as_python_constant() for k in self.set_items])
+
+    def python_type(self) -> type[OrderedSet[Any]]:
+        return OrderedSet
+
+    # pyrefly: ignore[bad-override]
+    def python_type_var(self) -> OrderedSetClassVariable:
+        return OrderedSetClassVariable()
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.load_import_from("torch.utils._ordered_set", "OrderedSet")
+        )
+        codegen.foreach([x.vt for x in self.set_items])
+        codegen.append_output(create_instruction("BUILD_LIST", arg=len(self.set_items)))
+        codegen.extend_output(create_call_function(1, False))
+
+
+class FrozensetVariable(SetVariable):
+    # PyFrozenSet_Type: https://github.com/python/cpython/blob/v3.13.0/Objects/setobject.c#L2526
+    _cpython_type = frozenset
+
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "frozenset()"
+        else:
+            items: list[str] = []
+            for k in self.items:
+                key_str = (
+                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
+                )
+                items.append(key_str)
+            return "{" + ",".join(items) + "}"
+
+    @property
+    def set_items(self) -> set["HashableTracker"]:
+        return set(self.items.keys())
+
+    def python_type(self) -> type:
+        return frozenset
+
+    def python_type_var(self) -> "BuiltinVariable":
+        return variables.BuiltinVariable(frozenset)
+
+    def as_python_constant(self) -> Any:
+        return frozenset({k.vt.as_python_constant() for k in self.set_items})
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        codegen.add_push_null(
+            lambda: codegen.extend_output(
+                [
+                    codegen.create_load_global("frozenset"),
+                ]
+            )
+        )
+        codegen.foreach([x.vt for x in self.set_items])
+        codegen.extend_output(
+            [
+                create_instruction("BUILD_LIST", arg=len(self.set_items)),
+                *create_call_function(1, False),
+            ]
+        )
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name in ["add", "pop", "update", "remove", "discard", "clear"]:
+            raise RuntimeError(f"Illegal call_method {name} on a frozenset")
+        elif name == "__init__":
+            # frozenset is immutable. Calling __init__ again shouldn't have any effect
+            return ConstantVariable.create(None)
+        elif name in (
+            "copy",
+            "difference",
+            "intersection",
+            "symmetric_difference",
+        ):
+            r = super().call_method(tx, name, args, kwargs)
+            return FrozensetVariable(r.items)  # type: ignore[attr-defined]
+        return super().call_method(tx, name, args, kwargs)
+
+    def is_python_hashable(self) -> Literal[True]:
+        """
+        Frozensets are immutable and hashable in Python.
+        """
+        return True
+
+    def get_python_hash(self) -> int:
+        return hash(self.as_python_constant())
+
+    def is_python_equal(self, other: object) -> bool:
+        return (
+            isinstance(other, VariableTracker)
+            and self.as_python_constant() == other.as_python_constant()
+        )
+
+
+class DictKeySetVariable(SetVariable):
+    def debug_repr(self) -> str:
+        if not self.items:
+            return "dict_keys([])"
+        else:
+            items: list[str] = []
+            for k in self.items:
+                key_str = (
+                    repr(k.vt.value) if hasattr(k.vt, "value") else k.vt.debug_repr()
+                )
+                items.append(key_str)
+            return "dict_keys([" + ",".join(items) + "])"
+
+    def install_set_contains_guard(
+        self, tx: "InstructionTranslator", args: list[VariableTracker]
+    ) -> None:
+        # Already EQUALS_MATCH guarded
+        pass
+
+    @property
+    def set_items(self) -> Any:
+        return self.items
+
+    def python_type(self) -> type:
+        from ..utils import dict_keys
+
+        return dict_keys
+
+    def as_python_constant(self) -> Any:
+        return dict.fromkeys(
+            {k.vt.as_python_constant() for k in self.set_items}, None
+        ).keys()
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        if name in ["add", "pop", "update", "remove", "discard", "clear"]:
+            raise RuntimeError(f"Illegal call_method {name} on a dict_keys")
+        return super().call_method(tx, name, args, kwargs)

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/streams.py

@@ -0,0 +1,724 @@
+import collections
+from collections.abc import Callable
+from typing import Any, Optional
+
+import torch
+from torch._dynamo.variables.dicts import ConstDictVariable
+from torch._dynamo.variables.lists import TupleVariable
+from torch.fx import has_side_effect, Proxy
+
+from .. import graph_break_hints
+from ..bytecode_transformation import create_call_function
+from ..exc import TYPE_CHECKING, unimplemented
+from ..graph_bytecode_inputs import (
+    CURRENT_STREAM_INDEX,
+    get_external_object_by_index,
+    register_graph_created_object,
+    register_user_object,
+    reset_user_object_tracking,
+)
+from ..source import CurrentStreamSource
+from .base import VariableTracker
+from .constant import ConstantVariable
+from .ctx_manager import FxTracebackAnnotateVariable
+from .lazy import LazyVariableTracker
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+    from ..codegen import PyCodegen
+
+from torch._library.custom_ops import custom_op
+
+
+Tensor = torch.Tensor
+
+
+def new_event(*args: Any, **kwargs: Any) -> int:
+    event = torch.Event(*args, **kwargs)
+    return register_graph_created_object(
+        event,
+        EventVariable.make_construct_in_graph_event_fn(
+            TupleVariable([]), ConstDictVariable({})
+        ),
+    )
+
+
+def new_stream(*args: tuple[Any], **kwargs: Any) -> int:
+    stream = torch.Stream(*args, **kwargs)  # type: ignore[no-matching-overload,call-overload]
+    return register_graph_created_object(
+        stream,
+        StreamVariable.make_construct_in_graph_stream_fn(
+            TupleVariable([]), ConstDictVariable({})
+        ),
+    )
+
+
+def _codegen_current_stream(device: torch.device, cg: "PyCodegen") -> None:
+    cg.add_push_null(
+        lambda: cg.load_import_from(
+            torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
+            "stash_graph_created_object",
+        )
+    )
+    cg(CurrentStreamSource(device))
+    cg.extend_output(create_call_function(1, False))
+
+
+def get_current_stream(device: torch.device) -> int:
+    stream = torch.accelerator.current_stream(device)
+    return register_graph_created_object(
+        stream, lambda _, cg: _codegen_current_stream(device, cg)
+    )
+
+
+def _get_stream_by_index(index: int) -> torch.Stream:
+    stream = get_external_object_by_index(index)
+    assert isinstance(stream, torch.Stream), (
+        f"Fork/join stream expected a stream object at index {index}"
+    )
+    return stream
+
+
+def _get_event_by_index(index: int) -> torch.Event:
+    event = get_external_object_by_index(index)
+    assert isinstance(event, torch.Event), (
+        f"Record/wait event expected an event object at index {index}"
+    )
+    return event
+
+
+@custom_op("streams::fork", mutates_args=())
+def fork_stream(
+    from_index: int,  # kept to make stream transitions clearer
+    to_index: int,
+) -> None:
+    torch.accelerator.set_stream(_get_stream_by_index(to_index))
+
+
+@fork_stream.register_fake
+def _(
+    from_index: int,  # kept to make stream transitions clearer
+    to_index: int,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.fork.default)
+
+
+@custom_op("streams::join", mutates_args=())
+def join_stream(from_index: int, to_index: int) -> None:
+    torch.accelerator.set_stream(_get_stream_by_index(to_index))
+
+
+@join_stream.register_fake
+def _(
+    from_index: int,
+    to_index: int,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.join.default)
+
+
+@custom_op("streams::record_event", mutates_args=())
+def record_event(event_index: int, stream_index: int) -> None:
+    event = _get_event_by_index(event_index)
+    stream = _get_stream_by_index(stream_index)
+    event.record(stream)
+
+
+@record_event.register_fake
+def _(
+    event_index: int,
+    stream_index: int,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.record_event.default)
+
+
+@custom_op("streams::wait_event", mutates_args=())
+def wait_event(event_index: int, stream_index: int) -> None:
+    event = _get_event_by_index(event_index)
+    stream = _get_stream_by_index(stream_index)
+    event.wait(stream)
+
+
+@wait_event.register_fake
+def _(
+    event_index: int,
+    stream_index: int,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.wait_event.default)
+
+
+@custom_op("streams::synchronize_event", mutates_args=())
+def synchronize_event(event_index: int) -> None:
+    event = _get_event_by_index(event_index)
+    event.synchronize()
+
+
+@synchronize_event.register_fake
+def _(event_index: int) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.synchronize_event.default)
+
+
+@custom_op("streams::synchronize_device", mutates_args=())
+def synchronize_device(device_type: str, device_index: int) -> None:
+    torch.accelerator.synchronize(torch.device(device_type, device_index))
+
+
+@synchronize_device.register_fake
+def _(device_type: str, device_index: int) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.synchronize_device.default)
+
+
+@custom_op("streams::synchronize_stream", mutates_args=())
+def synchronize_stream(stream_index: int) -> None:
+    stream = _get_stream_by_index(stream_index)
+    stream.synchronize()
+
+
+@synchronize_stream.register_fake
+def _(stream_index: int) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.synchronize_stream.default)
+
+
+@custom_op("streams::wait_stream", mutates_args=())
+def wait_stream(waiting_stream_index: int, waited_on_stream_index: int) -> None:
+    waiting = _get_stream_by_index(waiting_stream_index)
+    waited_on = _get_stream_by_index(waited_on_stream_index)
+    waiting.wait_stream(waited_on)
+
+
+@wait_stream.register_fake
+def _(
+    event_index: int,
+    stream_index: int,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.wait_stream.default)
+
+
+@custom_op("streams::sync_dealloc", mutates_args=())
+def sync_dealloc(
+    wait_event_index: int, src_stream_index: int, to_dealloc: torch.Tensor
+) -> None:
+    """An op which waits on an event and moves the last usage of to_dealloc
+    after the wait, so that after the sync occurs, the deallocation or
+    subsequent reuse of the tensor's memory will be guaranteed to happen
+    after a side stream is finished using it.
+    See https://docs.pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html#torch.Tensor.record_stream
+    for more details"""
+    torch.ops.streams.wait_event.default(wait_event_index, src_stream_index)
+
+
+@sync_dealloc.register_fake
+def _(
+    wait_event_index: int,
+    src_stream_index: int,
+    to_dealloc: torch.Tensor,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.sync_dealloc.default)
+
+
+@custom_op("streams::record_stream", mutates_args=())
+def record_stream(tensor: torch.Tensor, stream_index: int) -> None:
+    tensor.record_stream(_get_stream_by_index(stream_index))
+
+
+@record_stream.register_fake
+def _(
+    tensor: torch.Tensor,
+    stream_index: int,
+) -> None:
+    pass
+
+
+has_side_effect(torch.ops.streams.record_stream.default)
+
+
+class SymbolicStreamState:
+    """Track the currently entered stream if any"""
+
+    def __init__(self) -> None:
+        from ..source import CurrentStreamSource
+
+        cur_stack: list[StreamVariable] = []
+        if torch.accelerator.is_available():
+            # Reset the registry so the current stream is guaranteed index 0.
+            reset_user_object_tracking()
+            stream = torch.accelerator.current_stream()
+            source = CurrentStreamSource(stream.device)
+            # Register the current stream so it gets index 0 (registry is
+            # fresh at tracing start).  The inductor wrapper updates this
+            # entry at runtime so cudagraph capture uses the capture stream
+            # instead of this stale trace-time stream.
+            index = register_user_object(stream, source)
+            assert index == CURRENT_STREAM_INDEX, (
+                f"Current stream must be registered at index {CURRENT_STREAM_INDEX}, "
+                f"got {index}"
+            )
+            stream_var = LazyVariableTracker.create(stream, source=source)
+            # Set user_object_index as an instance attribute so accessing it
+            # does NOT trigger LazyVariableTracker realization.
+            stream_var.user_object_index = index  # type: ignore[union-attr]
+            cur_stack = [stream_var]  # type: ignore[list-item]
+
+        self.cur_stream_stack: collections.deque[StreamVariable] = collections.deque(
+            cur_stack
+        )
+
+    def enter_stream(self, stream: "StreamVariable") -> None:
+        self.cur_stream_stack.append(stream)
+
+    def exit_stream(self) -> None:
+        self.cur_stream_stack.pop()
+
+    def cur_stream(self, device: torch.device | None = None) -> "StreamVariable":
+        if device is not None:
+            for stream in reversed(self.cur_stream_stack):
+                if stream.device == device:
+                    return stream
+
+        return self.cur_stream_stack[-1]
+
+    def in_stream_context(self) -> bool:
+        return len(self.cur_stream_stack) > 0
+
+    def cur_stream_id(self) -> int:
+        """Get a Python object id for the current stream without realizing lazy variables."""
+        stream = self.cur_stream_stack[-1]
+        if isinstance(stream, LazyVariableTracker) and not stream.is_realized():
+            return id(stream.peek_value())
+        return id(stream.value)
+
+
+class StreamContextVariable(FxTracebackAnnotateVariable):
+    """This represents torch.cuda.StreamContext"""
+
+    @staticmethod
+    def create(
+        tx: "InstructionTranslator",
+        stream_to_enter: "StreamVariable",
+        **kwargs: dict[str, Any],
+    ) -> "StreamContextVariable":
+        return StreamContextVariable(
+            stream_to_enter,
+            **kwargs,
+        )
+
+    def __init__(self, stream: Optional["StreamVariable"], **kwargs: Any) -> None:
+        self.stream = stream
+        super().__init__(
+            target_values={"stream": self.get_stream().user_object_index},
+            initial_values=None,
+            **kwargs,
+        )
+
+    def enter(
+        self, tx: "InstructionTranslator", *args: VariableTracker
+    ) -> VariableTracker:
+        # to stream, from stream is the order of the arguments
+        # we are entering the target, and leaving the initial stream
+        tx.symbolic_stream_state.enter_stream(self.get_stream())
+        return super().enter(tx)
+
+    def exit(
+        self, tx: "InstructionTranslator", *args: VariableTracker
+    ) -> VariableTracker:
+        # to stream, from stream is the order of the arguments
+        # we are leaving the target, and entering the initial stream
+        tx.symbolic_stream_state.exit_stream()
+        return super().exit(tx, *args)
+
+    def python_type(self) -> type:
+        return torch.cuda.StreamContext
+
+    def supports_graph_breaks(self) -> bool:
+        return True
+
+    def get_stream(self) -> "StreamVariable":
+        assert self.stream, "Stream context should have a separate stream"
+        return self.stream
+
+
+class StreamVariable(StreamContextVariable):
+    """Represents the device-agnostic torch.Stream class"""
+
+    _cpython_type = torch.Stream
+
+    def __init__(
+        self,
+        proxy: Proxy,
+        value: torch.Stream,
+        user_object_index: int | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Index into the user object table
+        # used to pass arbitrary objects to the graph
+        if proxy is not None and "example_value" in proxy.node.meta:
+            assert proxy.node.meta["example_value"] == value
+
+        self.proxy = proxy
+        self.value = value
+        self.device = value.device
+
+        self.user_object_index = user_object_index
+        super().__init__(None, **kwargs)
+
+    def python_type(self) -> type:
+        return torch.Stream
+
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        assert hasattr(self.value, name), f"no stream method found named {name}"
+
+        from ..utils import cmp_name_to_op_mapping, proxy_args_kwargs
+        from .builder import wrap_fx_proxy_cls
+
+        if name == "wait_event":
+            event_arg = args[0]
+            assert isinstance(event_arg, EventVariable)
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.wait_event,
+                (event_arg.user_object_index, self.user_object_index),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "wait_stream":
+            other_stream = args[0]
+            assert isinstance(other_stream, StreamVariable)
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.wait_stream,
+                (self.user_object_index, other_stream.user_object_index),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "synchronize":
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.synchronize_stream,
+                (self.user_object_index,),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "query":
+            return wrap_fx_proxy_cls(
+                target_cls=ConstantVariable,
+                tx=tx,
+                proxy=tx.output.create_proxy(
+                    "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
+                ),
+            )
+        elif name == "record_event":
+            from .builder import wrap_fx_proxy
+
+            tx.output.check_event_record_after_input_mutation(id(self.value))
+            if args and isinstance(args[0], EventVariable):
+                event_var = args[0]
+                event = event_var.value
+                event_index = event_var.user_object_index
+            else:
+                event = self.value.record_event()
+                event_index = register_graph_created_object(
+                    event,
+                    EventVariable.make_construct_in_graph_event_fn(
+                        TupleVariable([]), ConstDictVariable({})
+                    ),
+                )
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.record_event,
+                (event_index, self.user_object_index),
+                {},
+            )
+            return wrap_fx_proxy(
+                tx=tx,
+                proxy=tx.output.create_proxy(
+                    "call_function",
+                    get_external_object_by_index,
+                    (event_index,),
+                    {},
+                ),
+            )
+        elif name in cmp_name_to_op_mapping and len(args) == 1 and not kwargs:
+            from ..guards import GuardBuilder, install_guard
+
+            if self.source:
+                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))
+
+            # NB : Checking for mutation is necessary because we compare
+            # constant values
+            other = args[0]
+            if not isinstance(other, StreamVariable):
+                return VariableTracker.build(tx, NotImplemented)
+
+            if other.source:
+                assert self.source is not None
+                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))
+            return VariableTracker.build(
+                tx,
+                cmp_name_to_op_mapping[name](self.value, other.value),  # type: ignore[arg-type]
+            )
+
+        return super().call_method(tx, name, args, kwargs)
+
+    def as_proxy(self) -> Proxy:
+        return self.proxy
+
+    def module_name(self) -> str:
+        return "torch._C"
+
+    def fn_name(self) -> str:
+        return "Stream"
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        # If we got here, this stream is fully subsumed by the graph - this means it is
+        # not an input or global
+        assert not self.source
+        if self.user_object_index is not None:
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.graph_bytecode_inputs.__name__,
+                    "get_external_object_by_index",
+                )
+            )
+            codegen.append_output(codegen.create_load_const(self.user_object_index))
+            codegen.extend_output(create_call_function(1, False))
+        else:
+            # This will support the legacy behavior
+            prefix = f"_stream_{self.device}"
+            name = codegen.tx.output.install_global_by_id(prefix, self.value)
+            codegen.append_output(codegen.create_load_global(name, add=True))
+
+    def get_stream(self) -> "StreamVariable":
+        return self
+
+    @staticmethod
+    def make_construct_in_graph_stream_fn(
+        args: TupleVariable, kwargs: ConstDictVariable
+    ) -> Callable[[int, "PyCodegen"], None]:
+        def fn(index: int, codegen: "PyCodegen") -> None:
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
+                    "stash_graph_created_object",
+                )
+            )
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.utils.__name__, "build_stream"
+                )
+            )
+            codegen(args)
+            codegen(kwargs)
+            codegen.extend_output(create_call_function(2, False))
+            codegen.extend_output(create_call_function(1, False))
+
+        return fn
+
+
+class CudaStreamVariable(StreamVariable):
+    """Represents torch.cuda.Stream, preserving device-specific type and attributes."""
+
+    _cpython_type = torch.cuda.Stream
+
+    def python_type(self) -> type:
+        return torch.cuda.Stream
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> "VariableTracker":
+        from . import ConstantVariable
+
+        if name == "cuda_stream":
+            from ..guards import GuardBuilder, install_guard
+
+            if self.source:
+                install_guard(self.source.make_guard(GuardBuilder.EQUALS_MATCH))
+
+            if hasattr(self.value, "cuda_stream"):
+                return ConstantVariable.create(self.value.cuda_stream)
+
+            if hasattr(self.value, "native_handle"):
+                return ConstantVariable.create(self.value.native_handle)
+
+        return super().var_getattr(tx, name)
+
+
+class EventVariable(VariableTracker):
+    def __init__(
+        self,
+        proxy: Proxy,
+        value: torch.Event,
+        user_object_index: int | None,
+        **kwargs: Any,
+    ) -> None:
+        if proxy is not None and "example_value" in proxy.node.meta:
+            assert proxy.node.meta["example_value"] == value
+        super().__init__(**kwargs)
+        self.proxy = proxy
+        self.value = value
+        self.user_object_index = user_object_index
+
+    def python_type(self) -> type:
+        return torch.Event
+
+    def get_real_python_backed_value(self) -> object:
+        return self.value
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> VariableTracker:
+        from ..utils import proxy_args_kwargs
+        from .builder import wrap_fx_proxy_cls
+
+        if name == "wait":
+            _, stream_index = EventVariable._get_stream_arg(tx, args, kwargs)
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.wait_event,
+                (
+                    self.user_object_index,
+                    stream_index,
+                ),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "record":
+            stream_arg, stream_index = EventVariable._get_stream_arg(tx, args, kwargs)
+            tx.output.check_event_record_after_input_mutation(id(stream_arg.value))
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.record_event,
+                (
+                    self.user_object_index,
+                    stream_index,
+                ),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "synchronize":
+            tx.output.create_proxy(
+                "call_function",
+                torch.ops.streams.synchronize_event,
+                (self.user_object_index,),
+                {},
+            )
+            return ConstantVariable.create(None)
+        elif name == "query":
+            return wrap_fx_proxy_cls(
+                target_cls=ConstantVariable,
+                tx=tx,
+                proxy=tx.output.create_proxy(
+                    "call_method", name, *proxy_args_kwargs([self] + args, kwargs)
+                ),
+            )
+        else:
+            method_name = (
+                f"{type(self.value).__module__}.{type(self.value).__qualname__}.{name}"
+            )
+            unimplemented(
+                gb_type="Unsupported event method",
+                context=str(name),
+                explanation=f"Dynamo doesn't support tracing the {method_name} method. "
+                f"We currently support wait, record, synchronize, and query.",
+                hints=[
+                    *graph_break_hints.SUPPORTABLE,
+                ],
+            )
+
+    def as_proxy(self) -> Proxy:
+        return self.proxy
+
+    @staticmethod
+    def _get_stream_arg(
+        tx: "InstructionTranslator",
+        args: list[VariableTracker],
+        kwargs: dict[str, VariableTracker],
+    ) -> tuple["StreamVariable", int]:
+        """Returns (stream_variable, stream_index_for_op).
+
+        The ambient current stream is registered at index 0 in the external
+        object registry.  The inductor wrapper updates index 0 at runtime so
+        that cudagraph capture sees the capture stream, not the stale
+        trace-time default stream.
+        """
+        stream_arg = None
+        if args:
+            stream_arg = args[0]
+        elif kwargs:
+            stream_arg = kwargs.get("stream")
+
+        if not stream_arg:
+            stream_var = tx.symbolic_stream_state.cur_stream()
+            return stream_var, stream_var.user_object_index  # type: ignore[return-value]
+
+        return stream_arg, stream_arg.user_object_index  # type: ignore[return-value]
+
+    @staticmethod
+    def make_construct_in_graph_event_fn(
+        args: TupleVariable, kwargs: ConstDictVariable
+    ) -> Callable[[int, "PyCodegen"], None]:
+        def fn(index: int, codegen: "PyCodegen") -> None:
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
+                    "stash_graph_created_object",
+                )
+            )
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.utils.__name__, "build_event"
+                )
+            )
+            codegen(args)
+            codegen(kwargs)
+            codegen.extend_output(create_call_function(2, False))
+            codegen.extend_output(create_call_function(1, False))
+
+        return fn
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        # If we got here, this event is fully subsumed by the graph - this means it is
+        # not an input or global
+        assert not self.source
+        # Similar to stream handling, we lift the event into a global and then codegen bytecode to load it from there.
+        prefix = "_event"
+        name = codegen.tx.output.install_global_by_id(prefix, self.value)
+        codegen.append_output(codegen.create_load_global(name, add=True))

kronos-venv/lib/python3.12/site-packages/torch/_dynamo/variables/torch_function.py

@@ -0,0 +1,781 @@
+"""TorchDynamo support for __torch_function__ tensor subclasses.
+
+This module implements support for tensor subclasses with __torch_function__ overrides.
+A tensor subclass instance is represented as a TensorWithTFOverrideVariable, which handles
+dispatching __torch_function__ on attribute accesses, method calls, and torch API calls.
+
+Unsupported features:
+- Triggering __torch_function__ on tensor subclass non-tensor custom attributes
+- Graph breaking on mutating guardable tensor properties within a __torch_function__ context
+  (can cause excessive recompiles in certain cases)
+- Matching exact eager behavior of ignoring __torch_function__ objects in non-tensor
+  argument positions of Torch API calls
+
+Supported features:
+- Static method implementations of __torch_function__ on custom objects (triggers on torch
+  API calls with the object as any argument)
+- Triggering __torch_function__ on torch API calls with tensor subclass arguments
+- __torch_function__ calls on base tensor attribute access and method calls for tensor
+  subclass instances
+- Matches dispatch ordering behavior of eager __torch_function__ with subclass/object
+  arguments in any position
+
+See https://docs.google.com/document/d/1WBxBSvW3NXhRp9ncmtokJloMLCtF4AYNhJaffvHe8Kw/edit#heading=h.vacn73lozd9w
+for more information on the design.
+"""
+
+import collections
+import contextlib
+import functools
+import inspect
+import operator
+from collections.abc import Generator, Iterable, Sequence
+from types import TracebackType
+from typing import Any, TYPE_CHECKING
+
+import torch._C
+import torch.utils._pytree as pytree
+from torch._guards import Source
+from torch.overrides import (
+    _get_overloaded_args,
+    get_default_nowrap_functions,
+    TorchFunctionMode,
+)
+from torch.utils._device import DeviceContext
+
+from .. import graph_break_hints
+from ..exc import unimplemented
+from ..guards import GuardBuilder, install_guard
+from ..polyfills import NoEnterTorchFunctionMode
+from ..source import AttrSource, GlobalSource, TorchFunctionModeStackSource, TypeSource
+from ..utils import (
+    class_has_getattribute,
+    clear_torch_function_mode_stack,
+    get_safe_global_name,
+    has_torch_function,
+    is_tensor_base_attr_getter,
+    set_torch_function_mode_stack,
+)
+from .base import VariableTracker
+from .constant import ConstantVariable
+from .ctx_manager import GenericContextWrappingVariable
+from .functions import UserMethodVariable
+from .lazy import LazyVariableTracker
+from .lists import TupleVariable
+from .tensor import TensorSubclassVariable, TensorVariable
+from .user_defined import UserDefinedObjectVariable
+
+
+if TYPE_CHECKING:
+    from torch._dynamo.codegen import PyCodegen
+    from torch._dynamo.symbolic_convert import InstructionTranslator
+
+
+bin_ops = [
+    operator.pow,
+    operator.mul,
+    operator.matmul,
+    operator.floordiv,
+    operator.truediv,
+    operator.mod,
+    operator.add,
+    operator.lt,
+    operator.gt,
+    operator.ge,
+    operator.le,
+    operator.ne,
+    operator.eq,
+    operator.sub,
+    operator.ipow,
+    operator.imul,
+    operator.imatmul,
+    operator.ifloordiv,
+    operator.itruediv,
+    operator.imod,
+    operator.iadd,
+    operator.isub,
+]
+
+bin_int_ops = [
+    operator.and_,
+    operator.or_,
+    operator.xor,
+    operator.iand,
+    operator.ixor,
+    operator.ior,
+]
+
+un_int_ops = [operator.invert]
+
+tensor_and_int_ops = [
+    operator.lshift,
+    operator.rshift,
+    operator.ilshift,
+    operator.irshift,
+    operator.getitem,
+]
+
+un_ops = [
+    operator.abs,
+    operator.pos,
+    operator.neg,
+    operator.not_,  # Note: this has a local scalar dense call
+    operator.length_hint,
+]
+
+
+banned_attrs = [
+    fn.__self__.__name__  # type: ignore[attr-defined]
+    for fn in get_default_nowrap_functions()
+    if is_tensor_base_attr_getter(fn)
+]
+
+
+@functools.cache
+def get_prev_stack_var_name() -> str:
+    from ..bytecode_transformation import unique_id
+
+    return unique_id("___prev_torch_function_mode_stack")
+
+
+class TorchFunctionModeVariable(GenericContextWrappingVariable):
+    @staticmethod
+    def is_supported_torch_function_mode(ty: type[TorchFunctionMode]) -> bool:
+        # Supported in this sense means we can support graph breaks under the
+        # context.
+        # We are able to trace custom modes but if there are graph breaks under them
+        # and they have a custom __enter__/__exit__ we don't handle this for the
+        # same reason we don't handle generic context managers: there may be side effects
+        # that are now affected by executing the function across two frames instead of one
+        # Today we support the enter/exit of the default TorchFunctionMode as well as
+        # DeviceContext (which is used for set_default_device)
+        return issubclass(ty, (NoEnterTorchFunctionMode, DeviceContext)) or (
+            not class_has_getattribute(ty)
+            and inspect.getattr_static(ty, "__enter__") is TorchFunctionMode.__enter__
+            and inspect.getattr_static(ty, "__exit__") is TorchFunctionMode.__exit__
+        )
+
+    def __init__(
+        self,
+        value: TorchFunctionMode | None,
+        source: Source | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if value is not None:
+            super().__init__(value, **kwargs)
+        self.value = value
+        # needed for BC with calling enter from CM code
+        self.cm_obj = value  # type: ignore[assignment]
+        self.source = source  # type: ignore[assignment]
+
+    def reconstruct(self, codegen: "PyCodegen") -> None:
+        # This shouldn't be called unless we have a source
+        assert self.source
+        self.source.reconstruct(codegen)
+
+    def module_name(self) -> str:
+        return self.value.__module__
+
+    def fn_name(self) -> str:
+        return type(self.value).__name__
+
+    def python_type(self) -> type:
+        return type(self.value)
+
+    def call_torch_function(
+        self,
+        tx: "InstructionTranslator",
+        fn: VariableTracker,
+        types: TupleVariable,
+        args: Iterable[Any],
+        kwargs: dict[str, Any],
+    ) -> VariableTracker:
+        return call_torch_function(
+            tx,
+            get_torch_function_fn(tx, self),  # type: ignore[arg-type]
+            fn,
+            types,
+            args,
+            kwargs,
+        )
+
+    def enter(self, tx: "InstructionTranslator") -> VariableTracker:
+        from .torch import TorchInGraphFunctionVariable
+
+        if isinstance(self.value, NoEnterTorchFunctionMode):
+            return ConstantVariable.create(None)
+
+        TorchInGraphFunctionVariable(
+            torch._C._push_on_torch_function_stack
+        ).call_function(tx, [self], {})
+        return ConstantVariable.create(None)
+
+    def exit(self, tx: "InstructionTranslator", *args: Any) -> VariableTracker:
+        from .torch import TorchInGraphFunctionVariable
+
+        TorchInGraphFunctionVariable(torch._C._pop_torch_function_stack).call_function(
+            tx, [], {}
+        )
+        return ConstantVariable.create(None)
+
+    def reconstruct_type(self, codegen: "PyCodegen") -> None:
+        ty = NoEnterTorchFunctionMode
+        codegen(
+            AttrSource(
+                codegen.tx.import_source(ty.__module__),
+                ty.__name__,
+            )
+        )
+
+    def supports_graph_breaks(self) -> bool:
+        return True
+
+    def exit_on_graph_break(self) -> bool:
+        return False
+
+
+# Used to clear/restore the python torch function mode stack and temporarily restore it as needed
+class TorchFunctionModeStackStateManager:
+    def __init__(self) -> None:
+        self.stack: list[Any] = []
+
+    def __enter__(self) -> None:
+        self.stack = torch.overrides._get_current_function_mode_stack()
+        clear_torch_function_mode_stack()
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        set_torch_function_mode_stack(self.stack)
+        self.stack = []
+
+    @contextlib.contextmanager
+    def temp_restore_stack(self) -> Generator[None, None, None]:
+        prev = torch.overrides._get_current_function_mode_stack()
+        set_torch_function_mode_stack(self.stack)
+        try:
+            yield
+        finally:
+            set_torch_function_mode_stack(prev)
+
+
+torch_function_mode_stack_state_mgr = TorchFunctionModeStackStateManager()
+
+
+class SymbolicTorchFunctionState:
+    def __init__(self, py_stack: Iterable[Any]) -> None:
+        # This is annoyingly complicated because of how the torch function subclass + mode C API was designed
+        # There are two exposed C knobs here as contexts: torch._C.DisableTorchFunction and torch._C.DisableTorchFunctionSubclass
+        # These are their definitions:
+        # 1) torch._C._is_torch_function_enabled indicates that neither of the above knobs have been entered
+        # (if either are entered, this will be False)
+        # 2) torch._C._is_torch_function_mode_enabled indicates that either the torch mode stack is empty OR
+        # torch._C.DisableTorchFunction has been entered
+        # To disambiguate these and keep myself sane I added a C API to check whether all torch function
+        # concepts (modes and subclasses) are enabled.
+        # This only returns true iff we have not entered torch._C.DisableTorchFunction and allows us to separate
+        # the stack length from the enablement state of torch function modes.
+        # This is important because now if a mode is pushed while dynamo is tracing, we know whether
+        # or not torch function modes are enabled and whether we should trace it.
+        self.torch_function_subclass_enabled = torch._C._is_torch_function_enabled()
+
+        # This differs from the C API of the same name
+        # this will only be false iff we have entered torch._C.DisableTorchFunction
+        # and does not take into account the mode stack length, while the C API bundles these
+        # two concepts
+        self.torch_function_mode_enabled = (
+            not torch._C._is_torch_function_all_disabled()
+        )
+
+        self.cur_mode = None
+
+        TorchFunctionModeStackVariable.reset()
+
+        self.mode_stack: collections.deque[TorchFunctionModeVariable] = (
+            collections.deque()
+        )
+
+        for i, val in enumerate(py_stack):
+            self.mode_stack.append(
+                LazyVariableTracker.create(val, source=TorchFunctionModeStackSource(i))  # type: ignore[arg-type]
+            )
+
+    def in_torch_function_mode(self) -> bool:
+        return len(self.mode_stack) > 0
+
+    def pop_torch_function_mode(self) -> TorchFunctionModeVariable:
+        return self.mode_stack.pop()
+
+    def push_torch_function_mode(self, mode_var: TorchFunctionModeVariable) -> None:
+        self.mode_stack.append(mode_var)
+
+    def call_torch_function_mode(
+        self,
+        tx: "InstructionTranslator",
+        fn: VariableTracker,
+        types: TupleVariable,
+        args: Iterable[Any],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        with self._pop_mode_for_inlining() as cur_mode:
+            return cur_mode.call_torch_function(tx, fn, types, args, kwargs)
+
+    @contextlib.contextmanager
+    def _pop_mode_for_inlining(
+        self,
+    ) -> Generator[TorchFunctionModeVariable, None, None]:
+        old_mode = self.cur_mode
+        self.cur_mode = self.pop_torch_function_mode()  # type: ignore[assignment]
+        try:
+            yield self.cur_mode  # type: ignore[misc]
+        finally:
+            mode = self.cur_mode
+            self.cur_mode = old_mode
+            self.push_torch_function_mode(mode)  # type: ignore[arg-type]
+
+
+class TorchFunctionModeStackVariable(VariableTracker):
+    """Fake VT to use as a dummy object, indicating the presence of torch function mode stack mutation"""
+
+    # singleton value representing the global torch function mode stack
+    # singleton (it exists in C++)
+    stack_value_singleton = object()
+
+    # offset is used to track if we have inserted/removed a
+    # device context which is always placed at the bottom of the stack
+    # if a device context is inserted, the graph will run this mutation
+    # so when we want to reconstruct any other modes on the stack
+    # their indices should be shifted right by 1 (+1)
+    # Conversely, if there was a device context on the stack, and the graph
+    # mutates the stack to remove that context (set default device to None)
+    # each of the indices of other modes should be shifted left by 1 (-1)
+    offset = 0
+
+    def __init__(
+        self,
+        source: Source,
+        symbolic_stack: collections.deque[TorchFunctionModeVariable],
+    ) -> None:
+        self.source = source
+        self.symbolic_stack = symbolic_stack
+
+    @classmethod
+    def reset(cls) -> None:
+        cls.offset = 0
+
+    @classmethod
+    def register_mutation(cls, tx: "InstructionTranslator") -> None:
+        if cls.stack_value_singleton not in tx.output.side_effects:
+            var = cls(
+                source=Source(),
+                symbolic_stack=tx.symbolic_torch_function_state.mode_stack,
+            )
+            tx.output.side_effects.track_mutable(cls.stack_value_singleton, var)
+            tx.output.side_effects.mutation(var)
+
+    @classmethod
+    def register_device_context_insertion(cls, tx: "InstructionTranslator") -> None:
+        stack = tx.symbolic_torch_function_state.mode_stack
+        if stack and cls.is_device_context(stack[0]):
+            return
+        else:
+            cls.offset += 1
+            stack.insert(
+                0,
+                TorchFunctionModeVariable(
+                    None, source=TorchFunctionModeStackSource(-cls.offset)
+                ),
+            )
+
+    @classmethod
+    def clear_default_device(cls, tx: "InstructionTranslator") -> None:
+        stack = tx.symbolic_torch_function_state.mode_stack
+        if stack and cls.is_device_context(stack[0]):
+            stack.popleft()
+            cls.offset -= 1
+
+    @staticmethod
+    def is_device_context(var: TorchFunctionModeVariable) -> bool:
+        return isinstance(var.value, DeviceContext) or var.value is None
+
+    @classmethod
+    def get_mode_index(cls, ind: int) -> int:
+        return ind + cls.offset
+
+
+def _get_all_args(
+    args: Iterable[Any], kwargs: dict[str, Any]
+) -> Iterable[VariableTracker]:
+    return _flatten_vts(pytree.arg_tree_leaves(*args, **kwargs))
+
+
+def _flatten_vts(vts: Iterable[VariableTracker]) -> list[VariableTracker]:
+    from collections import deque
+
+    from .dicts import ConstDictVariable
+    from .lists import ListVariable
+
+    vts = deque(vts)
+    output = []
+
+    while vts:
+        vt = vts.popleft()
+
+        if not vt.is_realized() and vt.peek_type() in (dict, list, tuple):  # type: ignore[attr-defined]
+            vt.realize()
+
+        if vt.is_realized():
+            if isinstance(vt, ListVariable):
+                vts.extend(vt.items)
+                continue
+            elif isinstance(vt, ConstDictVariable):
+                vts.extend(vt.items.values())
+                continue
+
+        output.append(vt)
+
+    return output
+
+
+def _get_subclass_type(var: VariableTracker) -> type:
+    assert isinstance(var, (TensorWithTFOverrideVariable, UserDefinedObjectVariable))
+    return var.python_type()
+
+
+def _get_subclass_type_var(
+    tx: "InstructionTranslator", var: VariableTracker
+) -> VariableTracker:
+    if isinstance(var, TensorWithTFOverrideVariable):
+        return var.class_type_var(tx)
+    elif isinstance(var, UserDefinedObjectVariable):
+        source = var.source and TypeSource(var.source)
+        return VariableTracker.build(tx, var.python_type(), source)
+    else:
+        raise AssertionError(f"Unexpected type {type(var)}")
+
+
+def _is_attr_overridden(
+    tx: "InstructionTranslator", var: VariableTracker, name: str
+) -> bool:
+    if not isinstance(var, (TensorWithTFOverrideVariable, UserDefinedObjectVariable)):
+        return False
+    import torch
+
+    overridden = False
+    try:
+        attr_val = inspect.getattr_static(var.python_type(), name)
+        overridden |= attr_val != getattr(torch.Tensor, name)
+    except AttributeError:
+        pass
+
+    return overridden
+
+
+def call_torch_function(
+    tx: "InstructionTranslator",
+    torch_function_var: VariableTracker,
+    fn: VariableTracker,
+    types: TupleVariable,
+    args: Iterable[Any],
+    kwargs: dict[str, Any],
+    *,
+    is_subclass_dispatch: bool = False,
+) -> Any:
+    # This emulates calling __torch_function__, which has a signature
+    #   def __torch_function__(cls, func, types, args=(), kwargs=None):
+    #
+    # Also notice the `cls` is not explicitly passed in the reference
+    # implementations:
+    # 1. https://github.com/pytorch/pytorch/blob/8d81806211bc3c0ee6c2ef235017bacf1d775a85/torch/csrc/utils/python_arg_parser.cpp#L368-L374  # noqa: B950
+    # 2. https://github.com/pytorch/pytorch/blob/8d81806211bc3c0ee6c2ef235017bacf1d775a85/torch/overrides.py#L1741-L1743
+    tf_args = [
+        fn,
+        types,
+        VariableTracker.build(tx, tuple(args)),
+        VariableTracker.build(tx, kwargs),
+    ]
+    # Mirror the C++ THPModule_disable_torch_function behavior: disable
+    # __torch_function__ subclass dispatch during the call to prevent
+    # re-entrant dispatch on operations inside __torch_function__.
+    # Only do this for subclass dispatch, not mode dispatch. Modes need
+    # subclass dispatch to remain enabled because the mode's
+    # __torch_function__ may re-dispatch to the subclass.
+    tf_state = tx.symbolic_torch_function_state
+    old_subclass_enabled = tf_state.torch_function_subclass_enabled
+    if is_subclass_dispatch and old_subclass_enabled:
+        tf_state.torch_function_subclass_enabled = False
+    try:
+        return torch_function_var.call_function(tx, tf_args, {})
+    finally:
+        tf_state.torch_function_subclass_enabled = old_subclass_enabled
+
+
+def get_torch_function_fn(
+    tx: "InstructionTranslator", vt: VariableTracker
+) -> VariableTracker:
+    # The underlying function could be a classmethod, staticmethod, regular
+    # function or a function with C-implementation. It doesn't matter as long as
+    # they satisfy the calling convention in `call_torch_function`.
+
+    args = [vt, VariableTracker.build(tx, "__torch_function__")]
+    func_vt = VariableTracker.build(tx, getattr).call_function(tx, args, {})
+    return func_vt
+
+
+def can_dispatch_torch_function(
+    tx: "InstructionTranslator", args: Iterable[Any], kwargs: dict[str, Any]
+) -> bool:
+    has_overridden_args = any(
+        has_torch_function(arg) for arg in _get_all_args(args, kwargs)
+    )
+    tf_state = tx.symbolic_torch_function_state
+    return (has_overridden_args and tf_state.torch_function_subclass_enabled) or (
+        tf_state.torch_function_mode_enabled and tf_state.in_torch_function_mode()
+    )
+
+
+def dispatch_torch_function(
+    tx: "InstructionTranslator",
+    fn: VariableTracker,
+    args: Iterable[Any],
+    kwargs: dict[str, Any],
+) -> Any:
+    """Gathers all args that are TensorWithTFOverrideVariable and dispatches based on the ordering in _get_overloaded_args"""
+
+    all_args = _get_all_args(args, kwargs)
+    overloaded_args = _get_overloaded_args(
+        [arg for arg in all_args if has_torch_function(arg)],
+        _get_subclass_type,
+    )
+
+    types = TupleVariable([_get_subclass_type_var(tx, arg) for arg in overloaded_args])
+
+    if tx.symbolic_torch_function_state.in_torch_function_mode():
+        res = tx.symbolic_torch_function_state.call_torch_function_mode(
+            tx, fn, types, args, kwargs
+        )
+        if not res.is_constant_match(NotImplemented):
+            return res
+
+    for arg in overloaded_args:
+        res = arg.call_torch_function(
+            tx,
+            fn,
+            types,
+            args,
+            kwargs,
+        )
+
+        if not res.is_constant_match(NotImplemented):
+            tx.output.torch_function_subclass_inlined = True
+            return res
+
+    unimplemented(
+        gb_type="All __torch_function__ overrides returned NotImplemented due to TypeError from user code",
+        context=f"{fn=}, {args=}, {kwargs=}",
+        explanation=f"All __torch_function__ overrides for for function {fn} returned NotImplemented",
+        hints=[
+            *graph_break_hints.USER_ERROR,
+        ],
+    )
+
+
+class TensorWithTFOverrideVariable(TensorVariable):
+    """
+    Represents a tensor subclass instance with a __torch_function__ override.
+    """
+
+    @classmethod
+    def from_tensor_var(
+        cls,
+        tx: "InstructionTranslator",
+        tensor_var: VariableTracker,
+        class_type: type,
+        cls_source: Source | None,
+    ) -> "TensorWithTFOverrideVariable":
+        # [Note: __torch_function__] coerce `tensor_var` into a
+        # TensorWithTFOverrideVariable. In eager, this is just a type change.
+        import torch
+
+        # This simulates shallow-copying the tensor object.
+        kwargs = dict(tensor_var.__dict__)
+        input_tensor_type = kwargs.pop("class_type")
+        assert input_tensor_type in (torch.Tensor, torch.nn.Parameter) or issubclass(
+            input_tensor_type, torch.Tensor
+        ), (
+            f"invalid class type {input_tensor_type} in TensorWithTFOverrideVariable.from_tensor_var"
+        )
+        var = cls(class_type=class_type, **kwargs)
+        var.install_global(tx)
+        return var
+
+    def install_global(self, tx: "InstructionTranslator") -> None:
+        # stash the subclass type to rewrap an output tensor if needed
+        # this is needed because the actual type needs to be available
+        # each time the compiled artifact is run and outputs a wrapped tensor.
+        if self.global_mangled_class_name(tx) not in tx.output.global_scope:
+            # Safe because global_mangled_class_name figures it out
+            tx.output.install_global_unsafe(
+                self.global_mangled_class_name(tx), self.class_type
+            )
+
+    def python_type(self) -> type:
+        return self.class_type
+
+    def class_type_var(self, tx: "InstructionTranslator") -> VariableTracker:
+        return TensorSubclassVariable(
+            self.class_type, source=GlobalSource(self.global_mangled_class_name(tx))
+        )
+
+    def global_mangled_class_name(self, tx: "InstructionTranslator") -> str:
+        return get_safe_global_name(
+            tx, f"__subclass_{self.class_type.__name__}", self.class_type
+        )
+
+    def var_getattr(self, tx: "InstructionTranslator", name: str) -> VariableTracker:
+        # [Note: __torch_function__] We currently only support attributes that are defined on
+        # base tensors, custom attribute accesses will graph break.
+        import torch
+
+        # I think only `_base` is breaking because we aren't modelling view
+        # relationship perfectly in some scenarios.
+        if name in banned_attrs:
+            unimplemented(
+                gb_type="Unsupported tensor subclass attribute access",
+                context=f"{name}",
+                explanation="`torch.compile` currently can't trace this",
+                hints=[
+                    f"Avoid accessing {name} of tensor subclass in torch.compile region",
+                    *graph_break_hints.SUPPORTABLE,
+                ],
+            )
+
+        # Handle non-overridden attributes inherited from `torch.Tensor`.
+        attr_is_overridden = _is_attr_overridden(tx, self, name)
+        if (
+            hasattr(torch.Tensor, name)
+            and not attr_is_overridden
+            and not inspect.ismethoddescriptor(getattr(torch.Tensor, name))
+        ):
+            args = [self]
+            kwargs: dict[Any, Any] = {}
+            if can_dispatch_torch_function(tx, args, kwargs):
+                get_fn = VariableTracker.build(tx, getattr(torch.Tensor, name).__get__)
+
+                return self.call_torch_function(
+                    tx,
+                    get_fn,
+                    TupleVariable([self.class_type_var(tx)]),
+                    args,
+                    kwargs,
+                )
+        else:
+            # `TensorVariable.var_getattr` doesn't handle user-defined
+            # function/attribute well, so we explicitly handle them here.
+            #
+            # TODO move this logic into `TensorVariable`, or try to merge it
+            # with similar logic in `UserDefinedObjectVariable`.
+            try:
+                attr = inspect.getattr_static(self.class_type, name)
+            except AttributeError:
+                pass
+            else:
+                import types
+
+                cls_source = GlobalSource(self.global_mangled_class_name(tx))
+                attr_source = AttrSource(cls_source, name)
+                if isinstance(attr, types.FunctionType):
+                    install_guard(attr_source.make_guard(GuardBuilder.CLOSURE_MATCH))
+                    return UserMethodVariable(attr, self)
+
+                elif isinstance(attr, property):
+                    getter_source = AttrSource(attr_source, "fget")
+                    getter = attr.fget
+                    getter_var = VariableTracker.build(
+                        tx, getter, source=getter_source, realize=True
+                    )
+                    return getter_var.call_function(tx, [self], {})
+
+                elif isinstance(attr, classmethod):
+                    return UserMethodVariable(
+                        attr.__func__, self.class_type_var(tx), source=attr_source
+                    )
+
+                elif attr_is_overridden:
+                    unimplemented(
+                        gb_type="Unsupported tensor subclass overridden attribute access",
+                        context=f"{name}",
+                        explanation="`torch.compile` only support tracing certain types of overridden tensor subclass attributes",
+                        hints=[
+                            f"Avoid accessing {name} of tensor subclass in torch.compile region",
+                            f"Renaming attribute `{name}` of type {self.class_type}",
+                            *graph_break_hints.SUPPORTABLE,
+                        ],
+                    )
+
+        return super().var_getattr(tx, name)
+
+    def call_torch_function(
+        self,
+        tx: "InstructionTranslator",
+        fn: VariableTracker,
+        types: TupleVariable,
+        args: Iterable[Any],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        # NOTE this assumes `__torch_function__` isn't modified during tracing.
+        if not hasattr(self, "torch_function_fn"):
+            self.torch_function_fn = get_torch_function_fn(tx, self)
+
+        return call_torch_function(
+            tx,
+            self.torch_function_fn,
+            fn,
+            types,
+            args,
+            kwargs,
+            is_subclass_dispatch=True,
+        )
+
+    def call_method(
+        self,
+        tx: "InstructionTranslator",
+        name: str,
+        args: Sequence[VariableTracker],
+        kwargs: "dict[str, VariableTracker]",
+    ) -> "VariableTracker":
+        # This code block implements inlining the __torch_function__ override
+        # of `call_method`.
+        tf_args = [self] + list(args)
+        if can_dispatch_torch_function(tx, tf_args, kwargs):
+            import torch
+
+            if _is_attr_overridden(tx, self, name):
+                unimplemented(
+                    gb_type="Tensor subclass overridden method call",
+                    context=f"{name}",
+                    explanation="`torch.compile` currently can't trace this",
+                    hints=[
+                        f"Avoid calling {name} of tensor subclass in torch.compile region",
+                        f"Renaming method `{name}` of type {self.class_type}",
+                        *graph_break_hints.SUPPORTABLE,
+                    ],
+                )
+
+            # [Note: __torch_function__] Currently we only support methods that are defined on tensor
+            # we will graph break in other cases this will need a bigger overhaul of extracting methods/comparing them for equality
+            # We've established with the above check that the method is not overridden, so we guard that the method is the same
+            # as the impl defined on tensor and retrieve it
+            if self.source:
+                source = AttrSource(AttrSource(self.source, "__class__"), name)
+                value = inspect.getattr_static(self.python_type(), name)
+            else:
+                source = None
+                value = getattr(torch.Tensor, name)
+            func_var = VariableTracker.build(tx, value, source)
+            return dispatch_torch_function(tx, func_var, tf_args, kwargs)
+        else:
+            return super().call_method(tx, name, args, kwargs)