【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单
• 重构优化
• 资料更新

关联 issue：https://gitcode.com/Ascend/pytorch/issues/1609

【修改方案】

请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列
如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容）

本 PR 针对 torch.fx.experimental.proxy_tensor 相关 API 新增专项验证用例，覆盖如下 API：

• torch.fx.experimental.proxy_tensor.get_proxy_mode
• torch.fx.experimental.proxy_tensor.handle_sym_dispatch
• torch.fx.experimental.proxy_tensor.make_fx
• torch.fx.experimental.proxy_tensor.maybe_disable_thunkify
• torch.fx.experimental.proxy_tensor.maybe_enable_thunkify

一、API 功能说明

1. torch.fx.experimental.proxy_tensor.get_proxy_mode

   该 API 用于获取当前活跃的 ProxyTorchDispatchMode。在普通执行环境下返回 None，在 make_fx tracing 过程中可获取到当前 proxy mode。

2. torch.fx.experimental.proxy_tensor.handle_sym_dispatch

   该 API 用于处理 proxy tensor tracing 过程中的 symbolic dispatch 逻辑，依赖当前 proxy mode 环境。该接口属于 proxy_tensor tracing 内部调度相关 helper。

3. torch.fx.experimental.proxy_tensor.make_fx

   该 API 用于对 Python 函数进行 tracing，并生成对应的 torch.fx.GraphModule。生成后的 GraphModule 可用于表示被 trace 函数的 ATen 图结构。

4. torch.fx.experimental.proxy_tensor.maybe_enable_thunkify

   该 API 用于在特定上下文中启用 thunkify 相关逻辑，可作为上下文管理器使用。

5. torch.fx.experimental.proxy_tensor.maybe_disable_thunkify

   该 API 用于在特定上下文中关闭 thunkify 相关逻辑，可作为上下文管理器使用。

二、上游社区用例分析与用例补齐说明

经检索，PyTorch 上游社区中已存在部分 make_fx 相关测试用例。例如在 test/export/test_export.py 中存在 test_stack_trace_make_fx 等用例，会调用 torch.fx.experimental.proxy_tensor.make_fx。

但经分析，现有 make_fx 用例主要验证的是 export / stack trace 等上层功能路径中 make_fx 的使用效果，测试重点并不是 make_fx API 本身的基础行为，也未集中覆盖本次 issue 要求的 torch.fx.experimental.proxy_tensor 相关 API 组合。具体来说，现有用例未充分覆盖以下内容：

1. 未独立验证 make_fx 返回对象是否为 torch.fx.GraphModule；
2. 未独立验证 make_fx 生成的 GraphModule 执行结果是否与原函数一致；
3. 未验证 get_proxy_mode 在 tracing 外返回 None、在 make_fx tracing 中可获取到 proxy mode 的行为；
4. 未验证 handle_sym_dispatch 作为 proxy tensor tracing 内部 helper 对 proxy mode 的依赖关系；
5. 未验证 maybe_enable_thunkify / maybe_disable_thunkify 作为上下文管理器在 make_fx tracing 中可稳定执行；
6. 未针对 Torch-NPU 场景验证测试张量迁移到当前 accelerator / NPU 侧运行。

同时，经检索，本次任务中的以下 API 未发现独立、直接的专项验证用例：

• torch.fx.experimental.proxy_tensor.get_proxy_mode
• torch.fx.experimental.proxy_tensor.handle_sym_dispatch
• torch.fx.experimental.proxy_tensor.maybe_disable_thunkify
• torch.fx.experimental.proxy_tensor.maybe_enable_thunkify

因此，现有上游社区用例对 make_fx 有调用覆盖，但不足以完整覆盖本次 issue 中 5 个 proxy_tensor API 的专项兼容性验证需求。基于此，本 PR 新增轻量级专项测试文件：

test/fx/test_proxy_tensor_api.py

该文件集中验证 torch.fx.experimental.proxy_tensor 相关 API 的基础行为，且测试张量通过 .to(device_type) 迁移到当前 accelerator / NPU 侧运行，满足 Torch-NPU API 一致性测试要求。

三、新增测试文件说明

本 PR 新增测试文件：

test/fx/test_proxy_tensor_api.py

该测试文件用于集中验证 torch.fx.experimental.proxy_tensor 相关 API 的基础兼容性。为便于后续维护和扩展，文件开头已补充英文说明，说明该文件用于补充 proxy_tensor API 的 NPU 侧验证用例，并可继续扩展同类 API 测试。

文件开头说明如下：

"""
Add validation cases for torch.fx.experimental.proxy_tensor APIs on NPU:
1. PyTorch community lacks sufficient and direct API validations for
   some proxy_tensor APIs, so this file is added.
2. This file validates get_proxy_mode, handle_sym_dispatch, make_fx,
   maybe_enable_thunkify, and maybe_disable_thunkify (extendable).
"""

新增测试覆盖以下场景：

1. make_fx tracing 基础场景

   通过 make_fx(fn)(x, y) 对函数进行 tracing，验证返回结果为 GraphModule，并验证生成的 GraphModule 执行结果与原函数一致。

