// torch/csrc/autograd/saved_variable.h
class TORCH_API SavedVariable {
 public:
  SavedVariable() = default;
  SavedVariable(const Variable& variable, bool is_output,
                bool is_inplace_on_view = false);

  // 恢复保存的变量
  Variable unpack(std::shared_ptr<Node> saved_for = nullptr) const;

  // 注册自定义 hooks
  void register_hooks(std::unique_ptr<SavedVariableHooks>&& hooks);

  void reset_data();

 private:
  at::Tensor data_;                      // 保存的张量数据
  std::shared_ptr<ForwardGrad> fw_grad_; // 前向梯度
  std::weak_ptr<Node> weak_grad_fn_;     // 弱引用（避免循环引用）

  uint32_t saved_version_ = 0;           // 保存时的版本号
  uint32_t output_nr_ = 0;              // 输出编号
  bool was_default_constructed_ = true;  // 是否默认构造
  bool is_inplace_on_view_ = false;     // 是否是 view 上的 inplace 操作
  bool saved_original_ = false;          // 是否保存了原始张量
  bool is_leaf_ = false;                // 是否是叶子节点
  bool is_output_ = false;              // 是否是某个 Node 的输出

  std::unique_ptr<SavedVariableHooks> hooks_;  // 自定义 pack/unpack hooks
  std::shared_ptr<Node> grad_fn_;              // 梯度函数
  std::shared_ptr<Node> grad_accumulator_;     // 梯度累积器
  bool requires_grad_ = false;
};

// torch/csrc/autograd/saved_variable.cpp
SavedVariable::SavedVariable(
    const Variable& variable, bool is_output, bool is_inplace_on_view) {
  if (variable.defined()) {
    // 1. 推理张量不能保存 - 会导致静默错误
    TORCH_CHECK(!variable.is_inference(),
        "Inference tensors cannot be saved for backward.");

    // 2. 记录版本号
    was_default_constructed_ = false;
    saved_version_ = variable._version();
    is_leaf_ = variable.is_leaf();
    is_output_ = is_output;

    // 3. 检查是否有默认 hooks
    std::unique_ptr<SavedVariableHooks> maybe_hooks =
        at::SavedTensorDefaultHooks::is_enabled()
            ? get_default_hooks() : nullptr;

    if (maybe_hooks && !variable.unsafeGetTensorImpl()->is_wrapped_number()) {
      save_metadata(variable);
      set_hooks_and_pack_data(std::move(maybe_hooks), variable);
      return;
    }

    // 4. 非输出或叶子节点 -> 直接保存原始张量（无循环引用风险）
    if (!is_output || is_leaf_) {
      saved_original_ = true;
      data_ = variable;
      return;
    }

    // 5. 输出张量 -> 保存 tensor_data()（避免循环引用）
    save_metadata(variable);
    data_ = variable.tensor_data();
  }
}

// torch/csrc/autograd/saved_variable.cpp
Variable SavedVariable::unpack(std::shared_ptr<Node> saved_for) const {
  // 1. 默认构造的 SavedVariable 返回空张量
  if (was_default_constructed_)
    return Variable();

  // 2. 数据已释放且无 hooks -> 二次 backward 错误
  if (!data_.defined())
    TORCH_CHECK(hooks_, ERR_BACKWARD_TWICE);

  // 3. 版本检查（仅在无 hooks 时检查）
  if (!hooks_) {
    auto current_version = impl::version_counter(data_).current_version();
    if (saved_version_ != current_version) {
      // 详细的错误信息：指出哪个张量被 inplace 修改
      TORCH_CHECK(false, "...modified by an inplace operation...");
    }
  }

  // 4. 如果保存了原始张量，直接返回
  if (!hooks_ && saved_original_)
    return data_;

  // 5. 通过 hooks 或 tensor_data 恢复
  auto data = hooks_ ? hooks_->call_unpack_hook() : data_;

  // 6. 重建 Variable 的 autograd 元数据
  Variable var;
  if (grad_fn)
    var = make_variable(data, Edge(std::move(grad_fn), output_nr_));
  else
    var = make_variable(data, requires_grad_);

  impl::set_grad_accumulator(var, grad_accumulator_);
  impl::set_version_counter(var, impl::version_counter(data));

  // 7. 恢复前向梯度
  if (fw_grad_ && !fw_grad_->empty()) {
    auto new_fw_grad = fw_grad_->value(0);
    var._set_fw_grad(new_fw_grad, 0, false);
  }

  return var;
}

if (saved_version_ != current_version) {
  // 构造详细的错误信息
  message << "one of the variables needed for gradient computation "
          << "has been modified by an inplace operation: ["
          << data_.toString() << ' ' << data_.sizes() << ']';
  if (grad_fn)
    message << ", which is output " << output_nr_ << " of " << grad_fn->name();
  message << " is at version " << current_version
          << "; expected version " << saved_version_;
}

