# Owner(s): ["module: custom-operators"]

import contextlib
import enum
import gc
import random
import re
import unittest
from contextlib import ExitStack
from dataclasses import dataclass

import torch
import torch.distributed as dist
import torch.utils._pytree as pytree
from torch._dynamo.functional_export import _dynamo_graph_capture_for_export
from torch._dynamo.testing import (
    AotEagerAndRecordGraphs,
    CompileCounter,
    CompileCounterWithBackend,
    EagerAndRecordGraphs,
    InductorAndRecordGraphs,
    normalize_gm,
)
from torch._dynamo.utils import counters, counters as dynamo_counters
from torch._functorch import config as functorch_config
from torch._functorch.aot_autograd import (
    aot_compile_joint_with_descriptors,
    aot_export_joint_with_descriptors,
    aot_export_module,
)
from torch._inductor import config as inductor_config
from torch._inductor.compile_fx import compile_fx, compile_fx_inner
from torch._inductor.test_case import run_tests, TestCase
from torch._inductor.utils import fresh_inductor_cache, run_and_get_code
from torch._library.effects import EffectType
from torch._library.fake_class_registry import FakeScriptObject, maybe_to_fake_obj
from torch._library.opaque_object import (
    _OPAQUE_TYPES,
    _OPAQUE_TYPES_BY_NAME,
    get_opaque_type_name,
    is_opaque_type,
    is_opaque_value_type,
    MemberType,
    OpaqueBase,
    register_opaque_type,
)
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.fx._graph_pickler import GraphPickler, Options
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch.fx.graph import _illegal_char_regex
from torch.testing._internal.common_utils import (
    instantiate_parametrized_tests,
    parametrize,
)
from torch.testing._internal.inductor_utils import HAS_GPU
from torch.utils._import_utils import import_dill


class Color(OpaqueBase):
    """Simulates a pybind11-style enum where class attributes are instances of the class."""

    def __init__(self, name: str, value: int) -> None:
        self._name = name
        self._value = value

    @property
    def name(self) -> str:
        return self._name

    @property
    def value(self) -> int:
        return self._value

    def __float__(self) -> float:
        return float(self._value)

    @staticmethod
    def RED_STATIC() -> "Color":
        return Color.RED


Color.RED = Color("RED", 1)
Color.GREEN = Color("GREEN", 2)
Color.BLUE = Color("BLUE", 3)
Color.DEFAULT_SCALE = 1.5  # Literal class attribute for testing inlining


class CustomDescriptor:
    def __get__(self, obj, objtype=None):
        return 42


# Create a class with an unsupported descriptor
class ColorWithDescriptor(OpaqueBase):
    def __init__(self, name: str, value: int) -> None:
        self._name = name
        self._value = value

    def __float__(self) -> float:
        return float(self._value)

    # This is an unsupported descriptor type
    custom_prop = CustomDescriptor()


ColorWithDescriptor.RED = ColorWithDescriptor("RED", 1)


class OpaqueQueue(OpaqueBase):
    def __init__(self, queue: list[torch.Tensor], init_tensor_: torch.Tensor) -> None:
        super().__init__()
        self.queue = queue
        self.init_tensor_ = init_tensor_

        # For testing purposes
        self._push_counter = 0
        self._pop_counter = 0
        self._size_counter = 0

    def push(self, tensor: torch.Tensor) -> None:
        self._push_counter += 1
        self.queue.append(tensor)

    def pop(self) -> torch.Tensor:
        self._pop_counter += 1
        if len(self.queue) > 0:
            return self.queue.pop(0)
        return self.init_tensor_

    def size(self) -> int:
        self._size_counter += 1
        return len(self.queue)


class NestedQueue(OpaqueBase):
    def __init__(self, q):
        self.q = q

    def get_q(self):
        return self.q

    def pop_q(self):
        return torch.ops._TestOpaqueObject.queue_pop(self.q)


class RNGState(OpaqueBase):
    def __init__(self, seed):
        self.seed = seed
        self.rng = random.Random(self.seed)
        self.dummy = lambda x: x  # test something not pickleable

    def get_seed(self):
        return self.seed

    def noisy_inject(self, x):
        return torch.ops._TestOpaqueObject.noisy_inject(x, self)


class OpaqueMultiplier(OpaqueBase):
    """Opaque object that holds a multiplier value for backward tests."""

    def __init__(self, multiplier: float):
        self.multiplier = multiplier


class Counter(OpaqueBase):
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def __eq__(self, other):
        return (
            isinstance(other, Counter)
            and self.start == other.start
            and self.end == other.end
        )

    def __hash__(self):
        return hash((self.start, self.end))

    def get_start_tensor(self):
        return torch.scalar_tensor(self.start, dtype=torch.int64)

    @property
    def counter(self):
        return torch.scalar_tensor(self.start, dtype=torch.int64)

    def increment_counter(self):
        self.start += 1


class NestedCounters(OpaqueBase):
    def __init__(self, c):
        self.c = c

    def __eq__(self, other):
        if not isinstance(other, NestedCounters):
            return False
        return self.c == other.c

    def get_c(self):
        return self.c

    def get_starts(self):
        if isinstance(self.c, list):
            return [c.start for c in self.c]
        else:
            return self.c.start

    def __getitem__(self, idx):
        counter = self.c[idx]
        # Create a new counter to match device mesh's __getitem__
        return Counter(counter.start, counter.end)


class AddModule(OpaqueBase, torch.nn.Module):
    def forward(self, x, y):
        return x * y


class ValueConfig(OpaqueBase):
    def __init__(self, mode: str):
        self.mode = mode

    def __eq__(self, other):
        return isinstance(other, ValueConfig) and self.mode == other.mode

    def __hash__(self):
        return hash(self.mode)

    def __fx_repr__(self):
        return f"ValueConfig(mode={self.mode!r})", {"ValueConfig": ValueConfig}

    def print_mode(self):
        print(self.mode)


class SizeStore(OpaqueBase):
    def __init__(self, size: int):
        self.size = size

    def __eq__(self, other):
        return isinstance(other, SizeStore) and self.size == other.size

    def __hash__(self):
        return hash(self.size)

    def __fx_repr__(self):
        # Return (repr_string, dict_mapping_name_to_type)
        return f"SizeStore(size={self.size!r})", {"SizeStore": SizeStore}

    def increment_size(self):
        return self.size + 1


class NestedValueSize(OpaqueBase):
    def __init__(self, size: SizeStore, config: ValueConfig):
        self.size = size
        self.config = config

    def __eq__(self, other):
        return self.size == other.size and self.config == other.config

    def __hash__(self):
        return hash(self.size) ^ hash(self.config)

    def __fx_repr__(self):
        # Recursively call __fx_repr__ on nested opaque objects
        size_eval, size_globals = self.size.__fx_repr__()
        config_eval, config_globals = self.config.__fx_repr__()

        # Combine repr and globals
        repr_str = f"NestedValueSize(size={size_eval}, config={config_eval})"
        all_globals = (
            {"NestedValueSize": NestedValueSize} | size_globals | config_globals
        )

        return repr_str, all_globals


class HoistedString(torch._opaque_base.OpaqueBase):
    def __init__(self, val):
        self.val = val

    def __eq__(self, other):
        return self.val == other.val

    def __hash__(self):
        return hash(self.val)

    def __fx_repr__(self):
        return (f"HoistedString('{self.val}')", {"HoistedString": HoistedString})


register_opaque_type(HoistedString, typ="value", hoist=True)


@torch.library.custom_op("mylib::op_with_string", mutates_args=())
def op_with_string(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
    if s.val == "double":
        return x * 2
    elif s.val == "square":
        return x**2
    raise AssertionError("expected double or square")


@op_with_string.register_fake
def _(x, s):
    return torch.empty_like(x)


register_opaque_type(OpaqueQueue, typ="reference")
register_opaque_type(
    RNGState,
    typ="reference",
    guard_fn=lambda obj: [obj.seed],
    members={
        "seed": MemberType.USE_REAL,
        "get_seed": MemberType.USE_REAL,
        "noisy_inject": MemberType.INLINED,
    },
)
register_opaque_type(
    Counter,
    typ="reference",
    guard_fn=lambda obj: [obj.start],
    members={"start": MemberType.USE_REAL, "get_start_tensor": MemberType.INLINED},
)
register_opaque_type(
    NestedCounters,
    typ="reference",
    members={
        "c": MemberType.USE_REAL,
        "get_c": MemberType.USE_REAL,
        "get_starts": MemberType.INLINED,
        "__getitem__": MemberType.INLINED,
    },
)
register_opaque_type(
    NestedQueue,
    typ="reference",
    members={
        "q": MemberType.USE_REAL,
        "get_q": MemberType.INLINED,
        "pop_q": MemberType.INLINED,
    },
)
register_opaque_type(AddModule, typ="reference")
register_opaque_type(ValueConfig, typ="value")
register_opaque_type(SizeStore, typ="value")
register_opaque_type(NestedValueSize, typ="value")
register_opaque_type(OpaqueMultiplier, typ="reference")
register_opaque_type(Color, typ="reference")
register_opaque_type(ColorWithDescriptor, typ="reference")


# A tensor subclass (similar to TwoTensor) that also holds an opaque Counter
# object
class TensorWithCounter(torch.Tensor):
    @staticmethod
    def __new__(cls, a, b, counter, size_store, outer_size=None, outer_stride=None):
        if outer_size is None:
            outer_size = a.size()
        if outer_stride is None:
            outer_stride = a.stride()

        if not (a.device == b.device and a.dtype == b.dtype):
            raise AssertionError(
                f"device/dtype mismatch: {a.device}/{a.dtype} vs {b.device}/{b.dtype}"
            )
        kwargs = {}
        kwargs["strides"] = outer_stride
        kwargs["storage_offset"] = a.storage_offset()
        kwargs["device"] = a.device
        kwargs["layout"] = a.layout
        kwargs["requires_grad"] = a.requires_grad
        kwargs["dtype"] = a.dtype
        out = torch.Tensor._make_wrapper_subclass(cls, outer_size, **kwargs)
        return out

    def __init__(self, a, b, counter, size_store, outer_size=None, outer_stride=None):
        self.a = a
        self.b = b
        self._counter = counter
        self._size_store = size_store

    def __repr__(self):
        return f"TensorWithCounter({self.a}, {self.b}, {self._counter}, {self._size_store})"

    def __tensor_flatten__(self):
        return ["a", "b", "_counter"], (self._size_store,)

    @staticmethod
    def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
        a, b = inner_tensors["a"], inner_tensors["b"]
        counter = inner_tensors["_counter"]
        (size_store,) = ctx
        return TensorWithCounter(a, b, counter, size_store, outer_size, outer_stride)

    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs):
        if kwargs is None:
            kwargs = {}

        def unwrap(x):
            return x.a if isinstance(x, TensorWithCounter) else x

        def wrap(x, counter, size_store):
            return (
                TensorWithCounter(x, x.clone(), counter, size_store)
                if isinstance(x, torch.Tensor)
                else x
            )

        # Get counter from first TensorWithCounter arg
        counter = None
        size_store = None
        for arg in torch.utils._pytree.tree_leaves(args):
            if isinstance(arg, TensorWithCounter):
                counter = arg._counter
                size_store = arg._size_store
                break

        unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
        unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)

        out = func(*unwrapped_args, **unwrapped_kwargs)
        return torch.utils._pytree.tree_map(lambda x: wrap(x, counter, size_store), out)

    @property
    def size_store(self):
        return self._size_store

    def get_size_store(self):
        return self._size_store

    def get_counter(self):
        return self._counter


