【合入来源】

[Usage]: torch.fx.node相关API NPU测试适配说明

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

【修改方案】

一、API功能说明

API范围：

• torch.fx.node._type_repr
• torch.fx.Node.all_input_nodes
• torch.fx.Node.append
• torch.fx.Node.args
• torch.fx.Node.format_node
• torch.fx.Node.insert_arg
• torch.fx.Node.is_impure
• torch.fx.Node.kwargs
• torch.fx.node.map_aggregate
• torch.fx.node.map_arg
• torch.fx.Node.next
• torch.fx.Node.normalized_arguments
• torch.fx.Node.prepend
• torch.fx.Node.prev
• torch.fx.Node.replace_all_uses_with

torch.fx.Node是torch.fx.Graph中的基本执行单元，以双向链表形式组织。上述15个API属于PyTorch torch.fx框架层的图节点操作与工具函数：

• torch.fx.node._type_repr是模块级私有函数，将Python类型转换为可读字符串，供图的类型注解生成使用；
• torch.fx.Node.all_input_nodes是args中所有Node类型元素的只读视图；
• torch.fx.Node.append/prepend调整节点在链表中的拓扑位置；
• torch.fx.Node.args/kwargs是节点的位置参数和关键字参数，赋值时自动维护use-def关系；
• torch.fx.Node.format_node返回节点的可读字符串表示；
• torch.fx.Node.insert_arg在指定位置插入新参数并更新use-def；
• torch.fx.Node.is_impure判断节点是否有副作用，是DCE等优化pass的判断依据；
• torch.fx.node.map_aggregate对任意嵌套结构（list/tuple/dict）递归应用函数；
• torch.fx.node.map_arg对节点args结构中每个Node递归应用函数；
• torch.fx.Node.next/prev是链表的后继/前驱指针属性；
• torch.fx.Node.normalized_arguments将节点args+kwargs按函数签名规范化为统一格式；
• torch.fx.Node.replace_all_uses_with将所有引用当前节点处替换为另一节点。
• 核心特性：上述15个API均为纯框架层图操作逻辑，操作对象是torch.fx.Graph的节点结构（Node链表、args/kwargs参数列表、use-def依赖关系等Python对象），不直接触发任何CPU/GPU/NPU硬件算子，不涉及Tensor数据的读写或设备相关计算，与计算设备完全解耦。其中torch.fx.node.map_aggregate是唯一一个以Tensor作为被操作数据的API（将函数递归应用于嵌套结构中的每个Tensor），其余14个API均在图结构层面工作，无论底层硬件为何均行为一致。
• NPU适配意义：torch.fx是PyTorch模型编译、图变换、量化、算子融合等核心优化流程的基础设施，在NPU环境下的模型编译与推理优化中被广泛调用。验证上述API在NPU环境下的行为一致性，可确保基于torch.fx构建的图变换pass、自定义优化器、模型导出工具在昇腾NPU上与CPU/GPU环境行为完全一致，避免因图操作层的隐性差异导致NPU上的模型编译失败或优化结果不符合预期。

二、测试用例完备性说明

（一）以下API在PyTorch官方test/test_fx.py中有对应测试函数，且用例全程操作的是图结构对象（Node/Graph），不创建也不执行任何Tensor计算，与设备完全无关，天然兼容NPU环境：

各测试函数无Tensor计算的具体依据：

• test_all_input_nodes：仅调用graph.placeholder/call_module/get_attr/call_function建立图结构，用assertEqual验证all_input_nodes列表内容，无Tensor创建。
• test_wrong_topo：通过nodes[3].append(nodes[2])制造拓扑错误，验证graph.lint()抛出RuntimeError，无Tensor创建。
• test_reassign_args_kwargs_uses：通过Proxy构建图后直接对node.args赋值，验证users字典同步更新，无Tensor创建或执行。注：该函数只操作node.args，未对node.kwargs做任何赋值或断言。
• test_pretty_print_node：对symbolic_trace的图调用format_node()，用FileCheck验证字符串内容；torch.rand(3,4)仅为Module初始化参数，symbolic_trace不执行真实计算。
• test_insert_arg：torch.tensor(0)仅作为torch.nn.Buffer初始值，全程只验证len(output_node.args)、a.users等图结构属性，无任何模型执行。
• test_prepend_does_not_leak：验证节点删除后weakref失效（内存释放），无Tensor计算。
• test_prepend_self：验证b.prepend(b)后图节点数不变，无Tensor计算。
• test_remove_uses / test_remove_uses_with_custom_filter：验证replace_all_uses_with后users字典变化，无Tensor计算。

（二）PyTorch官方有用例且需要NPU适配的API

