# torch/_dynamo/guards.py (简化)
def TENSOR_MATCH(self, guard: Guard, value: Optional[Any] = None) -> None:
    if match_on_id_for_tensor(guard):
        # 对于 FX 图模块中的 tensor，ID 匹配即可
        self.ID_MATCH(guard)
    else:
        value = value if value is not None else self.get(guard)
        pytype = type(value)
        dispatch_keys = torch._C._dispatch_keys(value)

        # 获取 size/stride 元数据
        metadata = output_graph.input_source_to_sizes_strides[
            guard.originating_source
        ]
        size = convert_to_concrete_values(metadata["size"])
        stride = convert_to_concrete_values(metadata["stride"])

        # 安装 C++ 层的 tensor match guard
        guard_manager.add_tensor_match_guard(
            value, size, stride, tensor_name,
            verbose_code_parts, user_stack,
            pytype, dispatch_keys,
        )

# 适用类型包括：
# int, float, bool, str, bytes, NoneType, type,
# tuple, frozenset, slice, range, torch.Size, ...

@torch.compile
def fn(x, mode="train"):
    if mode == "train":
        return x * 2
    return x

# Dynamo 会为 mode 生成 EQUALS_MATCH guard:
# guard: L['mode'] == 'train'

def TYPE_MATCH(self, guard: Guard) -> None:
    value = self.get(guard)
    t = type(value)
    obj_id = self.id_ref(t, f"type({guard.name})")
    # 安装类型检查 guard
    self.get_guard_manager(guard).add_type_match_guard(
        obj_id, verbose_code_parts, guard.user_stack,
    )

def ID_MATCH(self, guard: Guard, recompile_hint=None) -> None:
    # 如果 source 是 TypeSource，优化为 TYPE_MATCH
    if isinstance(guard.originating_source, TypeSource):
        return self.TYPE_MATCH(...)

    val = self.get(guard)
    id_val = self.id_ref(val, guard.name)
    # 安装 ID 检查 guard
    self.get_guard_manager(guard).add_id_match_guard(
        id_val, verbose_code_parts, guard.user_stack,
    )

# torch/_guards.py
class GuardSource(enum.Enum):
    LOCAL = 0                                    # 函数局部变量
    GLOBAL = 1                                   # 全局变量
    LOCAL_SPECIALIZED_NN_MODULE = 2              # 特化的本地 nn.Module
    GLOBAL_SPECIALIZED_NN_MODULE = 3             # 特化的全局 nn.Module
    CONSTANT = 4                                 # 常量
    RANDOM_VALUE = 5                             # 随机数相关
    SHAPE_ENV = 6                                # 符号形状环境
    LOCAL_FSDP_MODULE = 7                        # 本地 FSDP 模块
    GLOBAL_FSDP_MODULE = 8                       # 全局 FSDP 模块
    BACKWARD_STATE = 9                           # 反向传播状态
    EPHEMERAL = 10                               # 临时变量
    SYNTHETIC_LOCAL = 11                         # 合成局部变量
    LOCAL_UNSPECIALIZED_NN_MODULE = 12           # 未特化的本地 nn.Module
    GLOBAL_UNSPECIALIZED_NN_MODULE = 13          # 未特化的全局 nn.Module
    LOCAL_UNSPECIALIZED_BUILTIN_NN_MODULE = 14   # 未特化的内置 nn.Module
    GLOBAL_UNSPECIALIZED_BUILTIN_NN_MODULE = 15  # 未特化的内置 nn.Module
    TEMP_LOCAL = 16                              # 临时局部变量

# torch/_guards.py 中的 Guard.create 方法
def create(self, builder: GuardBuilderBase) -> Any:
    return self.create_fn(builder, self)

# 使用示例：追踪期间创建 guard
guard = Guard(source, GuardBuilder.TENSOR_MATCH)
# 后续 guard.create(builder) 会调用 builder.TENSOR_MATCH(guard)

# GuardManagerWrapper 中的 diff guard 逻辑
def collect_diff_guard_sources(self):
    # 遍历 guard 树，收集 fail_count > 0 的节点
    # 这些节点是"高风险"的，优先检查
    def visit(node):
        is_diff_guard_node = (
            node.get_source() in self.diff_guard_sources
            or node.fail_count() > 0
        )
        for child_mgr in node.get_child_managers():
            is_diff_guard_node |= visit(child_mgr)
        if is_diff_guard_node:
            self.diff_guard_sources.add(node.get_source())
        return is_diff_guard_node
    visit(self.root)

# torch/_dynamo/guards.py
def get_and_maybe_log_recompilation_reasons(cache_entry, frame, backend):
    reasons = []
    while cache_entry is not None:
        reason = get_guard_fail_reason(
            cache_entry.guard_manager,
            cache_entry.code,
            frame.f_locals,
            cache_entry.compile_id,
            backend,
        )
        if reason:
            reasons.append(reason)
        cache_entry = cache_entry.next
    return reasons

# torch/_dynamo/cache_size.py
# 缓存大小计算会考虑 id_matched_objs 的影响
# 只有同一个对象多次触发不同的编译时，才计入重编译次数

import torch

x = torch.randn(4, 8)
torch._dynamo.mark_dynamic(x, 0)  # 标记第 0 维为动态
torch._dynamo.mark_dynamic(x, 1, min=1, max=128)  # 标记第 1 维为动态，带约束

@torch.compile
def fn(x):
    return x.sum(dim=0)

fn(x)  # 编译时第 0 维是符号化的

