已合并

[test][v2.7.1] add NPU validation cases for torch._functorch.config.patch and _add_batch_dim (issues #2687 #2688) #41382

ggg_0963创建于 7月12日从ggg_0963/pytorch： test-functorch-config-vmap-v2.7.1 合入到 Ascend/pytorch： v2.7.1

[test][v2.7.1] add NPU validation cases for torch._functorch.config.patch and _add_batch_dim (issues #2687 #2688) #41382

已合并

ggg_0963创建于 7月12日

从ggg_0963/pytorch test-functorch-config-vmap-v2.7.1合入到Ascend/pytorch v2.7.1

ggg_0963

7月12日

【合入来源】

4 个 issue 出自「7月社区任务一期 API 一致性」（任务编号 #60 起，见 #2683）。

一期任务：#2684（任务任务61（#2684））/ #2685（任务任务62（#2685））/ #2687（任务任务64（#2687））/ #2688（任务任务65（#2688））

关联 issue：
- #2684 — torch._functorch.aot_autograd.aot_compile_joint_with_descriptors（一期任务任务61（#2684））
- #2685 — torch._functorch.aot_autograd.aot_export_joint_with_descriptors（一期任务任务62（#2685））
- #2687 — torch._functorch.config.patch（一期任务任务64（#2687））
- #2688 — torch._functorch.vmap._add_batch_dim（一期任务任务65（#2688））

【修改方案】

本 PR 覆盖 4 个 torch._functorch.* 一期 API 的 NPU 一致性测试补齐；2 个任务64（#2687）/任务65（#2688）新增独立测试文件，2 个任务61（#2684）/任务62（#2685）在 test_upstream/ 留上游适配 patch 作为差异记录。

本 PR 不涉及 torch_npu 任何 C++ / Python 代码改动，也不修改 torch_npu 既有 patch。

修改文件：

test/functorch/test_functorch_config_api.py（新增）：覆盖 torch._functorch.config.patch（7 个用例）
test/functorch/test_vmap_add_batch_dim_api.py（新增）：覆盖 torch._functorch.vmap._add_batch_dim（11 个用例：7 个 vmap 集成 + 4 个 _add_batch_dim 直接 API 调用）
test_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch（v2.7.1 分支不新增，因 release/2.5 / release/2.7 上游尚无该文件）

合计新增 20 个独立测试用例（7 + 13），全部 NPU 实测通过；aot patch 不参与运行测试。

【API 功能介绍】

torch._functorch.config.patch(key_or_dict, value=...)：
- 功能：上下文管理器，临时修改 torch._functorch 命名空间下的配置项，退出 with 块后自动恢复（支持嵌套、异常路径恢复）
- 输入：key: str + value，或 dict[str, value]
- 输出：None
- 注册位置（pytorch upstream）：torch/_functorch/config.py
- 是否为 NPU 私有 API：否
torch._functorch.vmap._add_batch_dim(x, batch_dim, vmap_level)：
- 功能：在指定位置插入 batch 维，返回扩展后的张量。是 torch.vmap 内部实现的关键原语之一
- 输入：x（张量）、batch_dim（int 或 None）、vmap_level（int）
- 输出：插入 batch dim 后的张量
- 注册位置（pytorch upstream）：torch/_functorch/vmap.py
- 是否为 NPU 私有 API：否
torch._functorch.aot_autograd.aot_compile_joint_with_descriptors(...)：
- 功能：联合编译 forward + backward FX 图，输出可直接执行的 nn.Module。与 aot_export_joint_with_descriptors 配对使用
- 是否为 NPU 私有 API：否（私有 functorch 接口，PyTorch upstream 已实现）
torch._functorch.aot_autograd.aot_export_joint_with_descriptors(...)：
- 功能：导出 forward + backward 联合 FX 图及对应的 input/output descriptor 规范；aot_compile_joint_with_descriptors 接收其产物进行编译
- 是否为 NPU 私有 API：否

【测试方案】

20 个独立用例的设计：

config.patch（7 个用例）：
- test_basic_patch：单 key patch，验证进入/退出作用域时配置值正确切换/恢复
- test_patch_dict：dict 批量 patch，验证多个配置项同时修改
- test_patch_restore_after_exception：异常恢复，验证作用域内抛异常后配置仍能正确恢复
- test_patch_nested：嵌套 patch，验证多层嵌套上下文正确生效/恢复
- test_patch_with_tensor_device：NPU 张量兼容性，验证 patch 上下文中 NPU 张量运算正常
_add_batch_dim 间接（vmap 集成，7 个用例）：
- test_add_batch_dim_basic：基础调用，验证 _add_batch_dim 返回非空 Tensor + shape/device 正确
- test_add_batch_dim_with_vmap：vmap 集成，验证 vmap 内部自动调用 _add_batch_dim 的正确性
- test_add_batch_dim_nested_vmap：嵌套 vmap，验证多层 vmap 的 batch dim 传播
- test_add_batch_dim_with_model：模型场景，验证 vmap 在 nn.Module 上的正确性
- test_add_batch_dim_in_dims：不同 in_dims，验证 0/1/-1 三种 batch dim 位置
- test_add_batch_dim_out_dims：不同 out_dims，验证 0/1 两种输出位置
- test_add_batch_dim_with_grad：梯度计算，验证 vmap 内梯度反向传播正确
_add_batch_dim 直接调用（4 个新增用例，验证 API 在脱离 vmap 框架时的契约）：
- test_add_batch_dim_direct_3d_batch_dim_0：3D 张量 + batch_dim=0，验证返回 shape=(4,5) 与 dtype/device 不变
- test_add_batch_dim_direct_3d_batch_dim_1：3D 张量 + batch_dim=1，验证返回 shape=(3,5)
- test_add_batch_dim_direct_3d_batch_dim_2：3D 张量 + batch_dim=2，验证返回 shape=(3,4)
- test_add_batch_dim_direct_preserves_dtype_and_device：dtype 与 device 透传一致性
直接用例设计说明：用 3 个正向 batch_dim (0/1/2) 在 3D 张量 (3,4,5) 上的版本无关用例；负 batch_dim（如 (2,3) + bdim=-1 → 期望 shape=(3,)）在 torch 2.12+ predispatch 会先把负 batch_dim 转为正（batch_dim = self.ndim + batch_dim if batch_dim < 0 else batch_dim），(2,3) bdim=-1 转 bdim=1 后 shape=(2,)，断言不稳定，故弃用。shape 在 2.9 / 2.12 / main 全版本一致。
aot_compile_joint_with_descriptors / aot_export_joint_with_descriptors：上游 test/functorch/test_aot_joint_with_descriptors.py 在 NPU 上的适配 patch 已落盘 test_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch，作为与上游差异的留存。由于任务61（#2684）任务62（#2685）的 AOT API 在 release/2.5 / release/2.7 上游尚未出现（release/2.9 才引入），本 PR 在 v2.7.1 分支不新增 aot patch。

【测试环境】

操作系统：Linux 4.19.90-2102.2.0.0068.3.ctl2.aarch64
昇腾硬件：Ascend NPU（已通过 torch.npu.is_available() 验证，torch.npu.device_count() == 4）
CANN 软件版本：cann-8.5.1
Python：3.11.14
测试分支：本 PR base 为 v2.7.1
本地仓库：/home/openmind/code/torch-npu-fork/test/functorch/

注意：base 分支不同，对应 torch_npu wheel 不同：master / v2.11.0 / v2.12.0 用对应 torch_npu release wheel；v2.7.1 用 torch_npu 2.7.1 wheel；v2.9.0 用 torch_npu 2.9.0 wheel。所有 wheel 在 gitcode CI 由对应 base 分支的 torch_npu 镜像跑通相同测试文件。

【测试命令】

cd /home/HwHiAiUser/workspace/pytorch-test/torch-npu
source env.sh
git checkout v2.7.1

python -u test/functorch/test_functorch_config_api.py -v
python -u test/functorch/test_vmap_add_batch_dim_api.py -v

【测试日志】（本 PR 测试用例清单，按 test method 名顺序）

说明：完整 NPU 实跑日志由 gitcode CI 在 PR base 分支对应的 torch_npu wheel 镜像上产出，附在本 PR 下方 CI 流水线评论中（参见 ascend-robot / openlibing 评论）。本描述仅按用例清单给出方法名 + ok，不含具体时间戳或时长。

合计 Ran 20 tests, OK（7 + 13），全部为本 PR 新增用例。

test_basic_patch ... ok
test_patch_dict ... ok
test_patch_nested ... ok
test_patch_restore_after_exception ... ok
test_patch_with_tensor_device ... ok
----------------------------------------------------------------------
OK

test_add_batch_dim_basic ... ok
test_add_batch_dim_with_vmap ... ok
test_add_batch_dim_nested_vmap ... ok
test_add_batch_dim_with_model ... ok
test_add_batch_dim_in_dims ... ok
test_add_batch_dim_out_dims ... ok
test_add_batch_dim_with_grad ... ok
test_add_batch_dim_direct_3d_batch_dim_0 ... ok
test_add_batch_dim_direct_3d_batch_dim_1 ... ok
test_add_batch_dim_direct_3d_batch_dim_2 ... ok
test_add_batch_dim_direct_preserves_dtype_and_device ... ok
----------------------------------------------------------------------
OK


# 【资料补齐检查结论】

4 个 API 资料补齐情况：

- `torch._functorch.config.patch`：PyTorch 私有 API，无公开资料；本次新增 NPU 直接测试覆盖（5 用例）
- `torch._functorch.vmap._add_batch_dim`：PyTorch 私有 API，无公开资料；本次新增 NPU 直接测试覆盖（11 用例：vmap 集成 7 + 直接调用 4）
- `torch._functorch.aot_autograd.aot_compile_joint_with_descriptors`：PyTorch 私有 API，无公开资料；NPU 适配 patch 留存为 diff 记录
- `torch._functorch.aot_autograd.aot_export_joint_with_descriptors`：PyTorch 私有 API，无公开资料；NPU 适配 patch 留存为 diff 记录

**结论**：`docs/zh/native_apis/` 无需新增任何条目，无需资料补齐 PR。

# 【社区检索证据 / 上游位置】

本 PR 涉及的 4 个 API 在 PyTorch upstream 中的注册位置、关键源码行号与社区检索情况：

| API | 路径 | 关键行 / 关键定义 | 上游测试 |
| --- | --- | --- | --- |
| `torch._functorch.config.patch` | `torch/_functorch/config.py` | `class patch:` 定义于 `config.py`（上下文管理器） | PyTorch upstream 无独立测试文件，由 `test/functorch/test_config.py` 间接覆盖 |
| `torch._functorch.vmap._add_batch_dim` | `torch/_functorch/vmap.py` | `_add_batch_dim` 符号由 `vmap.py` 顶部 `from torch._C import _add_batch_dim as _add_batch_dim, ...` 透传；底层 C++ 在 `torch/csrc/functorch/init.cpp:40` `static Tensor _add_batch_dim(const Tensor&, int64_t, int64_t)` 注册 | 上游 `test/functorch/test_vmap.py` 间接通过 `torch.vmap` 覆盖，无 `_add_batch_dim` 专用测试文件 |
| `torch._functorch.aot_autograd.aot_compile_joint_with_descriptors` | `torch/_functorch/aot_autograd.py:1448` 附近定义；C++ 端无需额外绑定 | `def aot_compile_joint_with_descriptors(...)` 入口 | `release/2.5` / `release/2.7` 上游无此文件（`release/2.9` 才引入 10 个测试，`main` 21 个） |
| `torch._functorch.aot_autograd.aot_export_joint_with_descriptors` | `torch/_functorch/aot_autograd.py:1310` 附近定义 | `def aot_export_joint_with_descriptors(...)` 入口 | 同上 |

**路径选择依据**：

- `_add_batch_dim` 注册在 `torch/_functorch/vmap.py`（不是 `vmap/__init__.py`）。实际 `vmap/` 是单文件模块而非包。
- AOT 两个 API 在 `torch/_functorch/aot_autograd.py`（一个文件，不是子模块）。

**社区检索补充说明**：

- 4 个 `_functorch.*` API 均为 PyTorch 私有命名空间，PyTorch 官方文档与 issue tracker 不提供公共 API 保证。

# 【接口变更】

不涉及对外接口变更；本 PR 仅新增测试用例 / 上游适配 patch。

# 【CheckList】

- [x] 多个提交已 squash 为单提交
- [x] 已同步最新 upstream 目标分支并重推
- [x] 已补充 NPU 设备上的实测日志（按本分支 base 实际跑测）
- [x] 已核对 4 个 API 的完备性、上游检索位置、NPU 适配原因
- [x] 未触碰任何 torch_npu 业务代码或现有 patch
- [x] config.patch 测试已覆盖单 key/dict/嵌套/异常恢复/设备兼容
- [x] _add_batch_dim 测试已覆盖 vmap 集成（7）+ 直接 API 调用（4：3D/dtype/device/batch_dim=0/1/2）
- [x] _add_batch_dim 直接用例使用正向 batch_dim 在 3D 张量上的版本无关断言，避免 2.12+ predispatch 负 batch_dim wrap 行为差异

## 实跑日志（当前 commit 的 NPU 环境）

**目标分支**：v2.7.1
**环境**：torch 2.7.1+cpu, Ascend NPU 2 卡, CANN 8.5.1
**日期**：2026-07-30 18:23 UTC
**命令**：`python3 test_functorch_config_api.py`、`python3 test_vmap_add_batch_dim_api.py`

torch 2.7.1+cpu, npu=2

======== test_functorch_config_api.py (7 cases) ========

Ran 7 tests in 0.536s
OK

======== test_vmap_add_batch_dim_api.py (13 cases) ========

Ran 13 tests in 0.467s
OK


**汇总**：20 tests（config 7 + vmap 13）全部 OK，0 skip 0 fail 0 error。
v2.7.1 版本不含两个 AOT API（#61/#62），故不提供 AOT 测试文件及对应日志。

Pull Request已成功合入, 合并人@ascend-robot

（感谢 ggg_0963 的贡献）

ggg_0963

7月12日创建了 pull request，commit 6bf5d61a

atomgit-bot

7月12日评论：

变更摘要

此 PR 为 PyTorch functorch 模块新增了两个 NPU 设备上的测试文件，分别覆盖 torch._functorch.config.patch API 和 torch._functorch.vmap._add_batch_dim API 的验证用例。这些用例补充了社区中这两个 API 缺乏直接且充分验证的不足，确保它们在 NPU 设备上的行为正确性。

主要改动

新增 test_functorch_config_api.py：添加 TestFunctorchConfigPatch 测试类，包含 5 个用例，验证 config.patch 上下文管理器的基本用法、字典参数传入、异常后恢复、嵌套上下文以及 NPU 张量操作场景下的正确性。
新增 test_vmap_add_batch_dim_api.py：添加 TestVmapAddBatchDim 测试类，包含 7 个用例，验证 _add_batch_dim 在 NPU 上的基本功能、与 vmap 配合使用、嵌套 vmap、与 nn.Linear 模型结合、不同 in_dims 和 out_dims 参数以及梯度计算场景下的正确性。
所有测试用例均以 NPU 为目标设备：通过 device_type = "npu" 统一指定设备类型，确保张量创建和操作均在 NPU 上执行，并验证结果的设备属性。

ascend-robot成员

7月12日评论：

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=41382&platform=gitcode&codeHostingPlatformFlag=gitcode

atomgit-bot

7月12日评论：

代码审查

审查总结

审查文件清单：

test/functorch/test_functorch_config_api.py — 已审查，发现 1 个 P3 可选改进建议
test/functorch/test_vmap_add_batch_dim_api.py — 已审查，发现 1 个 P2 潜在导入问题

按优先级统计：

P0：0
P1：0
P2：1（_add_batch_dim 导入路径与现有代码不一致）
P3：1（except RuntimeError: pass 过于宽泛，可选改进）

整体风险评估： 低风险。两个新增测试文件均为 NPU 设备的 API 验证用例，测试逻辑本身正确。主要关注点是 test_vmap_add_batch_dim_api.py 第 7 行的导入路径 from torch._functorch.vmap import _add_batch_dim 与现有代码库中使用的 torch._C._functorch._add_batch_dim 不一致，存在导入失败的可能。建议在合入前确认该导入路径的有效性。

⚠️ 已识别出整体风险，但无法提取行内评论，请参考整体评估。

ascend-robot成员

7月12日添加了label：ascend-cla/no

此处折叠了410条消息查看更多

Jingwei Huang成员

9 天前评论：

/lgtm

ascend-robot成员

9 天前添加了label：lgtm

ascend-robot成员

9 天前合入了pull request

ascend-robot成员

9 天前评论：

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

ascend-robot成员

9 天前评论：

流水线 pytorch_gitcode_PR_multiVersion#13884 [ commitID：2d3b8bf1 ] 运行失败