• torch.fx.node.map_aggregate
  用例文件：test/dynamo/test_functions.py
  测试函数：test_fx_map_aggregate（直接覆盖）
  用例中创建torch.randn(4)并执行y * 2乘法计算，同时通过torch.compile编译执行，Tensor是核心被测对象，需迁移到NPU。

以下2个API在官方测试中有Tensor计算，但经实际验证确认无需NPU适配：

• torch.fx.Node.is_impure
  用例文件：test/fx/test_dce_pass.py
  测试函数：test_impure_nodes_args、test_impure_kwargs、test_impure_custom
  is_impure()是纯Python属性判断（判断节点是否有副作用），与设备无关。测试中的m(*inputs)是验证DCE变换前后模型输出一致的兜底验证，不是验证is_impure本身的NPU行为。
  实际运行原始社区用例（CPU），三个测试函数全部通过，无需NPU适配。
  运行命令：
  python -m unittest
  fx.test_dce_pass.TestDCE.test_impure_nodes_args
  fx.test_dce_pass.TestDCE.test_impure_kwargs
  fx.test_dce_pass.TestDCE.test_impure_custom -v
  运行结果：Ran 3 tests / OK

• torch.fx.Node.normalized_arguments
  用例文件：test/test_fx_experimental.py
  测试函数：test_normalize_modules_exhaustive
  normalized_arguments将节点args+kwargs按函数签名规范化，为纯Python图操作，与设备无关。测试中的traced(*inputs)/mod(*inputs)是归一化变换后的兜底验证，不是验证normalized_arguments的NPU行为。尝试迁移到NPU时，因穷举所有torch.nn模块触发了两类无关报错：
  （1）PadV3算子库未加载（含padding的卷积模块）
  （2）complex128 dtype不支持aclnnIsClose（复数类模块）
  两个报错均与torch.fx.Node.normalized_arguments无关，证明Tensor是配角。
  社区原有CPU测试已充分覆盖该API全部功能，无需NPU适配。

（三）PyTorch官方无用例、需新增用例的API

以下API在PyTorch官方测试目录中无直接测试函数，经全量搜索确认：

• torch.fx.node._type_repr
• torch.fx.Node.kwargs
• torch.fx.node.map_arg
• torch.fx.Node.next
• torch.fx.Node.prev

新增用例文件：test/fx/test_fx_node_api.py
覆盖上述API的核心功能，已测试验证通过。

三、API适配方案

1. API源码无需修改：上述涉及API均为PyTorch框架层图操作，不涉及NPU kernel开发或算子注册。

2. 情况（一）涉及API：社区原有用例天然兼容NPU环境。

3. 情况（二）涉及API（torch.fx.node.map_aggregate）：
   在test_upstream/test/dynamo/test_functions.py.patch中，调用npu()，采用最小化原则，不影响文件内其他测试函数。

4. 情况（三）涉及API：新增test/fx/test_fx_node_api.py，涵盖_type_repr、Node.kwargs（getter/setter/use-def更新）、map_arg（节点收集/非节点穿透/嵌套结构/callable校验）、Node.next/prev（基本顺序/append-prepend后的顺序变化/next-prev一致性）。上述API均为纯图结构/字符串操作，不涉及Tensor，无需设备适配。

【资料变更】

已检查API支持清单。变更如下：

（1）以下涉及API在文档中已有记录但标注为"否"，本次更正为"是"：
torch.fx.Node.all_input_nodes、torch.fx.Node.args、torch.fx.Node.kwargs、
torch.fx.Node.next、torch.fx.Node.prev

（2）以下涉及API在文档中完全缺失，本次新增记录（"是"）：
torch.fx.node._type_repr、torch.fx.Node.insert_arg、
torch.fx.node.map_arg、torch.fx.node.map_aggregate

【接口变更】

不涉及

【功能验证】

验证环境如下
操作系统：Ubuntu
昇腾硬件：800I A2
CANN软件版本：8.5.0
2.9.0：torch 2.9.0+cpu / torch-npu 2.9.0

运行命令与结果：

(test29) [root@4e6c21e2fea3 ascend-pytorch-work]# pip list | grep torch
torch                  2.9.0+cpu
torch_npu              2.9.0
(test29) [root@4e6c21e2fea3 ascend-pytorch-work]# python official-pytorch-v2.9.0/test/dynamo/test_functions.py -k test_fx_map_aggregate
inline_call []
stats [('calls_captured', 6), ('unique_graphs', 2)]
.
----------------------------------------------------------------------
Ran 1 test in 1.785s

OK
(test29) [root@4e6c21e2fea3 test]# python fx/test_fx_node_api.py
...................
----------------------------------------------------------------------
Ran 19 tests in 0.046s

OK

【CheckList】

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