# torch/autograd/graph.py
class saved_tensors_hooks:
    """设置 saved tensors 的 pack/unpack hooks。

    pack_hook(tensor: Tensor) -> Any
    unpack_hook(Any) -> Tensor
    """
    def __init__(self, pack_hook, unpack_hook):
        self.pack_hook = pack_hook
        self.unpack_hook = unpack_hook

    def __enter__(self):
        torch._C._autograd._push_saved_tensors_default_hooks(
            self.pack_hook, self.unpack_hook)

    def __exit__(self, *args):
        torch._C._autograd._pop_saved_tensors_default_hooks()

# torch/autograd/graph.py
class save_on_cpu(saved_tensors_hooks):
    """将保存的张量转移到 CPU，backward 时再搬回 GPU。"""
    def __init__(self, pin_memory=False, device_type="cuda"):
        device_module = getattr(torch, device_type, torch.cuda)

        def pack_to_cpu(tensor):
            if not pin_memory:
                return (tensor.device, tensor.cpu())
            packed = torch.empty(
                tensor.size(), dtype=tensor.dtype, layout=tensor.layout,
                pin_memory=(device_module.is_available() and not tensor.is_sparse))
            packed.copy_(tensor)
            return (tensor.device, packed)

        def unpack_from_cpu(packed):
            device, tensor = packed
            return tensor.to(device, non_blocking=pin_memory)

        super().__init__(pack_to_cpu, unpack_from_cpu)

# 节省 GPU 显存：中间结果暂存到 CPU
with torch.autograd.graph.save_on_cpu(pin_memory=True):
    output = model(input)  # 中间激活值保存在 CPU pinned memory
loss = criterion(output, target)
loss.backward()  # backward 时中间值从 CPU 异步拷回 GPU

// saved_variable.cpp - 构造函数中
if (maybe_hooks && !variable.unsafeGetTensorImpl()->is_wrapped_number()) {
    save_metadata(variable);
    set_hooks_and_pack_data(std::move(maybe_hooks), variable);
    return;
}

// set_hooks_and_pack_data 实现
void SavedVariable::set_hooks_and_pack_data(
    std::unique_ptr<SavedVariableHooks>&& hooks, const Variable& data) {
  hooks_ = std::move(hooks);
  at::NoGradGuard guard;  // pack_hook 中不记录梯度
  const auto version = impl::version_counter(data).current_version();
  hooks_->call_pack_hook(saved_original_ ? data.detach() : data);
  // 检查 pack_hook 没有修改输入
  TORCH_CHECK(version == impl::version_counter(data).current_version(),
      "A saved tensor pack hook is modifying its input in place.");
}

def checkpoint(function, *args, use_reentrant=True, **kwargs):
    """对 function 的执行应用梯度检查点。

    前向传播在 no_grad 模式下执行（不保存中间值）。
    反向传播时重新执行前向传播以获取中间激活值。
    """
    if use_reentrant:
        return CheckpointFunction.apply(function, preserve, *args)
    else:
        # 非重入版本使用 saved_tensors_hooks
        ...

class CheckpointFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, run_function, preserve_rng_state, *args):
        # 1. 保存输入和函数引用（不保存中间激活值）
        ctx.run_function = run_function
        ctx.save_for_backward(*args)

        # 2. 保存 RNG 状态（确保重算时随机性一致）
        if preserve_rng_state:
            ctx.fwd_cpu_state = torch.random.get_rng_state()
            ctx.fwd_gpu_state = torch.cuda.get_rng_state()

        # 3. 在 no_grad 下执行前向传播
        with torch.no_grad():
            outputs = run_function(*args)

        return outputs

    @staticmethod
    def backward(ctx, *args):
        inputs = ctx.saved_tensors

        # 1. 恢复 RNG 状态
        rng_devices = []
        if ctx.preserve_rng_state:
            # ... 恢复 CPU 和 GPU 的随机状态

        # 2. 分离输入并启用梯度追踪
        detached_inputs = detach_variable(inputs)

        # 3. 重新执行前向传播（这次会记录梯度）
        with torch.enable_grad():
            outputs = ctx.run_function(*detached_inputs)

        # 4. 对重算的输出执行 backward
        torch.autograd.backward(outputs, args)

        # 5. 返回输入的梯度
        grads = tuple(
            inp.grad if isinstance(inp, torch.Tensor) else None
            for inp in detached_inputs
        )
        return (None, None) + grads