2. get_proxy_mode tracing 场景

   验证 get_proxy_mode 在 tracing 外返回 None，在 make_fx tracing 过程中可以获取到非空 proxy mode。

3. handle_sym_dispatch proxy mode 依赖场景

   验证 handle_sym_dispatch 为可调用对象，并验证其依赖当前 proxy mode 的行为。在非 proxy mode 环境下直接调用会触发断言，符合该 helper 的使用约束。

4. maybe_enable_thunkify / maybe_disable_thunkify 上下文管理器场景

   验证 maybe_enable_thunkify 和 maybe_disable_thunkify 可作为上下文管理器正常使用，并可在 make_fx tracing 过程中稳定执行。

四、NPU 适配说明

上述 API 均属于 torch.fx.experimental.proxy_tensor Python 层 tracing / proxy tensor 相关接口，不涉及新增 NPU 算子实现，也不涉及 C++、Kernel 或底层算子逻辑修改。

但测试用例中涉及张量输入，因此测试张量需要运行在 NPU / 当前 accelerator 侧。本 PR 使用如下方式获取当前 accelerator：

device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"

并通过如下方式将测试张量迁移到当前设备：

x = torch.randn(2, 3).to(device_type)
y = torch.randn(2, 3).to(device_type)

在 Torch-NPU 测试环境中，当前 accelerator 为 NPU，因此测试张量会通过 .to(device_type) 迁移到 NPU 侧运行，满足新增用例涉及张量需在 NPU 侧执行的要求。

五、涉及版本说明

本次用例补齐分别提交到以下目标分支：

• v2.7.1
• v2.9.0
• v2.10.0
• v2.11.0
• v2.12.0
• master

各版本均单独提交 PR，目标分支与源码分支一一对应。

本次根据当前任务要求提交上述分支，不涉及 v2.8.0 分支。

【资料变更】

请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写“不涉及”

不涉及。

本 PR 仅新增测试用例，不修改资料文档。
相关 API 的 docs/zh/native_apis 支持状态已通过单独文档 PR 补充。
本 PR 仅新增 test/fx/test_proxy_tensor_api.py 测试用例，不直接修改 docs/zh/native_apis 资料文档。
但已结合配套资料补齐任务，对本次涉及的 torch.fx.experimental.proxy_tensor 相关 API 在 docs/zh/native_apis 中的资料支持状态进行了同步核查。
资料支持情况核查结论如下：

1. 本次任务涉及的 API 均属于 torch.fx.experimental.proxy_tensor 下的 Python 层 tracing / proxy tensor 相关接口；
2. 相关 API 的资料支持状态需要结合目标版本实际能力在 docs/zh/native_apis/pytorch_*/torch-fx.md 中进行声明；
3. 本测试 PR 仅负责补齐 test 目录下的专项用例与行为验证，不涉及 torch_npu/ 目录下的功能代码修改，也不涉及资料文档修改；
4. 相关 API 的 docs/zh/native_apis 支持状态已通过单独资料 PR 补充；
5. 本 PR 的测试验证范围与当前代码实际可验证能力保持一致，避免“资料声明”和“测试验证”口径

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写“不涉及”

不涉及。

本 PR 不修改客户可见接口，不修改 API 行为，仅新增测试用例验证现有 API 兼容性。

【功能验证】

说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤
新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图

新增测试用例：

test/fx/test_proxy_tensor_api.py

测试命令：

python -u test/fx/test_proxy_tensor_api.py -v

测试场景：

1. 验证 make_fx 可正常生成 GraphModule；
2. 验证 make_fx 生成的 GraphModule 执行结果与原函数一致；
3. 验证 get_proxy_mode 在 tracing 外返回 None；
4. 验证 get_proxy_mode 在 make_fx tracing 中可获取到 proxy mode；
5. 验证 handle_sym_dispatch 的可调用性及其依赖 proxy mode 的行为；
6. 验证 maybe_enable_thunkify / maybe_disable_thunkify 上下文管理器可正常使用；
7. 验证测试中涉及的张量已通过 .to(device_type) 迁移到当前 accelerator / NPU 侧运行。

本地验证结果示例：

test_get_proxy_mode_during_make_fx (__main__.TestProxyTensorAPI.test_get_proxy_mode_during_make_fx) ... ok
test_handle_sym_dispatch_requires_proxy_mode (__main__.TestProxyTensorAPI.test_handle_sym_dispatch_requires_proxy_mode) ... ok
test_make_fx_returns_graph_module (__main__.TestProxyTensorAPI.test_make_fx_returns_graph_module) ... ok
test_thunkify_context_managers (__main__.TestProxyTensorAPI.test_thunkify_context_managers) ... ok

----------------------------------------------------------------------
Ran 4 tests in 0.333s

OK

已在以下版本执行该用例：

• v2.7.1：通过
• v2.9.0：通过
• v2.10.0：通过
• v2.11.0：通过
• v2.12.0：通过
• master：通过

说明：

该组 API 为 torch.fx.experimental.proxy_tensor Python 层 tracing 相关接口，不涉及 NPU 算子计算。本次测试主要验证 API 功能行为与 PyTorch 原生保持一致，同时确保测试中涉及的张量运行在当前 accelerator / NPU 侧。

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常