class TestOpaqueObject(TestCase):
    def setUp(self):
        self.lib = torch.library.Library("_TestOpaqueObject", "FRAGMENT")  # noqa: TOR901
        self._opaque_types_before_test = set(_OPAQUE_TYPES_BY_NAME.keys())

        torch.library.define(
            "_TestOpaqueObject::queue_push",
            f"({get_opaque_type_name(OpaqueQueue)} a, Tensor b) -> ()",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::queue_push", "CompositeExplicitAutograd", lib=self.lib
        )
        def push_impl(queue: OpaqueQueue, b: torch.Tensor) -> None:
            if not isinstance(queue, OpaqueQueue):
                raise AssertionError(f"expected OpaqueQueue, got {type(queue)}")
            queue.push(b)

        @torch.library.register_fake("_TestOpaqueObject::queue_push", lib=self.lib)
        def push_impl_fake(q: OpaqueQueue, b: torch.Tensor) -> None:
            pass

        torch.library._register_effectful_op(
            "_TestOpaqueObject::queue_push", EffectType.ORDERED
        )

        self.lib.define(
            f"queue_pop({get_opaque_type_name(OpaqueQueue)} a) -> Tensor",
        )

        def pop_impl(queue: OpaqueQueue) -> torch.Tensor:
            if not isinstance(queue, OpaqueQueue):
                raise AssertionError(f"expected OpaqueQueue, got {type(queue)}")
            return queue.pop()

        self.lib.impl("queue_pop", pop_impl, "CompositeExplicitAutograd")

        def pop_impl_fake(q: OpaqueQueue) -> torch.Tensor:
            # This is not accurate since the queue could have tensors that are
            # not rank 1
            ctx = torch.library.get_ctx()
            u0 = ctx.new_dynamic_size()
            return torch.empty(u0)

        self.lib._register_fake("queue_pop", pop_impl_fake)

        torch.library._register_effectful_op(
            "_TestOpaqueObject::queue_pop", EffectType.ORDERED
        )

        @torch.library.custom_op(
            "_TestOpaqueObject::queue_size",
            mutates_args=[],
        )
        def size_impl(queue: OpaqueQueue) -> int:
            if not isinstance(queue, OpaqueQueue):
                raise AssertionError(f"expected OpaqueQueue, got {type(queue)}")
            return queue.size()

        @size_impl.register_fake
        def size_impl_fake(q: OpaqueQueue) -> int:
            ctx = torch._custom_op.impl.get_ctx()
            u0 = ctx.new_dynamic_size()
            torch._check(u0 >= 0)
            return u0

        torch.library.define(
            "_TestOpaqueObject::noisy_inject",
            f"(Tensor x, {get_opaque_type_name(RNGState)} obj) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::noisy_inject", "CompositeExplicitAutograd", lib=self.lib
        )
        def noisy_inject(x: torch.Tensor, rng_state: RNGState) -> torch.Tensor:
            if not isinstance(rng_state, RNGState):
                raise AssertionError(f"expected RNGState, got {type(rng_state)}")
            out = x.clone()
            for i in range(out.numel()):
                out.view(-1)[i] += rng_state.rng.random()
            return out

        @torch.library.register_fake("_TestOpaqueObject::noisy_inject", lib=self.lib)
        def noisy_inject_fake(x: torch.Tensor, obj: RNGState) -> torch.Tensor:
            if not isinstance(obj, RNGState):
                raise AssertionError(f"expected RNGState, got {type(obj)}")
            if obj.seed < 0:
                raise AssertionError(f"seed must be >= 0, got {obj.seed}")
            return torch.empty_like(x)

        @torch.library.custom_op(
            "_TestOpaqueObject::increment_counter",
            mutates_args=["prev"],
        )
        def increment_counter_impl(c: Counter, prev: torch.Tensor) -> torch.Tensor:
            if not isinstance(c, Counter):
                raise AssertionError(f"expected Counter, got {type(c)}")
            prev.copy_(c.counter)
            c.increment_counter()
            return c.counter

        @increment_counter_impl.register_fake
        def increment_counter_fake(c: Counter, prev: torch.Tensor) -> torch.Tensor:
            return torch.empty_like(prev)

        torch.library.define(
            "_TestOpaqueObject::process_with_config",
            f"(Tensor x, {get_opaque_type_name(ValueConfig)} config) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::process_with_config",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def process_with_config_impl(
            x: torch.Tensor, config: ValueConfig
        ) -> torch.Tensor:
            if not isinstance(config, ValueConfig):
                raise AssertionError(f"expected ValueConfig, got {type(config)}")
            if config.mode == "square":
                return x * x
            elif config.mode == "double":
                return x + x
            else:
                return x.clone()

        @torch.library.register_fake(
            "_TestOpaqueObject::process_with_config", lib=self.lib
        )
        def process_with_config_fake(
            x: torch.Tensor, config: ValueConfig
        ) -> torch.Tensor:
            return torch.empty_like(x)

        torch.library.define(
            "_TestOpaqueObject::process_nested_config",
            f"(Tensor x, {get_opaque_type_name(NestedValueSize)} config) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::process_nested_config",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def process_nested_config_impl(
            x: torch.Tensor, config: NestedValueSize
        ) -> torch.Tensor:
            if not isinstance(config, NestedValueSize):
                raise AssertionError(f"expected NestedValueSize, got {type(config)}")
            if config.config.mode == "square":
                return x * x
            elif config.config.mode == "double":
                return x + x
            else:
                return x.clone()

        @torch.library.register_fake(
            "_TestOpaqueObject::process_nested_config", lib=self.lib
        )
        def process_nested_config_fake(
            x: torch.Tensor, config: ValueConfig
        ) -> torch.Tensor:
            return torch.empty_like(x)

        torch.library.define(
            "_TestOpaqueObject::process_multiple_sizes",
            f"(Tensor x, {get_opaque_type_name(SizeStore)}[]? sizes) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::process_multiple_sizes",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def process_multiple_sizes_impl(
            x: torch.Tensor, config: list[SizeStore] | None
        ) -> torch.Tensor:
            if config is None:
                return x.clone()
            else:
                x_res = x.clone()
                for size in config:
                    x_res += size.size
                return x_res

        @torch.library.register_fake(
            "_TestOpaqueObject::process_multiple_sizes", lib=self.lib
        )
        def process_multiple_sizes_fake(
            x: torch.Tensor, config: list[SizeStore] | None
        ) -> torch.Tensor:
            return torch.empty_like(x)

        opaque_multiplier_type = get_opaque_type_name(OpaqueMultiplier)
        color_type = get_opaque_type_name(Color)

        torch.library.define(
            "_TestOpaqueObject::apply_color_scale",
            f"({color_type} color, Tensor x) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::apply_color_scale",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def apply_color_scale_impl(color: Color, x: torch.Tensor) -> torch.Tensor:
            return x * float(color)

        @torch.library.register_fake(
            "_TestOpaqueObject::apply_color_scale", lib=self.lib
        )
        def apply_color_scale_fake(color: Color, x: torch.Tensor) -> torch.Tensor:
            return torch.empty_like(x)

        torch.library.define(
            "_TestOpaqueObject::mul_with_scale",
            f"({opaque_multiplier_type} scale_obj, Tensor x) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        torch.library.define(
            "_TestOpaqueObject::get_multiplier_tensor",
            f"({opaque_multiplier_type} scale_obj, Tensor tensor) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::mul_with_scale",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def mul_with_scale_impl(
            scale_obj: OpaqueMultiplier, x: torch.Tensor
        ) -> torch.Tensor:
            return x * scale_obj.multiplier

        @torch.library.register_fake("_TestOpaqueObject::mul_with_scale", lib=self.lib)
        def mul_with_scale_fake(
            scale_obj: OpaqueMultiplier, x: torch.Tensor
        ) -> torch.Tensor:
            return torch.empty_like(x)

        @torch.library.impl(
            "_TestOpaqueObject::get_multiplier_tensor",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def get_multiplier_tensor_impl(
            scale_obj: OpaqueMultiplier, tensor: torch.Tensor
        ) -> torch.Tensor:
            return tensor * scale_obj.multiplier

        @torch.library.register_fake(
            "_TestOpaqueObject::get_multiplier_tensor", lib=self.lib
        )
        def get_multiplier_tensor_fake(
            scale_obj: OpaqueMultiplier, tensor: torch.Tensor
        ) -> torch.Tensor:
            return torch.empty_like(tensor)

        def mul_setup_context(ctx, inputs, output):
            ctx.scale_obj = inputs[0]

        def mul_backward(ctx, grad) -> tuple[torch.Tensor, None]:
            scale = torch.ops._TestOpaqueObject.get_multiplier_tensor(
                ctx.scale_obj, grad
            )
            return None, scale

        torch.library.register_autograd(
            "_TestOpaqueObject::mul_with_scale",
            mul_backward,
            setup_context=mul_setup_context,
            lib=self.lib,
        )

        rng_state_type = get_opaque_type_name(RNGState)
        # Define forward custom op that takes two opaque objects
        torch.library.define(
            "_TestOpaqueObject::multi_rng",
            f"(Tensor x, {rng_state_type} rng1, {rng_state_type} rng2) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::multi_rng",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def multi_rng_impl(
            x: torch.Tensor, rng1: RNGState, rng2: RNGState
        ) -> torch.Tensor:
            val1 = rng1.rng.random()
            val2 = rng2.rng.random()
            return x + val1 + val2

        @torch.library.register_fake("_TestOpaqueObject::multi_rng", lib=self.lib)
        def multi_rng_fake(
            x: torch.Tensor, rng1: RNGState, rng2: RNGState
        ) -> torch.Tensor:
            return torch.zeros_like(x)

        # Define backward custom op that also takes two opaque objects
        torch.library.define(
            "_TestOpaqueObject::multi_rng_backward",
            f"(Tensor grad_output, Tensor saved, {rng_state_type} rng1, {rng_state_type} rng2) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::multi_rng_backward",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def multi_rng_backward_impl(
            grad_output: torch.Tensor,
            saved: torch.Tensor,
            rng1: RNGState,
            rng2: RNGState,
        ) -> torch.Tensor:
            return grad_output * 1.5

        @torch.library.register_fake(
            "_TestOpaqueObject::multi_rng_backward", lib=self.lib
        )
        def multi_rng_backward_fake(
            grad_output: torch.Tensor,
            saved: torch.Tensor,
            rng1: RNGState,
            rng2: RNGState,
        ) -> torch.Tensor:
            return grad_output

        def setup_context(ctx, inputs, output):
            """Save opaque objects for backward - tests correct positioning."""
            x, rng1, rng2 = inputs
            ctx.save_for_backward(output)
            ctx.rng1 = rng1
            ctx.rng2 = rng2

        def backward(ctx, grad_output: torch.Tensor):
            """Backward pass that uses the opaque objects."""
            saved = ctx.saved_tensors[0]
            grad_input = torch.ops._TestOpaqueObject.multi_rng_backward(
                grad_output, saved, ctx.rng1, ctx.rng2
            )
            return grad_input, None, None

        torch.library.register_autograd(
            "_TestOpaqueObject::multi_rng",
            backward,
            setup_context=setup_context,
            lib=self.lib,
        )

        torch.library.define(
            "_TestOpaqueObject::create_multiplier",
            f"(Tensor scale) -> {opaque_multiplier_type}",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::create_multiplier",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def create_multiplier_impl(scale: torch.Tensor) -> OpaqueMultiplier:
            return OpaqueMultiplier(scale.item())

        @torch.library.register_fake(
            "_TestOpaqueObject::create_multiplier", lib=self.lib
        )
        def create_multiplier_fake(scale: torch.Tensor) -> OpaqueMultiplier:
            return OpaqueMultiplier(0.0)

        counter_type = get_opaque_type_name(Counter)
        torch.library.define(
            "_TestOpaqueObject::create_counter",
            f"(SymInt start, SymInt end) -> {counter_type}",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::create_counter",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def create_counter_impl(start: int, end: int) -> Counter:
            return Counter(start, end)

        @torch.library.register_fake("_TestOpaqueObject::create_counter", lib=self.lib)
        def create_counter_fake(start: int, end: int) -> Counter:
            return Counter(start, end)

        @torch.library.custom_op(
            "_TestOpaqueObject::create_counter2",
            mutates_args=[],
        )
        def create_counter2_impl(start: int, end: int) -> tuple[Counter, Counter]:
            return Counter(start, end), Counter(end, start)

        @create_counter2_impl.register_fake
        def create_counter2_fake(start: int, end: int) -> tuple[Counter, Counter]:
            return Counter(start, end), Counter(end, start)

        counter_type = get_opaque_type_name(Counter)
        torch.library.define(
            "_TestOpaqueObject::counter_start",
            f"({counter_type} a) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::counter_start",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def counter_start_impl(a: Counter) -> torch.Tensor:
            return torch.scalar_tensor(a.start, dtype=torch.int64)

        @torch.library.register_fake("_TestOpaqueObject::counter_start", lib=self.lib)
        def counter_start_fake(a: Counter) -> torch.Tensor:
            return torch.scalar_tensor(0, dtype=torch.int64)

        super().setUp()

    def tearDown(self):
        self.lib._destroy()

        # Clean up any opaque types registered during the test
        types_to_remove = (
            set(_OPAQUE_TYPES_BY_NAME.keys()) - self._opaque_types_before_test
        )
        for name in types_to_remove:
            torch._C._unregister_opaque_type(name)
            _OPAQUE_TYPES_BY_NAME.pop(name, None)

        super().tearDown()

    def test_ops(self):
        queue = OpaqueQueue([], torch.zeros(3))

        torch.ops._TestOpaqueObject.queue_push(queue, torch.ones(3) + 1)
        size = torch.ops._TestOpaqueObject.queue_size(queue)
        self.assertEqual(size, 1)
        popped = torch.ops._TestOpaqueObject.queue_pop(queue)
        self.assertEqual(popped, torch.ones(3) + 1)
        size = torch.ops._TestOpaqueObject.queue_size(queue)
        self.assertEqual(size, 0)

    def test_custom_op_reference_return(self):
        @torch.compile(backend="inductor", fullgraph=True)
        def f(x):
            c1 = torch.ops._TestOpaqueObject.create_counter(0, 10)
            c2, c3 = torch.ops._TestOpaqueObject.create_counter2(0, 10)
            return (
                c1,
                torch.ops._TestOpaqueObject.increment_counter(c1, x),
                c2,
                torch.ops._TestOpaqueObject.increment_counter(c2, x),
                c3,
                torch.ops._TestOpaqueObject.increment_counter(c3, x),
            )

        x = torch.scalar_tensor(0, dtype=torch.int64)
        c1, r1, c2, r2, c3, r3 = f(x)
        self.assertIsInstance(c1, Counter)
        self.assertIsInstance(c2, Counter)
        self.assertIsInstance(c3, Counter)
        self.assertEqual(c3.start, 11)
        self.assertEqual(c3.end, 0)
        self.assertEqual(x, torch.scalar_tensor(10, dtype=torch.int64))
        self.assertEqual(r1, torch.scalar_tensor(1, dtype=torch.int64))
        self.assertEqual(r2, torch.scalar_tensor(1, dtype=torch.int64))
        self.assertEqual(r3, torch.scalar_tensor(11, dtype=torch.int64))

    def test_custom_op_value_return(self):
        with self.assertRaisesRegex(ValueError, "Return has unsupported type"):

            @torch.library.custom_op(
                "_TestOpaqueObject::create_value_config",
                mutates_args=[],
            )
            def create_value_config(mode: str) -> ValueConfig:
                return ValueConfig(mode)

        with self.assertRaisesRegex(ValueError, "Return has unsupported type"):

            @torch.library.custom_op(
                "_TestOpaqueObject::create_value_config_tuple",
                mutates_args=[],
            )
            def create_value_config_tuple(mode: str) -> tuple[ValueConfig, int]:
                return ValueConfig(mode), 0

    def test_fake_script_object_isinstance_per_type(self):
        queue = OpaqueQueue([], torch.zeros(3))
        rng = RNGState(42)

        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
        with fake_mode:
            fake_queue = maybe_to_fake_obj(fake_mode, queue)
            fake_rng = maybe_to_fake_obj(fake_mode, rng)

        self.assertIsInstance(fake_queue, OpaqueQueue)
        self.assertIsInstance(fake_rng, RNGState)

        self.assertNotIsInstance(fake_queue, RNGState)
        self.assertNotIsInstance(fake_rng, OpaqueQueue)

        self.assertIsInstance(fake_queue, FakeScriptObject)
        self.assertIsInstance(fake_rng, FakeScriptObject)

    @parametrize("make_fx_tracing_mode", ["fake", "symbolic"])
    def test_make_fx(self, make_fx_tracing_mode):
        class M(torch.nn.Module):
            def forward(self, queue, x):
                torch.ops._TestOpaqueObject.queue_push(queue, x.tan())
                torch.ops._TestOpaqueObject.queue_push(queue, x.cos())
                torch.ops._TestOpaqueObject.queue_push(queue, x.sin())
                pop1 = torch.ops._TestOpaqueObject.queue_pop(queue)
                size1 = torch.ops._TestOpaqueObject.queue_size(queue)
                pop2 = torch.ops._TestOpaqueObject.queue_pop(queue)
                size2 = torch.ops._TestOpaqueObject.queue_size(queue)
                x_cos = pop1 + size1
                x_sin = pop2 - size2
                return x_sin + x_cos

        q1 = OpaqueQueue([], torch.empty(0).fill_(-1))
        q2 = OpaqueQueue([], torch.empty(0).fill_(-1))

        x = torch.ones(2, 3)
        gm = make_fx(M(), tracing_mode=make_fx_tracing_mode)(q1, x)
        self.assertTrue(torch.allclose(gm(q1, x), M()(q2, x)))
        self.assertEqual(q1._push_counter, 3)
        self.assertEqual(q1._pop_counter, 2)
        self.assertEqual(q1._size_counter, 2)
        self.assertEqual(q1.size(), 1)
        self.assertExpectedInline(
            gm.code.strip("\n"),
            """\
def forward(self, arg0_1, arg1_1):
    tan = torch.ops.aten.tan.default(arg1_1)
    queue_push = torch.ops._TestOpaqueObject.queue_push.default(arg0_1, tan);  tan = queue_push = None
    cos = torch.ops.aten.cos.default(arg1_1)
    queue_push_1 = torch.ops._TestOpaqueObject.queue_push.default(arg0_1, cos);  cos = queue_push_1 = None
    sin = torch.ops.aten.sin.default(arg1_1);  arg1_1 = None
    queue_push_2 = torch.ops._TestOpaqueObject.queue_push.default(arg0_1, sin);  sin = queue_push_2 = None
    queue_pop = torch.ops._TestOpaqueObject.queue_pop.default(arg0_1)
    queue_size = torch.ops._TestOpaqueObject.queue_size.default(arg0_1)
    queue_pop_1 = torch.ops._TestOpaqueObject.queue_pop.default(arg0_1)
    queue_size_1 = torch.ops._TestOpaqueObject.queue_size.default(arg0_1);  arg0_1 = None
    add = torch.ops.aten.add.Tensor(queue_pop, queue_size);  queue_pop = queue_size = None
    sub = torch.ops.aten.sub.Tensor(queue_pop_1, queue_size_1);  queue_pop_1 = queue_size_1 = None
    add_1 = torch.ops.aten.add.Tensor(sub, add);  sub = add = None
    return add_1
    """,
        )

    @parametrize("make_fx_tracing_mode", ["fake", "symbolic"])
    def test_make_fx_value_type(self, make_fx_tracing_mode):
        def f(x, cfg):
            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)

        x = torch.randn(3, 3)
        cfg = ValueConfig("square")
        gm = make_fx(f, tracing_mode=make_fx_tracing_mode)(x, cfg)
        self.assertEqual(gm(x, cfg), f(x, cfg))

        self.assertExpectedInline(
            gm.code.strip("\n"),
            """\
def forward(self, x_1, cfg_1):
    process_with_config = torch.ops._TestOpaqueObject.process_with_config.default(x_1, ValueConfig(mode='square'));  x_1 = None
    return process_with_config
    """,  # noqa: B950
        )

    def test_subclass_opaque_output_reuses_input_proxy(self):
        # Regression test: when a tensor subclass's __torch_dispatch__ wraps
        # the output with the real OpaqueBase (not the FakeScriptObject proxy),
        # the AOTAutograd forward graph should still reference the opaque via
        # its input placeholder — not create a duplicate get_attr constant.
        #
        # This mirrors DTensor where C++ dispatch creates new DTensors storing
        # the real DeviceMesh, so output flattening calls maybe_to_fake_obj
        # and mints a fresh FakeScriptObject for the same underlying object.

        class TensorWithRealCounter(torch.Tensor):
            @staticmethod
            def __new__(cls, data, counter):
                return torch.Tensor._make_wrapper_subclass(
                    cls,
                    data.size(),
                    strides=data.stride(),
                    storage_offset=data.storage_offset(),
                    device=data.device,
                    dtype=data.dtype,
                    layout=data.layout,
                    requires_grad=data.requires_grad,
                )

            def __init__(self, data, counter):
                self._data = data
                self._counter = counter

            def __repr__(self):
                return "TensorWithRealCounter(...)"

            def __tensor_flatten__(self):
                return ["_data", "_counter"], ()

            @staticmethod
            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
                return TensorWithRealCounter(
                    inner_tensors["_data"], inner_tensors["_counter"]
                )

            @classmethod
            def __torch_dispatch__(cls, func, types, args, kwargs):
                if kwargs is None:
                    kwargs = {}

                counter = None
                for arg in torch.utils._pytree.tree_leaves(args):
                    if isinstance(arg, TensorWithRealCounter):
                        counter = arg._counter
                        break

                def unwrap(x):
                    return x._data if isinstance(x, TensorWithRealCounter) else x

                out = func(
                    *torch.utils._pytree.tree_map(unwrap, args),
                    **torch.utils._pytree.tree_map(unwrap, kwargs),
                )

                # Unwrap FakeScriptObject to real OpaqueBase, simulating what
                # happens in DTensor's C++ dispatch path.
                real_counter = counter
                if isinstance(counter, FakeScriptObject):
                    real_counter = counter.real_obj

                return torch.utils._pytree.tree_map(
                    lambda x: TensorWithRealCounter(x, real_counter)
                    if isinstance(x, torch.Tensor)
                    else x,
                    out,
                )

        counter = Counter(start=0, end=10)
        x = TensorWithRealCounter(torch.randn(4), counter)

        backend = AotEagerAndRecordGraphs()
        torch.compile(lambda x: x + 1, fullgraph=True, backend=backend)(x)

        fw = backend.fw_graphs[0]
        get_attr_nodes = [n for n in fw.graph.nodes if n.op == "get_attr"]
        self.assertEqual(
            get_attr_nodes,
            [],
            "Opaque output should reuse the input placeholder, not create a get_attr constant",
        )

    def test_guard_pickle_subclass_with_opaque_inner_attr(self):
        # Regression test: the guard state pickler serializes tensor subclasses
        # by iterating over __tensor_flatten__ inner attrs.  Opaque inner attrs
        # (e.g. Counter) must be handled correctly — they are pickled by the
        # normal pickle machinery alongside tensor inner attrs.
        from torch._dynamo.guards import GuardsStatePickler

        a = torch.randn(4)
        b = torch.randn(4)
        counter = Counter(start=0, end=10)
        size = SizeStore(4)
        x = TensorWithCounter(a, b, counter, size)

        import io

        buf = io.BytesIO()
        pickler = GuardsStatePickler({id(x): x}, {}, {}, buf)
        func, args = pickler.reducer_override(x)
        obj = func(*args)
        self.assertIsInstance(obj, torch.Tensor)

    @parametrize("make_fx_tracing_mode", ["fake", "symbolic"])
    def test_bad_fake(self, make_fx_tracing_mode):
        torch.library.define(
            "_TestOpaqueObject::bad_fake",
            f"(Tensor x, {get_opaque_type_name(RNGState)} obj) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        def f(q, x):
            torch.ops._TestOpaqueObject.bad_fake(x, q)
            return x.cos()

        def bad_fake1(x, rng_state) -> torch.Tensor:
            self.assertTrue(isinstance(rng_state, FakeScriptObject))
            out = x.clone()
            for i in range(out.numel()):
                out.view(-1)[i] += rng_state.rng.random()  # bad: accessing attributes
            return out

        torch.library.register_fake(
            "_TestOpaqueObject::bad_fake", bad_fake1, lib=self.lib, allow_override=True
        )

        with self.assertRaisesRegex(
            AttributeError,
            "Tried to call __getattr__ with attr",
        ):
            make_fx(f, tracing_mode=make_fx_tracing_mode)(RNGState(0), torch.ones(3))

        def bad_fake2(x, rng_state) -> torch.Tensor:
            rng_state.rng = "foo"
            return torch.empty_like(x)

        torch.library.register_fake(
            "_TestOpaqueObject::bad_fake", bad_fake2, lib=self.lib, allow_override=True
        )

        with self.assertRaisesRegex(
            AttributeError,
            "Tried to call __setattr__ with attr",
        ):
            make_fx(f, tracing_mode=make_fx_tracing_mode)(RNGState(0), torch.ones(3))

    def test_aot_export(self):
        class Model(torch.nn.Module):
            def __init__(self) -> None:
                super().__init__()

            def forward(self, rng_state, x):
                x = torch.ops._TestOpaqueObject.noisy_inject(x, rng_state)
                x = x * x
                x = torch.ops._TestOpaqueObject.noisy_inject(x, rng_state)
                x = x + x
                return (x,)

        mod = Model()
        rng = RNGState(0)
        x = torch.ones(2, 3)

        fake_mode = torch._subclasses.fake_tensor.FakeTensorMode()
        fake_rng = torch._library.fake_class_registry.maybe_to_fake_obj(fake_mode, rng)
        fake_x = fake_mode.from_tensor(x)
        gm = aot_export_module(mod, (fake_rng, fake_x), trace_joint=False)[0]

        # By default we don't register ops containing PyObjs as being effectful
        self.assertExpectedInline(
            gm.code.strip(),
            """\
def forward(self, arg0_1, arg1_1):
    noisy_inject = torch.ops._TestOpaqueObject.noisy_inject.default(arg1_1, arg0_1);  arg1_1 = None
    mul = torch.ops.aten.mul.Tensor(noisy_inject, noisy_inject);  noisy_inject = None
    noisy_inject_1 = torch.ops._TestOpaqueObject.noisy_inject.default(mul, arg0_1);  mul = arg0_1 = None
    add = torch.ops.aten.add.Tensor(noisy_inject_1, noisy_inject_1);  noisy_inject_1 = None
    return (add,)""",  # noqa: B950
        )

        torch.library._register_effectful_op(
            "_TestOpaqueObject::noisy_inject", EffectType.ORDERED
        )
        try:
            gm = aot_export_module(mod, (rng, fake_x), trace_joint=False)[0]
            # inputs: token, rng, x
            # return: token, res
            self.assertExpectedInline(
                gm.code.strip(),
                """\
def forward(self, arg0_1, arg1_1, arg2_1):
    with_effects = torch.ops.higher_order.with_effects(arg0_1, torch.ops._TestOpaqueObject.noisy_inject.default, arg2_1, arg1_1);  arg0_1 = arg2_1 = None
    getitem = with_effects[0]
    getitem_1 = with_effects[1];  with_effects = None
    mul = torch.ops.aten.mul.Tensor(getitem_1, getitem_1);  getitem_1 = None
    with_effects_1 = torch.ops.higher_order.with_effects(getitem, torch.ops._TestOpaqueObject.noisy_inject.default, mul, arg1_1);  getitem = mul = arg1_1 = None
    getitem_2 = with_effects_1[0]
    getitem_3 = with_effects_1[1];  with_effects_1 = None
    add = torch.ops.aten.add.Tensor(getitem_3, getitem_3);  getitem_3 = None
    return (getitem_2, add)""",  # noqa: B950
            )
        finally:
            torch.library._register_effectful_op(
                "_TestOpaqueObject::noisy_inject", None
            )

    def test_install_free_opaque_object(self):
        rng = RNGState(0)
        x = torch.ones(2, 3)

        def fn(x):
            return torch.ops._TestOpaqueObject.noisy_inject(x, rng)

        gm = _dynamo_graph_capture_for_export(fn)(x)
        self.assertExpectedInline(
            gm.code.strip(),
            """\
def forward(self, x):
    arg_0, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)
    l_flat_args_0_ = arg_0
    l__self____export_root___closure___0_cell_contents = self.L__self____export_root___closure___0_cell_contents
    res = torch.ops._TestOpaqueObject.noisy_inject(l_flat_args_0_, l__self____export_root___closure___0_cell_contents);  l_flat_args_0_ = l__self____export_root___closure___0_cell_contents = None
    return pytree.tree_unflatten((res,), self._out_spec)""",  # noqa: B950
        )

    def test_compile1(self):
        def foo(rng_state, x):
            if not isinstance(rng_state, RNGState):
                raise AssertionError(f"expected RNGState, got {type(rng_state)}")
            x = torch.ops._TestOpaqueObject.noisy_inject(x, rng_state)
            x = x * x
            x = torch.ops._TestOpaqueObject.noisy_inject(x, rng_state)
            x = x + x
            return x

        rng = RNGState(0)
        x = torch.ones(2, 3)

        res = torch.compile(foo, fullgraph=True, backend="inductor")(rng, x)
        self.assertFalse(torch.allclose(res, x * x + x))

        backend = AotEagerAndRecordGraphs()
        torch.compile(foo, fullgraph=True, backend=backend)(rng, x)

        # This is done in torch.fx's graph in _namespace.create_name() where it
        # sanitizes the name
        fx_class = _illegal_char_regex.sub("_", get_opaque_type_name(RNGState))
        self.assertExpectedInline(
            backend.graphs[0].code.strip(),
            f"""\
def forward(self, L_rng_state_ : {fx_class}, L_x_ : torch.Tensor):
    l_rng_state_ = L_rng_state_
    l_x_ = L_x_
    x = torch.ops._TestOpaqueObject.noisy_inject(l_x_, l_rng_state_);  l_x_ = None
    x_1 = x * x;  x = None
    x_2 = torch.ops._TestOpaqueObject.noisy_inject(x_1, l_rng_state_);  x_1 = l_rng_state_ = None
    x_3 = x_2 + x_2;  x_2 = None
    return (x_3,)""",  # noqa: B950
        )
        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1, arg1_1):
    noisy_inject = torch.ops._TestOpaqueObject.noisy_inject.default(arg1_1, arg0_1);  arg1_1 = None
    mul = torch.ops.aten.mul.Tensor(noisy_inject, noisy_inject);  noisy_inject = None
    noisy_inject_1 = torch.ops._TestOpaqueObject.noisy_inject.default(mul, arg0_1);  mul = arg0_1 = None
    add = torch.ops.aten.add.Tensor(noisy_inject_1, noisy_inject_1);  noisy_inject_1 = None
    return (add,)""",  # noqa: B950
        )

    def test_compile_inline_methods(self):
        def foo(rng_state, x):
            seed1 = rng_state.get_seed()
            seed2 = rng_state.seed
            x = torch.ops._TestOpaqueObject.noisy_inject(x, rng_state)
            x = x * (seed1 + seed2 + 1)
            x = rng_state.noisy_inject(x)
            x = x + x
            return x

        rng = RNGState(0)
        x = torch.ones(2, 3)

        backend = AotEagerAndRecordGraphs()
        torch.compile(foo, fullgraph=True, backend=backend)(rng, x)

        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1, arg1_1):
    noisy_inject = torch.ops._TestOpaqueObject.noisy_inject.default(arg1_1, arg0_1);  arg1_1 = None
    mul = torch.ops.aten.mul.Tensor(noisy_inject, 1);  noisy_inject = None
    noisy_inject_1 = torch.ops._TestOpaqueObject.noisy_inject.default(mul, arg0_1);  mul = arg0_1 = None
    add = torch.ops.aten.add.Tensor(noisy_inject_1, noisy_inject_1);  noisy_inject_1 = None
    return (add,)""",  # noqa: B950
        )

        res = torch.compile(foo, fullgraph=True, backend="inductor")(rng, x)
        self.assertFalse(torch.allclose(res, x * x + x))

    def test_reference_type_recompile(self):
        cnt = CompileCounter()

        def foo(counter, x):
            z = torch.ops._TestOpaqueObject.increment_counter(counter, x)
            x = x * z
            return x

        x = torch.ones(2, 3)

        opt_f = torch.compile(foo, backend=cnt, fullgraph=True)
        opt_f(Counter(1, 5), x)
        self.assertEqual(cnt.frame_count, 1)

        opt_f(Counter(1, 6), x)  # we only guard on the first number
        self.assertEqual(cnt.frame_count, 1)

        opt_f(Counter(2, 5), x)  # recompile!
        self.assertEqual(cnt.frame_count, 2)

    def test_nested_reference_recompile(self):
        def foo(nested_counter, x):
            c1 = nested_counter.c
            return c1.start + x

        cnt = CompileCounter()
        x = torch.ones(2, 3)
        inp = (NestedCounters(Counter(1, 5)), x)
        opt_f = torch.compile(foo, backend=cnt, fullgraph=True)
        res = opt_f(*inp)
        self.assertEqual(res, foo(*inp))
        self.assertEqual(cnt.frame_count, 1)

        inp = (NestedCounters(Counter(1, 6)), x)
        res = opt_f(*inp)
        self.assertEqual(res, foo(*inp))
        self.assertEqual(cnt.frame_count, 1)  # we only guard on the first number

        inp = (NestedCounters(Counter(2, 5)), x)
        res = opt_f(*inp)
        self.assertEqual(res, foo(*inp))
        self.assertEqual(cnt.frame_count, 2)  # recompile!

    def test_nested_reference_list_trace(self):
        def foo(nested_counter, x):
            for c in nested_counter.c:
                x = torch.ops._TestOpaqueObject.increment_counter(c, x)
            for start in nested_counter.get_starts():
                x = x + start
            return x

        backend = AotEagerAndRecordGraphs()
        inp = (NestedCounters([Counter(1, 5), Counter(2, 5)]), torch.ones(2, 3))
        torch.compile(foo, backend=backend, fullgraph=True)(*inp)

        fx_class = _illegal_char_regex.sub("_", get_opaque_type_name(Counter))
        self.assertExpectedInline(
            backend.graphs[0].code.strip(),
            f"""\
def forward(self, L_x_ : torch.Tensor, object_getattribute_L_nested_counter_c_0_ : {fx_class}, object_getattribute_L_nested_counter_c_1_ : {fx_class}):
    l_x_ = L_x_
    object_getattribute_l_nested_counter_c_0_ = object_getattribute_L_nested_counter_c_0_
    object_getattribute_l_nested_counter_c_1_ = object_getattribute_L_nested_counter_c_1_
    x = torch.ops._TestOpaqueObject.increment_counter(object_getattribute_l_nested_counter_c_0_, l_x_);  object_getattribute_l_nested_counter_c_0_ = l_x_ = None
    x_1 = torch.ops._TestOpaqueObject.increment_counter(object_getattribute_l_nested_counter_c_1_, x);  object_getattribute_l_nested_counter_c_1_ = x = None
    x_2 = x_1 + 1;  x_1 = None
    x_3 = x_2 + 2;  x_2 = None
    return (x_3,)""",  # noqa: B950
        )

    def test_nested_reference_trace(self):
        def foo(nested_queue, x):
            q1 = nested_queue.q
            torch.ops._TestOpaqueObject.queue_push(q1, x.tan())
            q2 = nested_queue.get_q()
            torch.ops._TestOpaqueObject.queue_push(q2, x.cos())
            pop1 = nested_queue.pop_q()
            pop2 = nested_queue.pop_q()
            return pop1 + pop2

        inp = (
            NestedQueue(OpaqueQueue([], torch.empty(0).fill_(-1))),
            torch.randn(2, 3),
        )
        backend = AotEagerAndRecordGraphs()
        res = torch.compile(foo, fullgraph=True, backend=backend)(*inp)
        self.assertEqual(res, foo(*inp))

        fx_class = _illegal_char_regex.sub("_", get_opaque_type_name(OpaqueQueue))
        self.assertExpectedInline(
            backend.graphs[0].code.strip(),
            f"""\
def forward(self, L_x_ : torch.Tensor, object_getattribute_L_nested_queue_q_ : {fx_class}):
    l_x_ = L_x_
    object_getattribute_l_nested_queue_q_ = object_getattribute_L_nested_queue_q_
    tan = l_x_.tan()
    queue_push = torch.ops._TestOpaqueObject.queue_push(object_getattribute_l_nested_queue_q_, tan);  tan = queue_push = None
    cos = l_x_.cos();  l_x_ = None
    queue_push_1 = torch.ops._TestOpaqueObject.queue_push(object_getattribute_l_nested_queue_q_, cos);  cos = queue_push_1 = None
    pop1 = torch.ops._TestOpaqueObject.queue_pop(object_getattribute_l_nested_queue_q_)
    sym_size_int = torch.ops.aten.sym_size.int(pop1, 0)
    ge = sym_size_int >= 0
    _assert_scalar_default = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u0 >= 0 on node 'ge'");  ge = _assert_scalar_default = None
    pop2 = torch.ops._TestOpaqueObject.queue_pop(object_getattribute_l_nested_queue_q_);  object_getattribute_l_nested_queue_q_ = None
    sym_size_int_1 = torch.ops.aten.sym_size.int(pop2, 0)
    ge_1 = sym_size_int_1 >= 0
    _assert_scalar_default_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u1 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_default_1 = None
    eq = sym_size_int == sym_size_int_1;  sym_size_int = sym_size_int_1 = None
    _assert_scalar_default_2 = torch.ops.aten._assert_scalar.default(eq, "Runtime assertion failed for expression Eq(u0, u1) on node 'eq'");  eq = _assert_scalar_default_2 = None
    add = pop1 + pop2;  pop1 = pop2 = None
    return (add,)""",  # noqa: B950
        )

        # inputs: (token, nested_queue.q, x)
        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1, arg1_1, arg2_1):
    tan = torch.ops.aten.tan.default(arg1_1)
    with_effects = torch.ops.higher_order.with_effects(arg0_1, torch.ops._TestOpaqueObject.queue_push.default, arg2_1, tan);  arg0_1 = tan = None
    getitem = with_effects[0];  with_effects = None
    cos = torch.ops.aten.cos.default(arg1_1);  arg1_1 = None
    with_effects_1 = torch.ops.higher_order.with_effects(getitem, torch.ops._TestOpaqueObject.queue_push.default, arg2_1, cos);  getitem = cos = None
    getitem_2 = with_effects_1[0];  with_effects_1 = None
    with_effects_2 = torch.ops.higher_order.with_effects(getitem_2, torch.ops._TestOpaqueObject.queue_pop.default, arg2_1);  getitem_2 = None
    getitem_4 = with_effects_2[0]
    getitem_5 = with_effects_2[1];  with_effects_2 = None
    sym_size_int = torch.ops.aten.sym_size.int(getitem_5, 0)
    ge = sym_size_int >= 0
    _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u0 >= 0 on node 'ge'");  ge = _assert_scalar = None
    with_effects_3 = torch.ops.higher_order.with_effects(getitem_4, torch.ops._TestOpaqueObject.queue_pop.default, arg2_1);  getitem_4 = arg2_1 = None
    getitem_6 = with_effects_3[0]
    getitem_7 = with_effects_3[1];  with_effects_3 = None
    sym_size_int_1 = torch.ops.aten.sym_size.int(getitem_7, 0)
    ge_1 = sym_size_int_1 >= 0
    _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u1 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
    eq_2 = sym_size_int == sym_size_int_1;  sym_size_int = sym_size_int_1 = None
    _assert_scalar_2 = torch.ops.aten._assert_scalar.default(eq_2, "Runtime assertion failed for expression Eq(u0, u1) on node 'eq'");  eq_2 = _assert_scalar_2 = None
    add_4 = torch.ops.aten.add.Tensor(getitem_5, getitem_7);  getitem_5 = getitem_7 = None
    return (getitem_6, add_4)""",  # noqa: B950
        )

    def test_compile_global(self):
        counter = Counter(0, 10)

        def foo(x, y):
            z = torch.ops._TestOpaqueObject.increment_counter(counter, y)
            x = x * z
            z = torch.ops._TestOpaqueObject.increment_counter(counter, y)
            x = x + z
            return x, counter

        inp = (torch.tensor(1), torch.tensor(0))
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        res = opt_f(*inp)
        self.assertEqual(res[0], torch.tensor(3))
        self.assertEqual(res[1].counter, torch.tensor(2))

        res = opt_f(*inp)
        self.assertEqual(res[0], torch.tensor(7))
        self.assertEqual(res[1].counter, torch.tensor(4))

        # counter is automatically lifted as an input
        # Even though we returned counter in the eager code, it does not get
        # returned in the graph because dynamo does not detect that the object
        # is mutated.
        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1, arg1_1, arg2_1):
    auto_functionalized_v2 = torch.ops.higher_order.auto_functionalized_v2(torch.ops._TestOpaqueObject.increment_counter.default, c = arg1_1, _prev_base_index = 0, _all_bases = [arg0_1])
    getitem = auto_functionalized_v2[0]
    getitem_1 = auto_functionalized_v2[1];  auto_functionalized_v2 = None
    mul = torch.ops.aten.mul.Tensor(arg2_1, getitem);  arg2_1 = getitem = None
    auto_functionalized_v2_1 = torch.ops.higher_order.auto_functionalized_v2(torch.ops._TestOpaqueObject.increment_counter.default, c = arg1_1, _prev_base_index = 0, _all_bases = [getitem_1]);  arg1_1 = getitem_1 = None
    getitem_2 = auto_functionalized_v2_1[0]
    getitem_3 = auto_functionalized_v2_1[1];  auto_functionalized_v2_1 = None
    add = torch.ops.aten.add.Tensor(mul, getitem_2);  mul = getitem_2 = None
    copy_ = torch.ops.aten.copy_.default(arg0_1, getitem_3);  arg0_1 = getitem_3 = copy_ = None
    return (add,)""",  # noqa: B950
        )

    def test_compile_create_intermediate(self):
        dynamo_counters.clear()

        def foo(x, y):
            counter = Counter(0, 10)
            z = torch.ops._TestOpaqueObject.increment_counter(counter, y)
            x = x * z
            return x

        inp = (torch.tensor(1), torch.tensor(0))
        torch.compile(foo)(*inp)
        self.assertEqual(len(dynamo_counters["graph_break"]), 1)
        self.assertTrue(
            "An opaque object was created in the middle of the program"
            in next(iter(dynamo_counters["graph_break"].keys())),
        )

    def test_compile_attribute(self):
        counter = Counter(0, 10)

        def foo(counter, x):
            x = x * x
            counter.increment_counter()
            return x

        with self.assertRaisesRegex(
            RuntimeError, "Attempted to access unregistered member on an OpaqueObject"
        ):
            torch.compile(foo, backend="eager")(counter, torch.ones(2, 3))

        def bar(counter, x):
            x = x * x
            x += counter.counter
            return x

        with self.assertRaisesRegex(
            RuntimeError, "Attempted to access unregistered member on an OpaqueObject"
        ):
            torch.compile(bar, backend="eager")(counter, torch.ones(2, 3))

        def foo(counter, x):
            return counter.get_c()

        with self.assertRaisesRegex(
            RuntimeError,
            "Opaque object member with method-type USE_REAL returned a reference-type opaque object.",
        ):
            torch.compile(foo, backend="eager")(
                NestedCounters(Counter(1, 5)), torch.ones(2, 3)
            )

    def test_compile_fixed_stride_order(self):
        hs_name = get_opaque_type_name(HoistedString)

        torch.library.define(
            "_TestOpaqueObject::stride_op",
            f"(Tensor x, {hs_name} s) -> Tensor",
            tags=(torch.Tag.needs_fixed_stride_order,),
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::stride_op",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def stride_op_impl(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
            return x * 2.0

        @torch.library.register_fake("_TestOpaqueObject::stride_op", lib=self.lib)
        def stride_op_fake(x, s):
            return torch.empty_like(x)

        def fn(x, s):
            return torch.ops._TestOpaqueObject.stride_op(x, s)

        s = HoistedString("double")
        x = torch.randn(4, 4)

        compiled_fn = torch.compile(fn, backend="inductor")
        result = compiled_fn(x, s)

        expected = x * 2.0
        self.assertEqual(result, expected)

    def test_compile_exact_strides(self):
        hs_name = get_opaque_type_name(HoistedString)

        torch.library.define(
            "_TestOpaqueObject::exact_op",
            f"(Tensor x, {hs_name} s) -> Tensor",
            tags=(torch.Tag.needs_exact_strides,),
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::exact_op",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def exact_op_impl(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
            return x * 3.0

        @torch.library.register_fake("_TestOpaqueObject::exact_op", lib=self.lib)
        def exact_op_fake(x, s):
            return torch.empty_like(x)

        def fn(x, s):
            return torch.ops._TestOpaqueObject.exact_op(x, s)

        s = HoistedString("double")
        x = torch.randn(4, 4)

        compiled_fn = torch.compile(fn, backend="inductor")
        result = compiled_fn(x, s)

        expected = x * 3.0
        self.assertEqual(result, expected)

    def test_compile_contiguous_strides(self):
        hs_name = get_opaque_type_name(HoistedString)

        torch.library.define(
            "_TestOpaqueObject::contig_op",
            f"(Tensor x, {hs_name} s) -> Tensor",
            tags=(torch.Tag.needs_contiguous_strides,),
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::contig_op",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def contig_op_impl(x: torch.Tensor, s: HoistedString) -> torch.Tensor:
            return x * 4.0

        @torch.library.register_fake("_TestOpaqueObject::contig_op", lib=self.lib)
        def contig_op_fake(x, s):
            return torch.empty_like(x)

        def fn(x, s):
            return torch.ops._TestOpaqueObject.contig_op(x, s)

        s = HoistedString("double")
        x = torch.randn(4, 4)

        compiled_fn = torch.compile(fn, backend="inductor")
        result = compiled_fn(x, s)

        expected = x * 4.0
        self.assertEqual(result, expected)

    def test_export_joint(self):
        torch.library.define(
            "_TestOpaqueObject::module_mul",
            f"({get_opaque_type_name(AddModule)} a, Tensor b, SymInt c) -> Tensor",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::module_mul", "CompositeExplicitAutograd", lib=self.lib
        )
        def module_mul_impl(m: AddModule, a: torch.Tensor, b: int) -> torch.Tensor:
            if not isinstance(m, AddModule):
                raise AssertionError(f"expected AddModule, got {type(m)}")
            return m(a, b)

        @torch.library.register_fake("_TestOpaqueObject::module_mul", lib=self.lib)
        def module_mul_fake(m: AddModule, a: torch.Tensor, b: int) -> torch.Tensor:
            return torch.empty_like(a)

        def module_mul_setup_context(ctx, inputs, output):
            m, a, b = inputs
            ctx.b = b

        def module_mul_backward(ctx, grad) -> torch.Tensor:
            return None, grad * ctx.b, None

        torch.library.register_autograd(
            "_TestOpaqueObject::module_mul",
            module_mul_backward,
            setup_context=module_mul_setup_context,
            lib=self.lib,
        )

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.moo = AddModule()

            def forward(self, x, y):
                b = y.item()
                return torch.ops._TestOpaqueObject.module_mul(self.moo, x, b)

        inp = (torch.randn(3, requires_grad=True), torch.tensor(4))
        with ExitStack() as stack:
            with FakeTensorMode(shape_env=ShapeEnv()):
                joint = aot_export_joint_with_descriptors(stack, M(), inp)
                self.assertExpectedInline(
                    joint.graph_module.code.strip(),
                    """\
def forward(self, primals, tangents):
    primals_1, primals_2, tangents_1, = fx_pytree.tree_flatten_spec([primals, tangents], self._in_spec)
    _local_scalar_dense = torch.ops.aten._local_scalar_dense.default(primals_2);  primals_2 = None
    _opaque_obj0 = self._opaque_obj0
    module_mul = torch.ops._TestOpaqueObject.module_mul.default(_opaque_obj0, primals_1, _local_scalar_dense);  _opaque_obj0 = primals_1 = None
    mul_1 = torch.ops.aten.mul.Tensor(tangents_1, _local_scalar_dense);  tangents_1 = _local_scalar_dense = None
    return pytree.tree_unflatten([module_mul, mul_1, None], self._out_spec)""",  # noqa: B950
                )
                compiled_fn = aot_compile_joint_with_descriptors(joint)

        self.assertEqual(compiled_fn(*inp), M()(*inp))

    def test_invalid_reference_type(self):
        # Test that classes without subclassing OpaqueBase are rejected
        class NoOpaqueBase:
            def __init__(self, x):
                self.x = x

        with self.assertRaisesRegex(
            TypeError,
            "must subclass torch._opaque_base.OpaqueBase",
        ):
            register_opaque_type(NoOpaqueBase, typ="reference")

        class BadMember(OpaqueBase):
            def __init__(self, x):
                self.x = x

        def foo(bad, y):
            return y + bad.x

        register_opaque_type(
            BadMember, typ="reference", members={"y": MemberType.USE_REAL}
        )
        with self.assertRaisesRegex(
            torch._dynamo.exc.InternalTorchDynamoError,
            f"Opaque object of type '{get_opaque_type_name(BadMember)}' was specified to have member 'y'",
        ):
            torch.compile(foo)(BadMember(1), torch.ones(1))

    def test_invalid_value_type(self):
        class NoEq(OpaqueBase):
            def __init__(self, x):
                self.x = x

        with self.assertRaisesRegex(
            TypeError, "expected to have a non-default `__eq__`"
        ):
            register_opaque_type(NoEq, typ="value")

        class NoHash(OpaqueBase):
            def __init__(self, x):
                self.x = x

            def __eq__(self, other):
                return self.x == other.x

        with self.assertRaisesRegex(
            TypeError, "expected to have a non-default `__hash__`"
        ):
            register_opaque_type(NoHash, typ="value")

        class NoRepr(OpaqueBase):
            def __init__(self, x):
                self.x = x

            def __eq__(self, other):
                return self.x == other.x

            def __hash__(self):
                return hash(self.x)

        with self.assertRaisesRegex(TypeError, "expected to have a `__fx_repr__`"):
            register_opaque_type(NoRepr, typ="value")

        class SpecifyMember(OpaqueBase):
            def __init__(self, x):
                self.x = x

            def __eq__(self, other):
                return self.x == other.x

            def __hash__(self):
                return hash(self.x)

            def __fx_repr__(self):
                return f"SpecifyMember({self.x})"

        with self.assertRaisesRegex(TypeError, "No need to specify `guard_fn`"):
            register_opaque_type(SpecifyMember, typ="value", guard_fn=lambda obj: [])

    def test_invalid_schema(self):
        with self.assertRaisesRegex(
            RuntimeError,
            "unknown type specifier",
        ):
            torch.library.define(
                "_TestOpaqueObject::invalid_op1",
                "(foo.bar.baz a) -> Tensor",
                tags=torch.Tag.pt2_compliant_tag,
                lib=self.lib,
            )

        with self.assertRaisesRegex(
            RuntimeError,
            r"expected \) but found 'dots' here",
        ):
            torch.library.define(
                "_TestOpaqueObject::invalid_op2",
                "(......... a) -> Tensor",
                tags=torch.Tag.pt2_compliant_tag,
                lib=self.lib,
            )

        with self.assertRaisesRegex(
            RuntimeError,
            "unknown type specifier",
        ):
            torch.library.define(
                "_TestOpaqueObject::invalid_op5",
                "(MyNamespace..MyClass a) -> Tensor",
                tags=torch.Tag.pt2_compliant_tag,
                lib=self.lib,
            )

    def test_invalid_opaque_obj_types(self):
        for t in [str, bool, int, float, torch.Tensor]:
            with self.assertRaisesRegex(ValueError, "Unable to register built-in type"):
                register_opaque_type(t, typ="reference")

        @dataclass
        class Bad1(OpaqueBase):
            x: int

        pytree.register_dataclass(Bad1)
        try:
            with self.assertRaisesRegex(
                ValueError,
                "cannot be registered as an opaque object as it has been registered as a pytree.",
            ):
                register_opaque_type(Bad1, typ="reference")
        finally:
            # Clean up pytree registration to avoid leaking state to other tests
            pytree.SUPPORTED_NODES.pop(Bad1, None)
            pytree.SUPPORTED_SERIALIZED_TYPES.pop(Bad1, None)
            pytree.CONSTANT_NODES.discard(Bad1)

        @dataclass
        class Bad2(OpaqueBase):
            x: int

        register_opaque_type(Bad2, typ="reference")
        with self.assertRaisesRegex(
            ValueError,
            "cannot be registered as a pytree as it has been registered as an opaque object.",
        ):
            pytree.register_dataclass(Bad2)

    def test_value_type_recompile(self):
        cnt = CompileCounter()

        def foo(x, cfg):
            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)

        x = torch.randn(3, 3)

        opt_f = torch.compile(foo, backend=cnt, fullgraph=True)
        res = opt_f(x, ValueConfig("square"))
        self.assertEqual(res, x * x)
        self.assertEqual(cnt.frame_count, 1)

        res = opt_f(x, ValueConfig("square"))
        self.assertEqual(res, x * x)
        self.assertEqual(cnt.frame_count, 1)

        # Recompilation!
        res = opt_f(x, ValueConfig("double"))
        self.assertEqual(res, x + x)
        self.assertEqual(cnt.frame_count, 2)

    @parametrize("backend", ["eager", "aot_eager", "inductor"])
    def test_value_type_graph_output(self, backend):
        def foo(x):
            return x * x, ValueConfig("square")

        x = torch.randn(3, 3)
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        res = opt_f(x)
        self.assertEqual(res[1], ValueConfig("square"))

        gm = _dynamo_graph_capture_for_export(foo)(x)
        res = gm(x)
        self.assertEqual(res[1], ValueConfig("square"))

    def test_value_type_graph_input(self):
        # Even though cfg is an input, it should not be an input to the dynamo
        # graph. Instead it should directly put in the graph argument as a
        # constant.

        def foo(x, cfg):
            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)

        x = torch.randn(3, 3)
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        opt_f(x, ValueConfig("square"))

        self.assertExpectedInline(
            backend.graphs[0].code.strip(),
            """\
def forward(self, L_x_ : torch.Tensor):
    l_x_ = L_x_
    process_with_config = torch.ops._TestOpaqueObject.process_with_config(l_x_, ValueConfig(mode='square'));  l_x_ = None
    return (process_with_config,)""",  # noqa: B950
        )
        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1):
    process_with_config = torch.ops._TestOpaqueObject.process_with_config.default(arg0_1, ValueConfig(mode='square'));  arg0_1 = None
    return (process_with_config,)""",  # noqa: B950
        )

        opt_f(x, ValueConfig("double"))

        self.assertExpectedInline(
            backend.fw_graphs[1].code.strip(),
            """\
def forward(self, arg0_1):
    process_with_config = torch.ops._TestOpaqueObject.process_with_config.default(arg0_1, ValueConfig(mode='double'));  arg0_1 = None
    return (process_with_config,)""",  # noqa: B950
        )

    def test_value_type_graph_intermediate(self):
        def foo(x, config):
            cfg = ValueConfig(config)
            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)

        x = torch.randn(3, 3)
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        opt_f(x, "square")

        self.assertExpectedInline(
            backend.graphs[0].code.strip(),
            """\
def forward(self, L_x_ : torch.Tensor):
    l_x_ = L_x_
    process_with_config = torch.ops._TestOpaqueObject.process_with_config(l_x_, ValueConfig(mode='square'));  l_x_ = None
    return (process_with_config,)""",  # noqa: B950
        )
        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1):
    process_with_config = torch.ops._TestOpaqueObject.process_with_config.default(arg0_1, ValueConfig(mode='square'));  arg0_1 = None
    return (process_with_config,)""",  # noqa: B950
        )

        opt_f(x, "double")
        self.assertExpectedInline(
            backend.fw_graphs[1].code.strip(),
            """\
def forward(self, arg0_1):
    process_with_config = torch.ops._TestOpaqueObject.process_with_config.default(arg0_1, ValueConfig(mode='double'));  arg0_1 = None
    return (process_with_config,)""",  # noqa: B950
        )

        opt_f = torch.compile(foo, fullgraph=True, backend="inductor")
        x = torch.randn(3, 3)
        self.assertEqual(opt_f(x, "square"), foo(x, "square"))
        self.assertEqual(opt_f(x, "double"), foo(x, "double"))

    def test_value_type_attr_access(self):
        def foo(x):
            size = SizeStore(x.shape[0])
            t1 = torch.ones((size.size,))
            t2 = torch.cat([x, t1])
            size.increment_size()
            return t2 + size.size

        x = torch.randn(3)
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        res = opt_f(x)
        self.assertEqual(res, foo(x))

        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1):
    ones = torch.ops.aten.ones.default([3], device = device(type='cpu'), pin_memory = False)
    cat = torch.ops.aten.cat.default([arg0_1, ones]);  arg0_1 = ones = None
    add = torch.ops.aten.add.Tensor(cat, 3);  cat = None
    return (add,)""",  # noqa: B950
        )

    def test_value_type_unregistered_method(self):
        # Unregistered methods on value types should inline (no error)
        def foo(x):
            cfg = ValueConfig("square")
            return x + len(cfg.mode)

        x = torch.randn(3)
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        res = opt_f(x)
        self.assertEqual(res, foo(x))

        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1):
    add = torch.ops.aten.add.Tensor(arg0_1, 6);  arg0_1 = None
    return (add,)""",
        )

    def test_weakref_cleanup(self):
        def register_tmp_class():
            class TmpClass(OpaqueBase):
                def __init__(self, value):
                    self.value = value

                def __eq__(self, other):
                    return self.value == other.value

                def __hash__(self):
                    return hash(self.value)

                def __fx_repr__(self):
                    return f"TmpClass(value={self.value!r})", {"TmpClass": TmpClass}

            register_opaque_type(TmpClass, typ="value")

            self.assertTrue(is_opaque_type(TmpClass))
            self.assertTrue(is_opaque_value_type(TmpClass))
            self.assertIn(TmpClass, _OPAQUE_TYPES)

            return get_opaque_type_name(TmpClass)

        # registers TmpClass as opaque
        tmp_class_name = register_tmp_class()

        # garbage collect TmpClass
        gc.collect()

        # Verify that the class is no longer registered
        for opaque_info in _OPAQUE_TYPES.values():
            self.assertFalse(tmp_class_name in opaque_info.class_name)

    def test_value_type_nested(self):
        def foo(x, config):
            size = SizeStore(x.shape[0])
            cfg = NestedValueSize(size, ValueConfig(config))
            return torch.ops._TestOpaqueObject.process_nested_config(x, cfg)

        x = torch.randn(3, 3)
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        opt_f(x, "square")

        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1):
    process_nested_config = torch.ops._TestOpaqueObject.process_nested_config.default(arg0_1, NestedValueSize(size=SizeStore(size=3), config=ValueConfig(mode='square')));  arg0_1 = None
    return (process_nested_config,)""",  # noqa: B950
        )

        opt_f = torch.compile(foo, fullgraph=True, backend="inductor")
        x = torch.randn(3, 3)
        self.assertEqual(opt_f(x, "square"), foo(x, "square"))
        self.assertEqual(opt_f(x, "double"), foo(x, "double"))

    def test_value_type_list(self):
        def foo(x):
            sizes = []
            for size in x.size():
                sizes.append(SizeStore(size))
            return torch.ops._TestOpaqueObject.process_multiple_sizes(x, sizes)

        x = torch.randn(3, 3)
        backend = AotEagerAndRecordGraphs()
        opt_f = torch.compile(foo, fullgraph=True, backend=backend)
        opt_f(x)

        self.assertExpectedInline(
            backend.fw_graphs[0].code.strip(),
            """\
def forward(self, arg0_1):
    process_multiple_sizes = torch.ops._TestOpaqueObject.process_multiple_sizes.default(arg0_1, [SizeStore(size=3), SizeStore(size=3)]);  arg0_1 = None
    return (process_multiple_sizes,)""",  # noqa: B950
        )

        opt_f = torch.compile(foo, fullgraph=True, backend="inductor")
        x = torch.randn(3, 3)
        self.assertEqual(opt_f(x), foo(x))

    def test_getitem(self):
        def foo(x):
            y = x * 2 + 1
            nested_counter = y.get_counter()
            s = y.size_store
            counter = nested_counter[s.size]
            return y * counter.get_start_tensor()

        x = TensorWithCounter(
            torch.tensor(1),
            torch.tensor(1),
            NestedCounters([Counter(1, 5), Counter(2, 5)]),
            SizeStore(1),
        )

        inp = (x,)
        backend = AotEagerAndRecordGraphs()
        res = torch.compile(foo, fullgraph=True, backend=backend)(*inp)
        self.assertEqual(res, torch.tensor(6))

        actual = normalize_gm(backend.graphs[0].print_readable(print_output=False))
        self.assertExpectedInline(
            actual,
            """\
class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "TensorWithCounter(i64[])"):
        l_x_ = L_x_

        mul: "TensorWithCounter(i64[])" = l_x_ * 2;  l_x_ = None
        y: "TensorWithCounter(i64[])" = mul + 1;  mul = None

        get_counter = y.get_counter()

        getitem = get_counter.__getitem__(1);  get_counter = None

        get_start_tensor: "i64[]" = getitem.get_start_tensor();  getitem = None
        mul_1: "TensorWithCounter(i64[])" = y * get_start_tensor;  y = get_start_tensor = None
        return (mul_1,)
""",  # noqa: B950
        )

    def test_tensor_subclass_with_opaque_attr(self):
        def fn(x):
            y = x * 2 + 1
            c1 = y._counter
            c2 = y.get_counter()
            s1 = y.size_store
            s2 = y.get_size_store()
            return y * c1.start + c2.start + s1.size + s2.size

        a = torch.rand(4, 4)
        b = torch.rand(4, 4)
        counter = Counter(start=3, end=10)
        size = SizeStore(4)
        x = TensorWithCounter(a, b, counter, size)

        backend = EagerAndRecordGraphs()
        cnt = CompileCounterWithBackend(backend)

        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
        opt_fn(x)

        actual = normalize_gm(backend.graphs[0].print_readable(print_output=False))
        self.assertExpectedInline(
            actual,
            """\
class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "TensorWithCounter(f32[4, 4])"):
        l_x_ = L_x_

        mul: "TensorWithCounter(f32[4, 4])" = l_x_ * 2;  l_x_ = None
        y: "TensorWithCounter(f32[4, 4])" = mul + 1;  mul = None

        getattr_1 = y._counter;  getattr_1 = None

        get_counter = y.get_counter();  get_counter = None

        get_size_store = y.get_size_store();  get_size_store = None

        mul_1: "TensorWithCounter(f32[4, 4])" = y * 3;  y = None
        add_1: "TensorWithCounter(f32[4, 4])" = mul_1 + 3;  mul_1 = None
        add_2: "TensorWithCounter(f32[4, 4])" = add_1 + 4;  add_1 = None
        add_3: "TensorWithCounter(f32[4, 4])" = add_2 + 4;  add_2 = None
        return (add_3,)
""",
        )
        self.assertEqual(cnt.frame_count, 1)

        a = torch.rand(4, 4)
        b = torch.rand(4, 4)
        counter = Counter(start=1, end=10)
        x = TensorWithCounter(a, b, counter, SizeStore(4))
        opt_fn(x)

        # Recompile since Counter has changed
        self.assertEqual(cnt.frame_count, 2)

        a = torch.rand(4, 4)
        b = torch.rand(4, 4)
        counter = Counter(start=1, end=10)
        x = TensorWithCounter(a, b, counter, SizeStore(5))
        opt_fn(x)

        # Recompile since SizeStore has changed
        self.assertEqual(cnt.frame_count, 3)

    def test_tensor_subclass_with_opaque_attr_backward(self):
        """Test opaque objects in tensor subclass are correctly remapped through backward."""

        def fn(x):
            return x * 2 + 1

        a = torch.rand(4, 4, requires_grad=True)
        b = torch.rand(4, 4, requires_grad=True)
        counter = Counter(start=3, end=10)
        size = SizeStore(4)
        x = TensorWithCounter(a, b, counter, size)

        cnt = CompileCounterWithBackend("aot_eager")
        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
        out = opt_fn(x)

        # Output should be TensorWithCounter with same opaque objects
        self.assertTrue(isinstance(out, TensorWithCounter))
        self.assertIs(out._counter, counter)
        self.assertIs(out._size_store, size)

        # Run backward
        out.sum().backward()
        self.assertIsNotNone(x.grad)
        self.assertTrue(isinstance(x.grad, TensorWithCounter))

        self.assertEqual(cnt.frame_count, 1)

        # Call again with different instances that are __eq__ equal.
        # Metadata guard uses ==, so equal values pass the guard.
        # But they're different objects, so assertIs tests the fix.
        a2 = torch.rand(4, 4, requires_grad=True)
        b2 = torch.rand(4, 4, requires_grad=True)
        counter2 = Counter(start=3, end=10)  # Same values, different instance
        size2 = SizeStore(4)  # Same value, different instance
        self.assertEqual(counter, counter2)  # Verify they're equal
        self.assertIsNot(counter, counter2)  # But not the same object
        x2 = TensorWithCounter(a2, b2, counter2, size2)

        out2 = opt_fn(x2)
        self.assertEqual(cnt.frame_count, 1)  # No recompilation

        self.assertTrue(isinstance(out2, TensorWithCounter))
        # Key checks: should use runtime instances, not traced instances
        self.assertIs(out2._counter, counter2)
        # Value-type opaques in ctx are baked at compile time
        self.assertEqual(out2._size_store, size2)

        out2.sum().backward()
        self.assertIsNotNone(x2.grad)
        self.assertTrue(isinstance(x2.grad, TensorWithCounter))
        self.assertIs(x2.grad._counter, counter2)
        self.assertEqual(x2.grad._size_store, size2)

    def test_tangent_primal_proxy_collision_for_opaque_inner_attr(self):
        """Regression test for tangent/primal proxy collision.

        When a tensor subclass has an opaque inner attr, joint graph tracing
        creates separate FakeScriptObject wrappers for the primal and tangent
        that share the same underlying real object.  set_proxy_slot must map
        the tangent wrapper to the *primal* proxy so that forward outputs
        don't spuriously depend on tangent placeholders (which would crash the
        partitioner with 'Node tangents_N was invalid, but is output').
        """
        from torch._library.fake_class_registry import FakeScriptObject
        from torch.fx.experimental.proxy_tensor import (
            _GraphAppendingTracerEx,
            set_proxy_slot,
        )

        counter = Counter(start=3, end=10)
        fso_primal = FakeScriptObject(counter, "Counter", counter)
        fso_tangent = FakeScriptObject(counter, "Counter", counter)
        # Sanity: different wrappers, same real_obj
        self.assertIsNot(fso_primal, fso_tangent)
        self.assertIs(
            object.__getattribute__(fso_primal, "real_obj"),
            object.__getattribute__(fso_tangent, "real_obj"),
        )

        graph = torch.fx.Graph()
        tracer = _GraphAppendingTracerEx(graph)

        primal_node = graph.placeholder("primals_1")
        tangent_node = graph.placeholder("tangents_1")
        primal_proxy = torch.fx.Proxy(primal_node, tracer)
        tangent_proxy = torch.fx.Proxy(tangent_node, tracer)

        # Register primal first, then tangent (mirrors joint graph tracing)
        set_proxy_slot(fso_primal, tracer, primal_proxy)
        set_proxy_slot(fso_tangent, tracer, tangent_proxy)

        # Both wrappers should resolve to the primal proxy
        self.assertIs(tracer.opaque_tracker[fso_primal].node, primal_node)
        self.assertIs(tracer.opaque_tracker[fso_tangent].node, primal_node)

    def test_opaque_produced_by_call_function_saved_for_backward(self):
        """Test that an opaque object produced by a call_function node
        (not a placeholder) is correctly saved for backward.

        Without is_opaque_node() in the min_cut partitioner, this crashes
        with 'AOT Autograd failed to partition' because the opaque node
        is a non-tensor call_function that can't be saved or recomputed."""

        def fn(x):
            scale = torch.tensor(2.5)
            multiplier = torch.ops._TestOpaqueObject.create_multiplier(scale)
            return torch.ops._TestOpaqueObject.mul_with_scale(multiplier, x)

        x = torch.randn(3, 3, requires_grad=True)
        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
        out = opt_fn(x)
        self.assertTrue(torch.allclose(out, x * 2.5))

        out.backward(torch.ones_like(out))
        self.assertIsNotNone(x.grad)
        expected_grad = torch.ones_like(x) * 2.5
        self.assertTrue(torch.allclose(x.grad, expected_grad))

    def test_tensor_subclass_opaque_backward_compiled_autograd(self):
        """Test opaque objects work with compiled autograd backward."""
        import torch._dynamo.compiled_autograd

        def fn(x):
            return x * 2

        a = torch.rand(4, 4, requires_grad=True)
        b = torch.rand(4, 4, requires_grad=True)
        counter = Counter(start=5, end=10)
        size = SizeStore(4)
        x = TensorWithCounter(a, b, counter, size)

        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
        out = opt_fn(x)

        self.assertTrue(isinstance(out, TensorWithCounter))
        self.assertIs(out._counter, counter)
        self.assertIs(out._size_store, size)

        with torch._dynamo.compiled_autograd._enable(torch.compile(backend="eager")):
            out.sum().backward()

        self.assertIsNotNone(x.grad)
        self.assertTrue(isinstance(x.grad, TensorWithCounter))

        # Second invocation with different-but-equal opaques
        a2 = torch.rand(4, 4, requires_grad=True)
        b2 = torch.rand(4, 4, requires_grad=True)
        counter2 = Counter(start=5, end=10)
        size2 = SizeStore(4)
        x2 = TensorWithCounter(a2, b2, counter2, size2)

        out2 = opt_fn(x2)
        self.assertTrue(isinstance(out2, TensorWithCounter))
        self.assertIs(out2._counter, counter2)
        self.assertEqual(out2._size_store, size2)

        with torch._dynamo.compiled_autograd._enable(torch.compile(backend="eager")):
            out2.sum().backward()

        self.assertIsNotNone(x2.grad)
        self.assertTrue(isinstance(x2.grad, TensorWithCounter))
        self.assertIs(x2.grad._counter, counter2)
        self.assertEqual(x2.grad._size_store, size2)

    def test_tensor_subclass_shared_opaque_remapping(self):
        """Test that shared opaques across inputs are correctly handled."""

        def fn(x, y):
            return x + y

        counter = Counter(start=3, end=10)  # Shared opaque
        size = SizeStore(4)

        a1 = torch.rand(4, 4, requires_grad=True)
        b1 = torch.rand(4, 4, requires_grad=True)
        x = TensorWithCounter(a1, b1, counter, size)

        a2 = torch.rand(4, 4, requires_grad=True)
        b2 = torch.rand(4, 4, requires_grad=True)
        y = TensorWithCounter(a2, b2, counter, size)  # Same counter

        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
        out = opt_fn(x, y)

        self.assertTrue(isinstance(out, TensorWithCounter))
        self.assertIs(out._counter, counter)

        # Second invocation with different-but-equal shared opaques
        counter2 = Counter(start=3, end=10)
        size2 = SizeStore(4)

        a3 = torch.rand(4, 4, requires_grad=True)
        b3 = torch.rand(4, 4, requires_grad=True)
        x2 = TensorWithCounter(a3, b3, counter2, size2)

        a4 = torch.rand(4, 4, requires_grad=True)
        b4 = torch.rand(4, 4, requires_grad=True)
        y2 = TensorWithCounter(a4, b4, counter2, size2)

        out2 = opt_fn(x2, y2)
        self.assertTrue(isinstance(out2, TensorWithCounter))
        self.assertIs(out2._counter, counter2)
        self.assertEqual(out2._size_store, size2)

    def test_shared_opaque_identity_guard(self):
        """When inputs share an opaque at trace time, breaking that sharing
        should trigger recompilation.

        Uses y + x so __torch_dispatch__ picks y's opaque (input 1), but
        first-occurrence-wins remapping points to x's position (input 0).
        With broken sharing, the output silently gets the wrong opaque.
        """

        def fn(x, y):
            return y + x

        counter = Counter(start=3, end=10)
        size = SizeStore(4)

        x = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter, size)
        y = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter, size)

        cnt = CompileCounterWithBackend("aot_eager")
        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
        out = opt_fn(x, y)
        self.assertEqual(cnt.frame_count, 1)
        # Shared opaques: output should have the shared counter
        self.assertIs(out._counter, counter)

        # Same sharing pattern → no recompile
        counter2 = Counter(start=3, end=10)
        size2 = SizeStore(4)
        x2 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter2, size2)
        y2 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter2, size2)
        out2 = opt_fn(x2, y2)
        self.assertEqual(cnt.frame_count, 1)
        self.assertIs(out2._counter, counter2)

        # Sharing broken → should recompile.
        # __torch_dispatch__ picks y's opaque, so the output should have
        # counter3b. Without the guard, the remapping uses x's position
        # and the output silently gets counter3a.
        counter3a = Counter(start=3, end=10)
        counter3b = Counter(start=3, end=10)
        size3 = SizeStore(4)
        x3 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter3a, size3)
        y3 = TensorWithCounter(torch.rand(4, 4), torch.rand(4, 4), counter3b, size3)
        out3 = opt_fn(x3, y3)
        self.assertIs(out3._counter, counter3b)
        self.assertEqual(cnt.frame_count, 2)

    def test_shared_direct_opaque_identity_guard(self):
        """When the same opaque is passed as multiple direct function
        arguments, breaking the sharing should trigger recompilation.

        Unlike test_shared_opaque_identity_guard where opaques are inside
        tensor subclass metadata, here opaques are direct function arguments.
        The tracer may conflate them when they're the same object, causing
        counter_start(counter_b) to return counter_a's start instead.
        """
        get_start = torch.ops._TestOpaqueObject.counter_start

        def fn(counter_a, counter_b, x):
            return get_start(counter_a) + x, get_start(counter_b) + x

        counter = Counter(start=3, end=10)
        x = torch.rand(4, 4)

        cnt = CompileCounterWithBackend("aot_eager")
        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
        out_a, out_b = opt_fn(counter, counter, x)
        self.assertEqual(cnt.frame_count, 1)
        self.assertEqual(out_a, 3 + x)
        self.assertEqual(out_b, 3 + x)

        # Same sharing → no recompile
        counter2 = Counter(start=3, end=10)
        out_a2, out_b2 = opt_fn(counter2, counter2, x)
        self.assertEqual(cnt.frame_count, 1)
        self.assertEqual(out_a2, 3 + x)
        self.assertEqual(out_b2, 3 + x)

        # Broken sharing → should recompile.
        # Both counters have start=3 so the existing guard_fn doesn't catch
        # the difference — only an identity guard would.
        counter3a = Counter(start=3, end=10)
        counter3b = Counter(start=3, end=10)
        out_a3, out_b3 = opt_fn(counter3a, counter3b, x)
        self.assertEqual(cnt.frame_count, 2)
        self.assertEqual(out_a3, 3 + x)
        self.assertEqual(out_b3, 3 + x)

    def test_tensor_subclass_shared_opaque_backward(self):
        """Test backward with multiple subclass inputs sharing the same opaque.

        This tests the case where grad_inputs share an opaque object.
        The bug was that create_subclass_meta(grad_inputs) would incorrectly
        build opaque_remappings between grad_inputs, causing an assertion
        failure in wrap_tensor_subclasses.
        """

        def fn(x, y):
            return x + y

        # Create two TensorWithCounter inputs sharing the SAME opaque objects
        counter = Counter(start=3, end=10)
        size = SizeStore(4)

        a1 = torch.rand(4, 4, requires_grad=True)
        b1 = torch.rand(4, 4, requires_grad=True)
        x = TensorWithCounter(a1, b1, counter, size)

        a2 = torch.rand(4, 4, requires_grad=True)
        b2 = torch.rand(4, 4, requires_grad=True)
        y = TensorWithCounter(a2, b2, counter, size)  # Same counter and size

        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)
        out = opt_fn(x, y)

        self.assertTrue(isinstance(out, TensorWithCounter))
        self.assertIs(out._counter, counter)

        out.sum().backward()
        self.assertIsNotNone(x.grad)
        self.assertIsNotNone(y.grad)

        # Second invocation with different-but-equal shared opaques
        counter2 = Counter(start=3, end=10)
        size2 = SizeStore(4)

        a3 = torch.rand(4, 4, requires_grad=True)
        b3 = torch.rand(4, 4, requires_grad=True)
        x2 = TensorWithCounter(a3, b3, counter2, size2)

        a4 = torch.rand(4, 4, requires_grad=True)
        b4 = torch.rand(4, 4, requires_grad=True)
        y2 = TensorWithCounter(a4, b4, counter2, size2)

        out2 = opt_fn(x2, y2)
        self.assertTrue(isinstance(out2, TensorWithCounter))
        self.assertIs(out2._counter, counter2)
        self.assertEqual(out2._size_store, size2)

        out2.sum().backward()
        self.assertIsNotNone(x2.grad)
        self.assertIsNotNone(y2.grad)
        self.assertTrue(isinstance(x2.grad, TensorWithCounter))
        self.assertIs(x2.grad._counter, counter2)
        self.assertEqual(x2.grad._size_store, size2)

    def test_deeply_nested_tensor_subclass_with_opaque(self):
        """Test opaque objects in nested tensor subclasses are correctly remapped."""

        # Outer subclass that wraps TensorWithCounter
        class NestedSubclass(torch.Tensor):
            @staticmethod
            def __new__(cls, inner, scale):
                # inner is a TensorWithCounter
                kwargs = {
                    "strides": inner.stride(),
                    "storage_offset": inner.storage_offset(),
                    "device": inner.device,
                    "layout": inner.layout,
                    "requires_grad": inner.requires_grad,
                    "dtype": inner.dtype,
                }
                return torch.Tensor._make_wrapper_subclass(cls, inner.size(), **kwargs)

            def __init__(self, inner, scale):
                self._inner = inner
                self._scale = scale

            def __repr__(self):
                return f"NestedSubclass({self._inner}, scale={self._scale})"

            def __tensor_flatten__(self):
                return ["_inner"], (self._scale,)

            @staticmethod
            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
                (scale,) = ctx
                return NestedSubclass(inner_tensors["_inner"], scale)

            @classmethod
            def __torch_dispatch__(cls, func, types, args, kwargs):
                if kwargs is None:
                    kwargs = {}

                def unwrap(x):
                    return x._inner if isinstance(x, NestedSubclass) else x

                def get_scale(x):
                    if isinstance(x, NestedSubclass):
                        return x._scale
                    return None

                scale = None
                for arg in torch.utils._pytree.tree_leaves(args):
                    scale = get_scale(arg)
                    if scale is not None:
                        break

                unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
                unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)
                out = func(*unwrapped_args, **unwrapped_kwargs)

                def wrap(x):
                    if isinstance(x, TensorWithCounter):
                        return NestedSubclass(x, scale)
                    return x

                return torch.utils._pytree.tree_map(wrap, out)

        def fn(x):
            return x * 2 + 1

        # Create nested structure: NestedSubclass -> TensorWithCounter -> tensors
        a = torch.rand(4, 4, requires_grad=True)
        b = torch.rand(4, 4, requires_grad=True)
        counter = Counter(start=3, end=10)
        size = SizeStore(4)
        inner = TensorWithCounter(a, b, counter, size)
        x = NestedSubclass(inner, scale=2.0)

        cnt = CompileCounterWithBackend("aot_eager")
        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
        out = opt_fn(x)

        # Output should preserve nested structure with same opaque objects
        self.assertTrue(isinstance(out, NestedSubclass))
        self.assertTrue(isinstance(out._inner, TensorWithCounter))
        self.assertIs(out._inner._counter, counter)
        self.assertIs(out._inner._size_store, size)
        self.assertEqual(out._scale, 2.0)

        # Test backward
        out.sum().backward()
        self.assertIsNotNone(x.grad)
        self.assertTrue(isinstance(x.grad, NestedSubclass))

        self.assertEqual(cnt.frame_count, 1)

        # Call again with different instances that are __eq__ equal.
        a2 = torch.rand(4, 4, requires_grad=True)
        b2 = torch.rand(4, 4, requires_grad=True)
        counter2 = Counter(start=3, end=10)  # Same values, different instance
        size2 = SizeStore(4)  # Same value, different instance
        inner2 = TensorWithCounter(a2, b2, counter2, size2)
        x2 = NestedSubclass(inner2, scale=2.0)

        out2 = opt_fn(x2)
        self.assertEqual(cnt.frame_count, 1)  # No recompilation

        self.assertTrue(isinstance(out2, NestedSubclass))
        # Key checks: should use runtime instances, not traced instances
        self.assertIs(out2._inner._counter, counter2)
        # Value-type opaques in ctx are baked at compile time, so identity
        # is not preserved across re-invocations.
        self.assertEqual(out2._inner._size_store, size2)
        self.assertEqual(out2._scale, 2.0)

        out2.sum().backward()
        self.assertIsNotNone(x2.grad)
        self.assertTrue(isinstance(x2.grad, NestedSubclass))
        self.assertTrue(isinstance(x2.grad._inner, TensorWithCounter))
        self.assertIs(x2.grad._inner._counter, counter2)
        self.assertEqual(x2.grad._inner._size_store, size2)

    def test_value_type_opaque_in_tensor_attrs_errors(self):
        """Value-type opaques must go in ctx, not tensor attrs."""

        class BadSubclass(torch.Tensor):
            @staticmethod
            def __new__(cls, data, size_store):
                kwargs = {
                    "strides": data.stride(),
                    "storage_offset": data.storage_offset(),
                    "device": data.device,
                    "layout": data.layout,
                    "requires_grad": data.requires_grad,
                    "dtype": data.dtype,
                }
                return torch.Tensor._make_wrapper_subclass(cls, data.size(), **kwargs)

            def __init__(self, data, size_store):
                self._data = data
                self._size_store = size_store

            def __tensor_flatten__(self):
                # BUG: SizeStore is value-type, should be in ctx not tensor attrs
                return ["_data", "_size_store"], ()

            @staticmethod
            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
                return BadSubclass(inner_tensors["_data"], inner_tensors["_size_store"])

            @classmethod
            def __torch_dispatch__(cls, func, types, args, kwargs):
                if kwargs is None:
                    kwargs = {}

                def unwrap(x):
                    return x._data if isinstance(x, BadSubclass) else x

                size_store = None
                for arg in torch.utils._pytree.tree_leaves(args):
                    if isinstance(arg, BadSubclass):
                        size_store = arg._size_store
                        break

                unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
                unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)
                out = func(*unwrapped_args, **unwrapped_kwargs)

                def wrap(x):
                    if isinstance(x, torch.Tensor):
                        return BadSubclass(x, size_store)
                    return x

                return torch.utils._pytree.tree_map(wrap, out)

        def fn(x):
            return x * 2

        x = BadSubclass(torch.rand(4, 4), SizeStore(4))
        opt_fn = torch.compile(fn, backend="aot_eager", fullgraph=True)

        with self.assertRaisesRegex(
            RuntimeError,
            "SizeStore.*tensor attrs.*Only tensors and reference-type opaques",
        ):
            opt_fn(x)

    def test_multi_reference_opaque_identity(self):
        """Test identity preservation with multiple reference-type opaques in tensor attrs."""

        class Tag(OpaqueBase):
            def __init__(self, name):
                self.name = name

            def __eq__(self, other):
                return isinstance(other, Tag) and self.name == other.name

            def __hash__(self):
                return hash(self.name)

        register_opaque_type(Tag, typ="reference")

        class TwoRefSubclass(torch.Tensor):
            @staticmethod
            def __new__(cls, data, counter, tag):
                kwargs = {
                    "strides": data.stride(),
                    "storage_offset": data.storage_offset(),
                    "device": data.device,
                    "layout": data.layout,
                    "requires_grad": data.requires_grad,
                    "dtype": data.dtype,
                }
                return torch.Tensor._make_wrapper_subclass(cls, data.size(), **kwargs)

            def __init__(self, data, counter, tag):
                self._data = data
                self._counter = counter
                self._tag = tag

            def __repr__(self):
                return f"TwoRefSubclass({self._data})"

            def __tensor_flatten__(self):
                return ["_data", "_counter", "_tag"], ()

            @staticmethod
            def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
                return TwoRefSubclass(
                    inner_tensors["_data"],
                    inner_tensors["_counter"],
                    inner_tensors["_tag"],
                )

            @classmethod
            def __torch_dispatch__(cls, func, types, args, kwargs):
                if kwargs is None:
                    kwargs = {}

                def unwrap(x):
                    return x._data if isinstance(x, TwoRefSubclass) else x

                counter = None
                tag = None
                for arg in torch.utils._pytree.tree_leaves(args):
                    if isinstance(arg, TwoRefSubclass):
                        counter = arg._counter
                        tag = arg._tag
                        break

                unwrapped_args = torch.utils._pytree.tree_map(unwrap, args)
                unwrapped_kwargs = torch.utils._pytree.tree_map(unwrap, kwargs)
                out = func(*unwrapped_args, **unwrapped_kwargs)

                def wrap(x):
                    if isinstance(x, torch.Tensor):
                        return TwoRefSubclass(x, counter, tag)
                    return x

                return torch.utils._pytree.tree_map(wrap, out)

        def fn(x):
            return x * 2 + 1

        counter = Counter(start=3, end=10)
        tag = Tag("test")
        x = TwoRefSubclass(torch.rand(4, 4, requires_grad=True), counter, tag)

        cnt = CompileCounterWithBackend("aot_eager")
        opt_fn = torch.compile(fn, backend=cnt, fullgraph=True)
        out = opt_fn(x)

        self.assertIsInstance(out, TwoRefSubclass)
        self.assertIs(out._counter, counter)
        self.assertIs(out._tag, tag)

        out.sum().backward()
        self.assertIsNotNone(x.grad)
        self.assertEqual(cnt.frame_count, 1)

        # Second invocation with different-but-equal opaques
        counter2 = Counter(start=3, end=10)
        tag2 = Tag("test")
        self.assertEqual(counter, counter2)
        self.assertIsNot(counter, counter2)
        x2 = TwoRefSubclass(torch.rand(4, 4, requires_grad=True), counter2, tag2)

        out2 = opt_fn(x2)
        self.assertEqual(cnt.frame_count, 1)  # No recompilation

        self.assertIsInstance(out2, TwoRefSubclass)
        self.assertIs(out2._counter, counter2)
        self.assertIs(out2._tag, tag2)

        out2.sum().backward()
        self.assertIsNotNone(x2.grad)
        self.assertIsInstance(x2.grad, TwoRefSubclass)
        self.assertIs(x2.grad._counter, counter2)
        self.assertIs(x2.grad._tag, tag2)

    def test_opaque_obj_saved_for_backward(self):
        """Test that opaque objects are correctly saved and passed to backward."""
        import torch._dynamo.compiled_autograd

        def foo(scale_obj, x):
            result = torch.ops._TestOpaqueObject.mul_with_scale(scale_obj, x)
            result = result * 2
            return result

        def compile_and_run_with_backend(backend):
            scale_obj = OpaqueMultiplier(2.5)
            x = torch.randn(3, 3, requires_grad=True)

            opt_f = torch.compile(foo, fullgraph=True, backend=backend)
            out = opt_f(scale_obj, x)
            expected = x * 2.5 * 2
            self.assertTrue(torch.allclose(out, expected))

            upstream_grad = torch.ones_like(out) * 5
            out.backward(upstream_grad)
            self.assertIsNotNone(x.grad)
            expected_grad = torch.ones_like(x) * 5 * 5
            self.assertTrue(torch.allclose(x.grad, expected_grad))

        backend = InductorAndRecordGraphs()
        compile_and_run_with_backend(backend)
        self.assertTrue(len(backend.graphs) > 0)
        fw_graph = backend.graphs[0]
        self.assertExpectedInline(
            fw_graph.code.strip(),
            f"""\
def forward(self, L_x_ : torch.Tensor, L_scale_obj_ : {_illegal_char_regex.sub("_", get_opaque_type_name(OpaqueMultiplier))}):
    l_x_ = L_x_
    l_scale_obj_ = L_scale_obj_
    result = torch.ops._TestOpaqueObject.mul_with_scale(l_scale_obj_, l_x_);  l_scale_obj_ = l_x_ = None
    result_1 = result * 2;  result = None
    return (result_1,)""",  # noqa: B950
        )

        backend = AotEagerAndRecordGraphs()
        compile_and_run_with_backend(backend)
        self.assertTrue(len(backend.fw_graphs) > 0)
        fw_graph = backend.fw_graphs[0]
        self.assertExpectedInline(
            fw_graph.code.strip(),
            """\
def forward(self, primals_1, primals_2):
    mul_with_scale = torch.ops._TestOpaqueObject.mul_with_scale.default(primals_2, primals_1);  primals_1 = None
    mul = torch.ops.aten.mul.Tensor(mul_with_scale, 2);  mul_with_scale = None
    return (mul, primals_2)""",
        )
        self.assertTrue(len(backend.bw_graphs) > 0)
        bw_graph = backend.bw_graphs[0]
        self.assertExpectedInline(
            bw_graph.code.strip(),
            """\
def forward(self, primals_2, tangents_1):
    mul_1 = torch.ops.aten.mul.Tensor(tangents_1, 2);  tangents_1 = None
    get_multiplier_tensor = torch.ops._TestOpaqueObject.get_multiplier_tensor.default(primals_2, mul_1);  primals_2 = mul_1 = None
    return (get_multiplier_tensor, None)""",
        )

        for use_compiled_autograd in [False, True]:
            with self.subTest(use_compiled_autograd=use_compiled_autograd):
                torch._dynamo.reset()

                def run_with_multiplier(multiplier_value: float):
                    scale_obj = OpaqueMultiplier(multiplier_value)
                    x = torch.randn(3, 3, requires_grad=True)

                    opt_f = torch.compile(foo, fullgraph=True, backend="aot_eager")
                    out = opt_f(scale_obj, x)
                    expected = x * multiplier_value * 2
                    self.assertTrue(torch.allclose(out, expected))

                    upstream_grad = torch.ones_like(out) * 5
                    with (
                        torch._dynamo.compiled_autograd._enable(
                            torch.compile(backend="eager")
                        )
                        if use_compiled_autograd
                        else contextlib.nullcontext()
                    ):
                        out.backward(upstream_grad)

                    self.assertIsNotNone(x.grad)
                    # gradient = upstream_grad * 2 * multiplier_value
                    expected_grad = torch.ones_like(x) * 5 * 2 * multiplier_value
                    self.assertTrue(torch.allclose(x.grad, expected_grad))
                    return x.grad.clone()

                grad1 = run_with_multiplier(2.5)
                torch._dynamo.reset()
                grad2 = run_with_multiplier(3.0)
                self.assertFalse(torch.allclose(grad1, grad2))

    def test_subgraph_tracer_create_arg_with_fake_script_object(self):
        """Test that opaque class attribute access works correctly.

        This tests the code path where:
        1. An opaque class (like Color) is accessed via OpaqueObjectClassVariable
        2. Attribute access (Color.RED) goes through var_getattr with static getattr
        3. The opaque object is correctly lifted as a graph input
        """
        from torch._library.opaque_object import is_opaque_reference_type

        self.assertTrue(is_opaque_reference_type(Color))
        self.assertTrue(is_opaque_reference_type(type(Color.RED)))

        captured = {"graph": None, "example_inputs": None}

        def capture_backend(gm, example_inputs):
            captured["graph"] = gm
            captured["example_inputs"] = example_inputs
            return gm

        @torch.compile(fullgraph=True, backend=capture_backend)
        def fn(x):
            return torch.ops._TestOpaqueObject.apply_color_scale(Color.GREEN, x)

        x = torch.randn(3, 3)
        result = fn(x)

        self.assertIsNotNone(captured["graph"])
        print(captured["graph"].code.strip())

        expected = x * float(Color.GREEN.value)
        self.assertTrue(torch.allclose(result, expected))

        graph_code = captured["graph"].code.strip()
        self.assertExpectedInline(
            graph_code,
            f"""\
def forward(self, L_x_ : torch.Tensor, G_Color_GREEN : {_illegal_char_regex.sub("_", get_opaque_type_name(Color))}):
    l_x_ = L_x_
    g_color_green = G_Color_GREEN
    apply_color_scale = torch.ops._TestOpaqueObject.apply_color_scale(g_color_green, l_x_);  g_color_green = l_x_ = None
    return (apply_color_scale,)""",  # noqa: B950
        )

    def test_hoist_basic(self):
        def f(x):
            return op_with_string(x, HoistedString("double"))

        x = torch.tensor(3.0)

        # backend = "eager"
        backend = EagerAndRecordGraphs()
        _ = torch.compile(f, fullgraph=True, backend=backend)(x)
        graph = backend.graphs[0]
        self.assertNotIn("double", str(graph))

        # backend = "aot_eager"
        backend = AotEagerAndRecordGraphs()
        _ = torch.compile(f, fullgraph=True, backend=backend)(x)
        graph = backend.graphs[0]
        self.assertNotIn("double", str(graph))

        # backend = "inductor"
        y = torch.compile(f, fullgraph=True)(x)
        self.assertEqual(y, 2 * x)

    @functorch_config.patch({"enable_autograd_cache": True})
    @inductor_config.patch(
        {
            "fx_graph_cache": True,
            "fx_graph_remote_cache": False,
        }
    )
    def test_hoist_cache_hits(self):
        torch._dynamo.reset()
        counters.clear()

        with fresh_inductor_cache():
            # Because HoistedString should not be in the graph, the following
            # two functions should share AOTAutogradCache and FXGraphCache entries

            @torch.compile(fullgraph=True)
            def f(x):
                return op_with_string(x, HoistedString("double"))

            @torch.compile(fullgraph=True)
            def g(x):
                return op_with_string(x, HoistedString("square"))

            x = torch.tensor(3.0)

            f(x)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
            self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
            self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)

            g(x)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_saved"], 1)
            self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 1)
            self.assertEqual(counters["inductor"]["fxgraph_cache_miss"], 1)

    def test_hoisted_value_type_make_fx(self):
        def foo(x, hoisted_str):
            return op_with_string(x, hoisted_str)

        x = torch.randn(3, 3)
        gm = make_fx(foo)(x, HoistedString("double"))

        self.assertExpectedInline(
            gm.code.strip(),
            """\
def forward(self, x_1, hoisted_str_1):
    op_with_string = torch.ops.mylib.op_with_string.default(x_1, hoisted_str_1);  x_1 = hoisted_str_1 = None
    return op_with_string""",  # noqa: B950
        )
        self.assertEqual(gm(x, HoistedString("double")), x * 2)
        self.assertEqual(gm(x, HoistedString("square")), x * x)

    def test_opaque_class_literal_attribute_inlined(self):
        """Test that literal attributes on opaque classes are inlined without source tracking.

        When accessing a literal class attribute (like Color.DEFAULT_SCALE = 1.5),
        the value should be constant-folded directly without creating a graph input.
        """
        captured = {"graph": None}

        def capture_backend(gm, example_inputs):
            captured["graph"] = gm
            return gm

        @torch.compile(fullgraph=True, backend=capture_backend)
        def fn(x):
            return x * Color.DEFAULT_SCALE

        x = torch.randn(3, 3)
        result = fn(x)

        expected = x * 1.5
        self.assertTrue(torch.allclose(result, expected))

        self.assertIsNotNone(captured["graph"])
        graph_code = captured["graph"].code.strip()
        self.assertExpectedInline(
            graph_code,
            """\
def forward(self, L_x_ : torch.Tensor):
    l_x_ = L_x_
    mul = l_x_ * 1.5;  l_x_ = None
    return (mul,)""",
        )

    def test_opaque_class_missing_attribute_graph_break(self):
        """Test that accessing a non-existent attribute on an opaque class causes a graph break.

        This tests GB7685: "Attribute not found on opaque class"
        """

        @torch.compile(fullgraph=True)
        def fn(x):
            return x * Color.NONEXISTENT

        x = torch.randn(3, 3)

        with self.assertRaisesRegex(
            torch._dynamo.exc.Unsupported,
            "Attribute not found on opaque class",
        ):
            fn(x)

    def test_opaque_class_unsupported_descriptor_graph_break(self):
        """Test that accessing an unsupported descriptor on an opaque class causes a graph break.

        This tests GB9567: "Unsupported descriptor on opaque class"
        """

        @torch.compile(fullgraph=True)
        def fn(x):
            return x * ColorWithDescriptor.custom_prop

        x = torch.randn(3, 3)

        with self.assertRaisesRegex(
            torch._dynamo.exc.Unsupported,
            "Unsupported descriptor on opaque class",
        ):
            fn(x)

    def test_opaque_class_staticmethod(self):
        """Test that accessing a staticmethod on an opaque class works correctly.

        This verifies that OpaqueObjectClassVariable.var_getattr properly handles
        staticmethod descriptors (instead of raising 'Unsupported descriptor').
        """
        captured = {"graph": None}

        def capture_backend(gm, _):
            captured["graph"] = gm
            return gm

        @torch.compile(fullgraph=True, backend=capture_backend)
        def fn(x):
            _ = Color.RED_STATIC
            return x * 2

        x = torch.randn(3, 3)
        result = fn(x)
        expected = x * 2
        self.assertTrue(torch.allclose(result, expected))
        self.assertIsNotNone(captured["graph"])

    def test_opaque_class_property(self):
        """Test that accessing a property descriptor on an opaque class works correctly.

        This verifies that OpaqueObjectClassVariable.var_getattr properly handles
        property descriptors. When accessing a property on the class (not instance),
        you get the property object back.
        """
        captured = {"graph": None}

        def capture_backend(gm, _):
            captured["graph"] = gm
            return gm

        @torch.compile(fullgraph=True, backend=capture_backend)
        def fn(x):
            _ = Color.name
            return x * 2

        x = torch.randn(3, 3)
        result = fn(x)
        expected = x * 2
        self.assertTrue(torch.allclose(result, expected))
        self.assertIsNotNone(captured["graph"])

    def test_multiple_opaque_objects_in_backward(self):
        """Test that multiple opaque objects are correctly positioned in backward pass.

        This test verifies that when setup_context saves multiple opaque objects,
        AOTAutograd correctly positions them when calling the backward function.
        This addresses a bug where opaque objects were placed in the wrong position.
        """
        # Test eager mode
        x = torch.randn(1, requires_grad=True)
        rng1 = RNGState(123)
        rng2 = RNGState(456)

        result = torch.ops._TestOpaqueObject.multi_rng(x, rng1, rng2)
        grad_o = torch.rand_like(result)
        result.backward(grad_o)
        self.assertEqual(x.grad, grad_o * 1.5)
        x.grad = None

        # Test compiled mode - this should not crash with incorrect opaque object positioning
        x2 = torch.randn(1, requires_grad=True)
        compiled_fn = torch.compile(
            torch.ops._TestOpaqueObject.multi_rng,
            fullgraph=True,
        )
        result = compiled_fn(x2, rng1, rng2)
        grad_o = torch.rand_like(result)
        result.backward(grad_o)
        self.assertEqual(x2.grad, grad_o * 1.5)

    def test_invoke_subgraph(self):
        @torch.compiler.nested_compile_region
        def fn(scale_obj, x):
            result = torch.ops._TestOpaqueObject.mul_with_scale(scale_obj, x)
            result = result * 2
            return result

        def gn(scale_obj, x):
            z = fn(scale_obj, x)
            return z + z

        scale_obj = OpaqueMultiplier(3.0)
        x = torch.randn(2, 2, requires_grad=True)
        x_clone = x.detach().clone().requires_grad_(True)

        ref = gn(scale_obj, x_clone)
        ref.sum().backward()

        backend = AotEagerAndRecordGraphs()
        opt_outer = torch.compile(gn, fullgraph=True, backend=backend)
        res = opt_outer(scale_obj, x)
        res.sum().backward()

        actual = normalize_gm(backend.graphs[0].print_readable(print_output=False))
        # Normalize module-qualified opaque type names since they differ
        # depending on how the test is invoked (__main__ vs test_opaque_obj_v2).
        actual = re.sub(r"\w+_OpaqueMultiplier", "OpaqueMultiplier", actual)
        self.assertExpectedInline(
            actual,
            """\
class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[2, 2]", L_scale_obj_ : OpaqueMultiplier):
        l_x_ = L_x_
        l_scale_obj_ = L_scale_obj_

        subgraph_0 = self.subgraph_0
        invoke_subgraph = torch.ops.higher_order.invoke_subgraph(subgraph_0, 'subgraph_0', l_scale_obj_, l_x_);  subgraph_0 = l_scale_obj_ = l_x_ = None
        getitem: "f32[2, 2]" = invoke_subgraph[0];  invoke_subgraph = None

        add: "f32[2, 2]" = getitem + getitem;  getitem = None
        return (add,)

    class subgraph_0(torch.nn.Module):
        def forward(self, l_scale_obj_ : OpaqueMultiplier, l_x_: "f32[2, 2]"):
            result: "f32[2, 2]" = torch.ops._TestOpaqueObject.mul_with_scale(l_scale_obj_, l_x_);  l_scale_obj_ = l_x_ = None

            result_1: "f32[2, 2]" = result * 2;  result = None
            return (result_1,)
""",  # noqa: B950
        )

        self.assertEqual(ref, res)
        self.assertEqual(ref.grad, res.grad)

    def test_value_type_subclass(self):
        """Subclasses of registered opaque types should be recognized as opaque."""

        class ExtendedConfig(ValueConfig):
            def __init__(self, mode: str, scale: float):
                super().__init__(mode)
                self.scale = scale

            def __eq__(self, other):
                return (
                    isinstance(other, ExtendedConfig)
                    and self.mode == other.mode
                    and self.scale == other.scale
                )

            def __hash__(self):
                return hash((self.mode, self.scale))

            def __fx_repr__(self):
                return (
                    f"ExtendedConfig(mode={self.mode!r}, scale={self.scale!r})",
                    {"ExtendedConfig": ExtendedConfig},
                )

        self.assertTrue(is_opaque_type(ExtendedConfig))
        self.assertTrue(is_opaque_value_type(ExtendedConfig))

        cfg = ExtendedConfig("square", 2.0)

        def foo(x):
            return torch.ops._TestOpaqueObject.process_with_config(x, cfg)

        x = torch.randn(3, 3)
        opt_f = torch.compile(foo, fullgraph=True, backend="aot_eager")
        res = opt_f(x)
        self.assertEqual(res, foo(x))

    def test_opaque_object_with_inductor_backend(self):
        """Test that opaque objects work correctly with inductor's get_attr handling."""

        class TestModule(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.color = Color.RED

            def forward(self, x):
                return torch.ops._TestOpaqueObject.apply_color_scale(self.color, x)

        mod = TestModule()
        x = torch.randn(3, 3)

        gm = torch.fx.symbolic_trace(mod)

        has_get_attr = any(node.op == "get_attr" for node in gm.graph.nodes)
        self.assertTrue(has_get_attr, "expected get_attr node for opaque object")

        compiled_fn = compile_fx(gm, [x])
        result = compiled_fn(x)

        expected = x * float(Color.RED.value)
        self.assertTrue(torch.allclose(result, expected))

    def test_captured_opaque_object_export(self):
        """Test that opaque objects captured from closures work with torch.export.

        When an opaque object is captured via a module hook closure, it gets a
        CellContentsSource instead of GetItemSource. The export path needs to
        handle this gracefully.
        """
        rng = RNGState(42)

        class Model(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(4, 4)

            def forward(self, x):
                torch.ops._TestOpaqueObject.noisy_inject(x, rng)
                return self.linear(x)

        ep = torch.export.export(Model(), (torch.randn(4, 4),), strict=True)
        self.assertExpectedInline(
            ep.graph_module.code.strip(),
            """\
def forward(self, p_linear_weight, p_linear_bias, obj_lifted_custom_0, x):
    noisy_inject = torch.ops._TestOpaqueObject.noisy_inject.default(x, obj_lifted_custom_0);  obj_lifted_custom_0 = noisy_inject = None
    linear = torch.ops.aten.linear.default(x, p_linear_weight, p_linear_bias);  x = p_linear_weight = p_linear_bias = None
    return (linear,)""",  # noqa: B950
        )

    def test_hoist_no_recompile_on_different_string(self):
        cnt = CompileCounter()

        def f(x, label):
            return op_with_string(x, HoistedString(label))

        opt_f = torch.compile(f, backend=cnt, fullgraph=True)
        x = torch.tensor(3.0)

        res1 = opt_f(x, "double")
        self.assertEqual(res1, f(x, "double"))
        self.assertEqual(cnt.frame_count, 1)

        res2 = opt_f(x, "square")
        self.assertEqual(res2, f(x, "square"))
        self.assertEqual(cnt.frame_count, 1)

    def test_opaque_multi_output_not_tensor_irnode(self):
        """OpaqueMultiOutput must not be classified as a tensor IR node.

        _is_tensor_irnode is used to select nodes for stride constraint
        functions (require_stride1, require_contiguous, etc). If
        OpaqueMultiOutput passes the check, those functions would call
        get_stride/get_dtype/make_loader on it and crash."""
        from unittest.mock import MagicMock

        from torch._inductor import ir
        from torch._inductor.lowering import _is_tensor_irnode

        opaque_multi = ir.OpaqueMultiOutput.__new__(ir.OpaqueMultiOutput)
        self.assertIsInstance(opaque_multi, ir.IRNode)
        self.assertFalse(_is_tensor_irnode(opaque_multi))

        non_tensor = MagicMock(spec=ir.NonTensorObj)
        non_tensor.__class__ = ir.NonTensorObj
        self.assertFalse(_is_tensor_irnode(non_tensor))

        self.assertFalse(_is_tensor_irnode(42))

    @unittest.skipIf(not dist.is_available(), "requires distributed")
    def test_fake_script_object_process_group_pybind(self):
        """FakeScriptObject wrapping ProcessGroup must be unwrapped in the
        pybind toIValue before casting to c10::intrusive_ptr<ProcessGroup>.

        During Dynamo tracing with CooR, mesh_get_process_group returns a
        FakeScriptObject-wrapped ProcessGroup. When that flows into a C++
        custom op, toIValue needs to unwrap real_obj before the pybind
        cast. This test exercises that path by calling
        mesh_get_process_group (which returns a wrapped PG during tracing)
        and passing it to another op that consumes the ProcessGroup."""
        from torch.distributed.device_mesh import (
            _register_distributed_opaque_types,
            DeviceMesh,
        )
        from torch.testing._internal.distributed.fake_pg import FakeStore

        already_initialized = dist.is_initialized()
        if already_initialized:
            dist.destroy_process_group()

        dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=2)
        try:
            _register_distributed_opaque_types()
            mesh = DeviceMesh("cpu", torch.arange(2))

            def f(mesh, x):
                pg = torch.ops._dtensor.mesh_get_process_group(mesh, 0)
                return x + pg.size()

            x = torch.randn(4)
            compiled_f = torch.compile(f, backend="aot_eager", fullgraph=True)
            result = compiled_f(mesh, x)
            expected = f(mesh, x)
            self.assertEqual(result, expected)
        finally:
            dist.destroy_process_group()
            if already_initialized:
                dist.init_process_group("fake", store=FakeStore(), rank=0, world_size=2)

    def test_enum_export(self):
        class Direction(enum.Enum):
            UP = 0
            DOWN = 1

        class Mod(torch.nn.Module):
            def forward(self, x, d):
                return x + d.value

        ep = torch.export.export(Mod(), (torch.randn(4, 4), Direction.UP), strict=False)
        self.assertEqual(
            ep.module()(torch.ones(4, 4), Direction.UP),
            torch.ones(4, 4) + Direction.UP.value,
        )
        self.assertExpectedInline(
            normalize_gm(ep.graph_module.print_readable(False)),
            """\
class GraphModule(torch.nn.Module):
    def forward(self, x: "f32[4, 4]", d):
        add: "f32[4, 4]" = torch.ops.aten.add.Tensor(x, 0);  x = None
        return (add,)
""",
        )

        backend = EagerAndRecordGraphs()
        opt_fn = torch.compile(Mod(), backend=backend)
        x = torch.randn(4, 4)
        res = opt_fn(x, Direction.UP)
        self.assertEqual(
            res,
            x + Direction.UP.value,
        )
        self.assertExpectedInline(
            normalize_gm(backend.graphs[0].print_readable(False)),
            """\
class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[4, 4]"):
        l_x_ = L_x_

        add: "f32[4, 4]" = l_x_ + 0;  l_x_ = None
        return (add,)
""",
        )

    def test_enum_custom_op(self):
        def get_color():
            class Color(enum.Enum):
                RED = 0
                GREEN = 1
                BLUE = 2

            return Color

        Color = get_color()

        @torch.library.custom_op("test_enum::add_color", mutates_args=())
        def add_color(x: torch.Tensor, c: Color) -> torch.Tensor:
            return x + c.value

        @add_color.register_fake
        def _(x, c):
            return torch.empty_like(x)

        def fn(x, c):
            return add_color(x, c)

        x = torch.randn(4, 4)
        ref = fn(x, Color.GREEN)
        backend = EagerAndRecordGraphs()
        opt_fn = torch.compile(fn, backend=backend, fullgraph=True)
        res = opt_fn(x, Color.GREEN)
        self.assertEqual(ref, res)
        self.assertExpectedInline(
            normalize_gm(backend.graphs[0].print_readable(False)),
            """\
class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[4, 4]"):
        l_x_ = L_x_

        add_color_default: "f32[4, 4]" = torch.ops.test_enum.add_color.default(l_x_, Color.GREEN);  l_x_ = None
        return (add_color_default,)
""",
        )

    def test_subclass_parametrization_with_opaque_attrs(self):
        """unwrap_tensor_subclass_parameters should handle non-tensor attrs."""
        from torch._functorch._aot_autograd.subclass_parametrization import (
            unwrap_tensor_subclass_parameters,
        )

        class M(torch.nn.Module):
            def __init__(self):
                super().__init__()
                a = torch.randn(3, 4)
                b = torch.randn(3, 4)
                self.p = torch.nn.Parameter(
                    TensorWithCounter(a, b, Counter(1, 10), SizeStore(4))
                )

            def forward(self, x):
                return x + self.p

        m = M()
        unwrap_tensor_subclass_parameters(m)
        x = torch.randn(3, 4)
        m(x)

    def test_export_non_strict_with_opaque_attrs(self):
        """torch.export with strict=False should handle subclasses with non-tensor attrs."""

        class M(torch.nn.Module):
            def forward(self, x):
                return x * 2

        a = torch.randn(4, 4)
        b = torch.randn(4, 4)
        x = TensorWithCounter(a, b, Counter(1, 10), SizeStore(4))
        torch.export.export(M(), (x,), strict=False)

    @unittest.skipIf(not dist.is_available(), "requires distributed")
    def test_get_untyped_storages_with_opaque_attrs(self):
        """get_untyped_storages should skip non-tensor attrs without warning."""
        from torch.distributed._tools.common_utils import get_untyped_storages

        a = torch.randn(4, 4)
        b = torch.randn(4, 4)
        x = TensorWithCounter(a, b, Counter(1, 10), SizeStore(4))
        import warnings

        with warnings.catch_warnings():
            warnings.simplefilter("error")
            storages = get_untyped_storages(x)
        # Should get storages from the two inner tensors (a and b)
        self.assertTrue(len(storages) > 0)

    @functorch_config.patch({"enable_autograd_cache": True})
    @inductor_config.patch(
        {
            "fx_graph_cache": True,
            "fx_graph_remote_cache": False,
        }
    )
    def test_subclass_opaque_attrs_cache_hit(self):
        """opaque_attrs in SubclassCreationMeta should survive AOTAutograd cache round-trip."""
        torch._dynamo.reset()
        counters.clear()

        with fresh_inductor_cache():

            @torch.compile(fullgraph=True)
            def fn(x):
                return x * 2

            a = torch.randn(4, 4)
            b = torch.randn(4, 4)
            counter = Counter(start=3, end=10)
            size = SizeStore(4)
            x = TensorWithCounter(a, b, counter, size)

            out = fn(x)
            self.assertIsInstance(out, TensorWithCounter)
            self.assertIs(out._counter, counter)
            self.assertIs(out._size_store, size)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_miss"], 1)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 0)

            # Reset dynamo to force a fresh compilation that should hit the cache
            torch._dynamo.reset()

            @torch.compile(fullgraph=True)
            def fn2(x):
                return x * 2

            a2 = torch.randn(4, 4)
            b2 = torch.randn(4, 4)
            counter2 = Counter(start=3, end=10)
            size2 = SizeStore(4)
            x2 = TensorWithCounter(a2, b2, counter2, size2)

            out2 = fn2(x2)
            self.assertIsInstance(out2, TensorWithCounter)
            # Runtime opaque values should be used, not the serialized ones
            self.assertIs(out2._counter, counter2)
            self.assertEqual(out2._size_store, size2)
            self.assertEqual(counters["aot_autograd"]["autograd_cache_hit"], 1)

    def test_op_passthrough_counter_in_tuple(self):
        # When a fake kernel returns its Counter input directly, the getitem
        # proxy's example_value is already a FakeScriptObject.
        counter_type = get_opaque_type_name(Counter)
        torch.library.define(
            "_TestOpaqueObject::passthrough_counter",
            f"({counter_type} c, Tensor x) -> ({counter_type}, Tensor)",
            tags=torch.Tag.pt2_compliant_tag,
            lib=self.lib,
        )

        @torch.library.impl(
            "_TestOpaqueObject::passthrough_counter",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def passthrough_impl(c: Counter, x: torch.Tensor):
            return c, x * c.start

        @torch.library.register_fake(
            "_TestOpaqueObject::passthrough_counter", lib=self.lib
        )
        def passthrough_fake(c: Counter, x: torch.Tensor):
            return c, torch.empty_like(x)

        def fn(c, x):
            out_c, out_x = torch.ops._TestOpaqueObject.passthrough_counter(c, x)
            return torch.ops._TestOpaqueObject.counter_start(out_c) + out_x

        c = Counter(3, 10)
        x = torch.randn(4)
        ref = fn(c, x)
        opt_fn = torch.compile(fn, fullgraph=True, backend="eager")
        res = opt_fn(c, x)
        self.assertEqual(ref, res)

    @torch._dynamo.config.patch(skip_fwd_side_effects_in_bwd_under_checkpoint=True)
    def test_script_object_intermediate_exposed_from_checkpoint(self):
        # A TorchScriptObjectVariable created inside an AC region and accessed
        # outside via a list side effect must be exposed as a subgraph output.
        import torch.utils.checkpoint

        def gn(x, results):
            counter = torch.ops._TestOpaqueObject.create_counter(x.shape[0], x.shape[0])
            results.append(counter)
            return x * 2

        def fn(x):
            results = []
            out = torch.utils.checkpoint.checkpoint(gn, x, results, use_reentrant=False)
            return torch.ops._TestOpaqueObject.counter_start(results[0]) + out

        x = torch.randn(3, 4)
        ref = fn(x)
        opt_fn = torch.compile(fn, fullgraph=True, backend="aot_eager_decomp_partition")
        res = opt_fn(x)
        self.assertEqual(ref, res)

    def test_reference_type_opaque_object_state(self):
        """When compile_fx_inner receives an FX graph whose placeholder meta['val']
        is a raw opaque reference type (not wrapped in FakeScriptObject), inductor
        handles it via OpaqueObjectState. This codepath is used by compile-on-one-rank
        (COOR) where the FX graph is constructed with real objects."""
        m = OpaqueMultiplier(2.0)
        x = torch.ones(3)

        graph = torch.fx.Graph()
        m_node = graph.placeholder("m")
        m_node.meta["val"] = m
        fake_mode = FakeTensorMode()
        x_node = graph.placeholder("x")
        x_node.meta["val"] = fake_mode.from_tensor(x)
        out = graph.call_function(
            torch.ops._TestOpaqueObject.mul_with_scale.default, (m_node, x_node)
        )
        out.meta["val"] = fake_mode.from_tensor(x)
        graph.output((out,))

        gm = torch.fx.GraphModule({}, graph)
        compiled = compile_fx_inner(gm, [m, x])
        result = compiled([m, x])
        self.assertEqual(result, (x * 2,))

    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
    def test_benchmark_harness_no_pickle_for_opaque_inputs(self):
        """Opaque graph inputs must not be pickled in the benchmark harness."""
        a = torch.randn(4, 4, device="cuda")
        b = torch.randn(4, 4, device="cuda")
        twc = TensorWithCounter(a, b, Counter(0, 10), SizeStore(4))

        def fn(x):
            return x + 1

        compiled_fn = torch.compile(fn, backend="inductor", fullgraph=True)
        _, codes = run_and_get_code(compiled_fn, twc)
        self.assertGreater(len(codes), 0)
        for code in codes:
            self.assertNotIn("pickle", code)

    def test_opaque_object_state_in_graph_output(self):
        """When compile_fx_inner receives a graph where a raw opaque reference
        type appears in the outputs (not just inputs), inductor must handle it.
        This happens in CooR (compile-on-one-rank) precompile: aot_autograd
        partitions the joint graph so the forward graph saves opaque objects
        (e.g. DeviceMesh) for the backward graph by including them in forward
        outputs.  Inductor must:
          1. Accept OpaqueObjectState in the output type assertion
             (GraphLowering.output).
          2. Handle NonTensorObj graph inputs in the memory planner
             (get_dep_size_hint) without calling get_numel/get_dtype on them."""
        m = OpaqueMultiplier(2.0)
        x = torch.ones(3)

        graph = torch.fx.Graph()
        m_node = graph.placeholder("m")
        m_node.meta["val"] = m
        fake_mode = FakeTensorMode()
        x_node = graph.placeholder("x")
        x_node.meta["val"] = fake_mode.from_tensor(x)
        out = graph.call_function(
            torch.ops._TestOpaqueObject.mul_with_scale.default, (m_node, x_node)
        )
        out.meta["val"] = fake_mode.from_tensor(x)
        # Include the opaque object in the output tuple, simulating how
        # aot_autograd's forward graph passes saved-for-backward objects
        # through as outputs.
        graph.output((out, m_node))

        gm = torch.fx.GraphModule({}, graph)
        compiled = compile_fx_inner(gm, [m, x])
        result = compiled([m, x])
        self.assertEqual(result[0], x * 2)
        self.assertIs(result[1], m)

    def test_reconstruct_fn_sets_meta_val(self):
        """Opaque nodes created via reconstruct_fn have meta['val'] set.

        When _try_reconstruct_opaque dispatches through a custom op via
        Proxy.__torch_function__, the resulting node does not get
        meta['val'] set automatically (unlike the __torch_dispatch__
        path which calls track_tensor_tree).  The fix in
        _try_reconstruct_opaque ensures set_meta is called so that
        downstream consumers like the min-cut partitioner can classify
        the node correctly via is_opaque_node().
        """
        from torch._functorch._aot_autograd.graph_compile import is_opaque_node

        # Use the already-registered OpaqueMultiplier type.
        # Register a reconstruct_fn that derives one multiplier from
        # another via a custom op — mirrors how DeviceMesh submeshes
        # are derived from a parent mesh via _get_submesh.
        multiplier_type = get_opaque_type_name(OpaqueMultiplier)

        self.lib.define(
            f"derive_multiplier({multiplier_type} parent, float scale)"
            f" -> {multiplier_type}",
        )

        @torch.library.impl(
            "_TestOpaqueObject::derive_multiplier",
            "CompositeExplicitAutograd",
            lib=self.lib,
        )
        def derive_impl(parent, scale):
            return OpaqueMultiplier(parent.multiplier * scale)

        @torch.library.register_fake(
            "_TestOpaqueObject::derive_multiplier", lib=self.lib
        )
        def derive_fake(parent, scale):
            return OpaqueMultiplier(0.0)

        parent = OpaqueMultiplier(3.0)
        # The "child" is derived at eager time before tracing.
        # It's NOT passed as an input — it's captured in a closure,
        # simulating how DeviceMesh submeshes are captured in DTensor
        # backward closures.  This forces _try_reconstruct_opaque to
        # be called when make_fx encounters the child.
        child = OpaqueMultiplier(parent.multiplier * 0.5)
        child._parent = parent
        child._scale = 0.5

        from torch._library.opaque_object import _OPAQUE_TYPES

        original_reconstruct_fn = _OPAQUE_TYPES[OpaqueMultiplier].reconstruct_fn

        def multiplier_reconstruct_fn(obj, get_tracked_proxy, tracer):
            if not hasattr(obj, "_parent"):
                return None
            parent_proxy = get_tracked_proxy(obj._parent)
            if parent_proxy is None:
                return None
            return torch.ops._TestOpaqueObject.derive_multiplier(
                parent_proxy, obj._scale
            )

        _OPAQUE_TYPES[OpaqueMultiplier].reconstruct_fn = multiplier_reconstruct_fn
        try:
            # child is captured from the closure, NOT an input
            def fn(parent, x):
                return torch.ops._TestOpaqueObject.mul_with_scale(child, x)

            fake_mode = FakeTensorMode(shape_env=ShapeEnv())
            with fake_mode:
                fake_x = torch.randn(4)

            gm = make_fx(fn, tracing_mode="fake")(parent, fake_x)

            # Find the derive_multiplier node (created by reconstruct_fn)
            derive_nodes = [
                n
                for n in gm.graph.nodes
                if n.op == "call_function" and "derive_multiplier" in str(n.target)
            ]
            self.assertGreater(len(derive_nodes), 0, "No derive_multiplier node found")

            for node in derive_nodes:
                self.assertIn(
                    "val",
                    node.meta,
                    f"Node {node.name} created via reconstruct_fn is missing "
                    f"meta['val']. This would cause the partitioner to fail "
                    f"with 'Expected {node.name} to be a tensor'.",
                )
                self.assertTrue(
                    is_opaque_node(node),
                    f"Node {node.name} should be classified as opaque",
                )
        finally:
            _OPAQUE_TYPES[OpaqueMultiplier].reconstruct_fn = original_reconstruct_fn

    def test_partitioner_must_save_opaque_node(self):
        """Opaque nodes tagged MUST_SAVE go to saved_opaque_nodes.

        When activation checkpointing tags an opaque node with
        CheckpointPolicy.MUST_SAVE, the default_partition function
        must route it to saved_opaque_nodes (not saved_values).
        Otherwise the runtime assertion in save_from_forward will
        fail because it expects all saved_values to be Tensors.
        """
        from functorch.compile import default_partition
        from torch._functorch._aot_autograd.graph_compile import is_opaque_node
        from torch.utils.checkpoint import CheckpointPolicy

        m = OpaqueMultiplier(2.0)

        fake_mode = FakeTensorMode(shape_env=ShapeEnv())

        # Build a joint fwd+bwd graph where the opaque node is tagged
        # MUST_SAVE — simulates activation checkpointing.
        graph = torch.fx.Graph()
        m_node = graph.placeholder("m")
        m_node.meta["val"] = m
        x_node = graph.placeholder("x")
        with fake_mode:
            x_node.meta["val"] = torch.randn(3)

        mul_node = graph.call_function(
            torch.ops._TestOpaqueObject.mul_with_scale.default,
            (m_node, x_node),
        )
        with fake_mode:
            mul_node.meta["val"] = torch.randn(3)

        # Tag the opaque node as MUST_SAVE
        m_node.meta["recompute"] = CheckpointPolicy.MUST_SAVE

        # Create backward portion: identity backward
        bw_node = graph.call_function(torch.ops.aten.mul.Scalar, (mul_node, 1.0))
        with fake_mode:
            bw_node.meta["val"] = torch.randn(3)

        graph.output((mul_node, bw_node))
        joint_gm = torch.fx.GraphModule({}, graph)

        # Run default_partition — this should NOT raise even though
        # the opaque node is tagged MUST_SAVE.
        num_fwd_outputs = 1
        fw_module, bw_module = default_partition(
            joint_gm, [], num_fwd_outputs=num_fwd_outputs
        )

        # Verify the opaque node ended up in the forward graph as
        # a pass-through (not as a saved tensor that would fail the
        # save_from_forward assertion).
        fw_opaque_nodes = [
            n
            for n in fw_module.graph.nodes
            if n.op == "placeholder" and is_opaque_node(n)
        ]
        self.assertGreater(
            len(fw_opaque_nodes),
            0,
            "Forward graph should have opaque placeholder nodes",
        )

    def test_hoisted_opaque_with_graphpickler(self):
        if import_dill() is None:
            self.skipTest("dill not available")

        backend = EagerAndRecordGraphs()

        def fn(x):
            return torch.ops.mylib.op_with_string(x, HoistedString("double"))

        compiled = torch.compile(fn, fullgraph=True, backend=backend)
        compiled(torch.randn(3, 4))

        self.assertEqual(len(backend.graphs), 1)
        captured_gm = backend.graphs[0]

        placeholders = [n for n in captured_gm.graph.nodes if n.op == "placeholder"]
        opaque_ph = None
        for ph in placeholders:
            ev = ph.meta.get("example_value")
            if isinstance(ev, FakeScriptObject):
                opaque_ph = ph
                break
        self.assertIsNotNone(opaque_ph, "no placeholder with FakeScriptObject found")

        fake_mode = FakeTensorMode(shape_env=ShapeEnv())
        options = Options(ops_filter=None, node_metadata_key_filter=None)
        data = GraphPickler.dumps(captured_gm, options)
        restored_gm = GraphPickler.loads(data, fake_mode)

        restored_phs = [n for n in restored_gm.graph.nodes if n.op == "placeholder"]
        restored_opaque = None
        for ph in restored_phs:
            ev = ph.meta.get("example_value")
            if isinstance(ev, FakeScriptObject):
                restored_opaque = ph
                break
        restored_real = restored_opaque.meta["example_value"].real_obj
        self.assertIsInstance(restored_real, HoistedString)
        self.assertEqual(restored_real.val, "double")


instantiate_parametrized_tests(TestOpaqueObject)


if __name__ == "__main__":
    run_tests()