# TENSOR_MATCH 中的动态维度 guard
if hasattr(value, "_dynamo_dynamic_indices"):
    dynamic_indices = value._dynamo_dynamic_indices
    code_part = (
        f"(({tensor_name}._dynamo_dynamic_indices"
        f".issubset({dynamic_indices})) "
        f"if hasattr({tensor_name}, '_dynamo_dynamic_indices') "
        f"else True)"
    )

# 是否启用自动动态形状（默认 True）
torch._dynamo.config.automatic_dynamic_shapes = True

# 自动动态标记的模式："dynamic" 或 "unbacked"
torch._dynamo.config.automatic_dynamic_shapes_mark_as = "dynamic"

# 是否默认假设静态形状（默认 True）
torch._dynamo.config.assume_static_by_default = True

class MutationTracker:
    db: ExactWeakKeyDictionary = ExactWeakKeyDictionary()

    def on_mutation(self, name: str) -> None:
        self.mutation_count += 1
        # 通知所有关联的 guarded_code 失效
        for ref in self.watchers:
            guarded = ref()
            if guarded is not None:
                guarded.invalidate(ref)

def ensure_patched(cls):
    """Monkey-patch cls.__setattr__ 以拦截属性修改"""
    original_setattr = cls.__setattr__

    def custom_setattr(self, key, value):
        try:
            MutationTracker.db[self].on_mutation(key)
        except KeyError:
            pass
        return original_setattr(self, key, value)

    cls.__setattr__ = custom_setattr

class GenerationTracker:
    generation: int = 0

    @classmethod
    def tag(cls, obj):
        cls.generation_values[obj] = cls.generation

    @classmethod
    def check(cls, obj):
        return (
            obj in cls.generation_values
            and cls.generation_values[obj] == cls.generation
        )

import torch

@torch.compile
def fn(x, flag):
    if flag:
        return x * 2
    return x + 1

# 第一次调用：编译并生成 guard
x = torch.randn(4, 4)
result1 = fn(x, True)
# Guard 集合包含：
#   TENSOR_MATCH(L['x']): shape=(4,4), dtype=float32, ...
#   EQUALS_MATCH(L['flag']): True

# 第二次调用：guard 全部通过，复用编译结果
x2 = torch.randn(4, 4)
result2 = fn(x2, True)  # 快速路径

# 第三次调用：flag 变化，EQUALS_MATCH 失败，触发重编译
result3 = fn(x2, False)  # 重编译
# 现在缓存中有两个条目，分别对应 flag=True 和 flag=False

import torch

@torch.compile
def fn(x):
    return x.sum()

# 静态 shape 编译
x1 = torch.randn(4, 8)
fn(x1)

# shape 变化，触发重编译
# 如果 automatic_dynamic_shapes=True，第 0 维和第 1 维会被自动标记为动态
x2 = torch.randn(8, 16)
fn(x2)

# 此后任何 shape 都不会再触发重编译（只要满足基本约束）
x3 = torch.randn(32, 64)
fn(x3)  # 复用第二次的编译结果

import torch._dynamo.config

# 增大缓存限制（允许更多重编译）
torch._dynamo.config.cache_size_limit = 64

# 在超过限制时立即报错（而非静默回退）
torch._dynamo.config.fail_on_recompile_limit_hit = True

# 在任何重编译时报错（用于调试）
torch._dynamo.config.error_on_recompile = True

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, use_relu=True):
        super().__init__()
        self.linear = nn.Linear(10, 10)
        self.use_relu = use_relu

    def forward(self, x):
        x = self.linear(x)
        if self.use_relu:
            x = torch.relu(x)
        return x

model = MyModel()
compiled = torch.compile(model)

x = torch.randn(4, 10)
compiled(x)

# 修改模块属性会触发 MutationTracker，使缓存失效
model.use_relu = False
compiled(x)  # 重编译，因为 mutation guard 检测到属性变化

# config.guard_nn_modules 控制是否对 nn.Module 的属性安装 guard
# 默认为 True
print(torch._dynamo.config.guard_nn_modules)  # True

# GuardManagerWrapper.find_tag_safe_roots 方法
# 识别 "tag safe" 节点：
# 1. 不可变值（tensor、基本类型等）
# 2. 嵌套的 tag safe 字典
# 3. 纯 nn.Module（仅暴露 __dict__）
#
# 对于 tag safe root，只需检查顶层字典的 tag 即可跳过整个子树

def sort_key(self):
    return (
        is_duplicate_input,      # 重复输入 guard 放最后
        self.source.value,       # 按 source 类型排序
        len(self.name),          # 短名称优先
        self.name,               # 字典序
        self.inner_create_fn().__code__.co_firstlineno,  # 定义顺序
    )

# shape_env guard 被分离到 builder.shape_env_code 中
# 确保在 tensor match guard 之后运行
self.shape_env_code: list[GuardCodeList] = []

# 最终安装
if builder.shape_env_code:
    install_symbolic_shape_guard(...)

import torch._dynamo.config as config

config.cache_size_limit                  # 每帧最大缓存条目数 (默认 8)
config.accumulated_cache_size_limit      # 累计编译次数上限 (默认 256)
config.automatic_dynamic_shapes          # 自动动态形状 (默认 True)
config.assume_static_by_default          # 默认假设静态形状 (默认 True)
config.guard_nn_modules                  # 是否对 nn.Module 属性安装 guard (默认 True)
config.error_on_recompile                # 重编译时报错 (默认 False)
config.fail_on_recompile_limit_hit       # 超限时报错 (默认 False)
config.specialize_int                    # 是否特化 int 输入 (默认 False)
config.specialize_float                  # 是否特化 float 输入 (默认 False)