def checkpoint_sequential(functions, segments, input, **kwargs):
    """将 functions 列表分成 segments 段，每段应用 checkpoint。"""
    # 将层列表平均分段
    segment_size = len(functions) // segments

    def run_segment(start, end, *inputs):
        for j in range(start, end):
            inputs = functions[j](*inputs)
        return inputs

    # 对每段应用 checkpoint
    for start in range(0, len(functions), segment_size):
        end = min(start + segment_size, len(functions))
        input = checkpoint(run_segment, start, end, input, **kwargs)

    return input

model = nn.Sequential(
    nn.Linear(784, 512), nn.ReLU(),
    nn.Linear(512, 256), nn.ReLU(),
    nn.Linear(256, 128), nn.ReLU(),
    nn.Linear(128, 10),
)

# 分成 4 段，每段应用 checkpoint
output = checkpoint_sequential(model, segments=4, input=x)

class CheckpointPolicy(enum.Enum):
    ALWAYS = "always"  # 总是保存
    NEVER = "never"    # 从不保存（重算）
    PREFER_RECOMPUTE = "prefer_recompute"  # 倾向于重算
    PREFER_SAVE = "prefer_save"            # 倾向于保存

# torch/autograd/anomaly_mode.py
class detect_anomaly:
    """启用 autograd 异常检测的上下文管理器。

    功能：
    1. 前向传播时记录调用栈，backward 出错时打印
    2. 检测 backward 中产生的 NaN 值
    """
    def __init__(self, check_nan=True):
        self.prev = torch.is_anomaly_enabled()
        self.check_nan = check_nan
        self.prev_check_nan = torch.is_anomaly_check_nan_enabled()
        warnings.warn("Anomaly Detection has been enabled. "
                      "This mode will increase the runtime.")

    def __enter__(self):
        torch.set_anomaly_enabled(True, self.check_nan)

    def __exit__(self, *args):
        torch.set_anomaly_enabled(self.prev, self.prev_check_nan)

// torch/csrc/autograd/engine.cpp - evaluate_function()
if (AnomalyMode::is_enabled() && AnomalyMode::should_check_nan()) {
  AutoGradMode grad_mode(false);
  for (const auto i : c10::irange(num_outputs)) {
    auto& output = outputs[i];
    at::OptionalDeviceGuard guard(device_of(output));
    TORCH_CHECK(
        !output.defined() || !isnan(output)._is_any_true().item<bool>(),
        "Function '", fn.name(), "' returned nan values in its ",
        i, "th output.");
  }
}

// graph_task.h - set_exception_without_signal()
void GraphTask::set_exception_without_signal(const std::shared_ptr<Node>& fn) {
  if (!has_error_.exchange(true)) {
    if (AnomalyMode::is_enabled() && fn) {
      fn->metadata()->print_stack(fn->name());  // 打印前向传播的栈轨迹
    }
  }
}

# 场景：调试 NaN 梯度
with torch.autograd.detect_anomaly():
    x = torch.randn(3, requires_grad=True)
    y = x ** 2
    z = torch.log(y)     # 如果 y 含负数，log 会产生 NaN
    z.sum().backward()
    # detect_anomaly 会报告 NaN 产生的位置和前向调用栈

class set_detect_anomaly:
    def __init__(self, mode: bool, check_nan: bool = True):
        self.prev = torch.is_anomaly_enabled()
        self.prev_check_nan = torch.is_anomaly_check_nan_enabled()
        torch.set_anomaly_enabled(mode, check_nan)

    def __enter__(self):
        pass

    def __exit__(self, *args):
        torch.set_anomaly_enabled(self.prev, self.prev_check_nan)

x = torch.randn(3, requires_grad=True)
y = x * 2
x.add_(1)      # inplace 修改了 x
y.backward()    # 报错：版本不匹配

x = torch.randn(3, requires_grad=True)
y = x ** 2
y.sum().backward()  # 第一次 backward 后中间值被释放
y.sum().backward()  # 报错：中间值已释放

# 方案 1：retain_graph=True
y.sum().backward(retain_graph=True)
y.sum().backward()  # 第二次可以成功

# 方案 2：使用 .detach() 而非 inplace 操作
x_new = x.detach() + 1  # 创建新张量，不影响原始 x

import torch
from torch.utils.checkpoint import checkpoint_sequential
from torch.autograd.graph import save_on_cpu

class LargeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.ModuleList([
            torch.nn.Linear(4096, 4096) for _ in range(48)
        ])

    def forward(self, x):
        # 策略 1：对整个模型使用 checkpoint_sequential
        # 分成 12 段，每段仅保存输入
        return checkpoint_sequential(self.layers, segments=12, input=x)

# 策略 2：结合 save_on_cpu 将中间值卸载到 CPU
model = LargeModel().cuda()
optimizer = torch.optim.Adam(model.parameters())

for batch in dataloader:
    input = batch.cuda()

    with save_on_cpu(pin_memory=True):
        output = model(input)

    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()