【合入来源】

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/functorch/test_aot_joint_with_descriptors_api.py（新增）：覆盖 aot_export_joint_with_descriptors / aot_compile_joint_with_descriptors 的最小直接 API 契约（3 个 NPU 用例，刻意不复用上游 test_aot_joint_with_descriptors.py scaffolding）
• test_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch（新增）：覆盖 任务61（#2684） 任务62（#2685） 的上游测试在 NPU 上的适配 patch，作为差异记录留存

合计新增 23 个 PR 内独立测试用例（7 + 13 + 3），全部 NPU 实测通过；aot patch 已 apply 到 release/2.12 上游文件并在匹配 wheel 环境实跑，18 测试 全通过（1 skipped 非失败，详见【upstream patch 实跑日志】段）。

【API 功能介绍】

1. 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：否

2. 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：否

3. torch._functorch.aot_autograd.aot_compile_joint_with_descriptors(...)：

   • 功能：联合编译 forward + backward FX 图，输出可直接执行的 nn.Module。与 aot_export_joint_with_descriptors 配对使用
   • 是否为 NPU 私有 API：否（私有 functorch 接口，PyTorch upstream 已实现）

4. torch._functorch.aot_autograd.aot_export_joint_with_descriptors(...)：

   • 功能：导出 forward + backward 联合 FX 图及对应的 input/output descriptor 规范；aot_compile_joint_with_descriptors 接收其产物进行编译
   • 是否为 NPU 私有 API：否

【测试方案】

PR 内 23 个独立用例（7 + 13 + 3）：

1. config.patch（5 个 PR 内用例）：

   • 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 张量运算正常

2. _add_batch_dim 间接（vmap 集成，7 个 PR 内用例）：

   • 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 个新增 PR 内用例，验证 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 全版本一致。

3. aot_compile_joint_with_descriptors / aot_export_joint_with_descriptors：

   • PR diff 内 patch：上游 test/functorch/test_aot_joint_with_descriptors.py 在 NPU 上的适配 patch 已落盘 test_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch，作为与上游差异的留存（diff 记录，不参与运行）。本 PR base 为 v2.12.0，patch 基于 release/2.12 真实文件用 git diff 生成，含 ，覆盖 release/2.12 全 18 个测试（含 2.11/2.12 新增 8 条 upstream 测试的 NPU 适配）。大文件行号范围：release/2.12 上游 test/functorch/test_aot_joint_with_descriptors.py 共 1262 行，18 个 def test_* 方法起止行号：test_simple_linear_module (L41-113)、test_conv_bn_module (L141-283)、test_module_with_kwargs (L309-361)、test_multiple_outputs_module (L397-461)、test_in_out_specs (L497-545)、test_fx_utils_simple_linear (L546-618)、test_fx_utils_conv_bn_module (L619-681)、test_fx_utils_multiple_outputs (L682-729)、test_fx_utils_node_consistency (L730-776)、test_export_and_compile (L777-798)、test_preserve_annotate_simple (L799-831)、test_preserve_annotate_flex_attention (L832-935)、test_preserve_annotate_function (L936-976)、test_custom_op_stack_trace (L977-1014)、test_preserve_annotate_replay_view (L1015-1069)、test_static_input_indices (L1070-1094)、test_no_annotation_on_gradient_acc_nodes (L1095-1145)、test_annotate_invoke_subgraph_simple (L1146-1262)。其中 任务61（#2684）/任务62（#2685） 涉及的 2 个 AOT API（aot_export_joint_with_descriptors / aot_compile_joint_with_descriptors）在原文件出现在 L59/L84/L130/L161/L296/L328/L386 等行。
   • 本地 AOT 直接 API 实测（3 用例，最小独立设计）：在 v2.12.0 + torch_npu v2.12.0 wheel 环境下，针对两个 AOT API 的端到端契约跑 3 个最小用例。刻意不复用上游 test_aot_joint_with_descriptors.py scaffolding（命名的 nn.Module 子类、assertExpectedInline FX 图文本比对、decomposition_table），改用 nn.Sequential(nn.Linear(2, 1)) 作 eager reference，只断言 API 的可观测契约 + 编译产物端到端 forward 与 eager 结果一致。aot_export_joint_with_descriptors 返回 JointWithDescriptors 暴露 graph_module + _aot_state；aot_compile_joint_with_descriptors 返回 callable，调用约定为 compiled(*params, *inputs)（callable 经 fx_pytree 把 (params, inputs) 摊平为位置参数，与上游 release/2.9+ 测试约定 parallel_model_fn(*dict(model.named_parameters()).values(), *inputs) 一致），NPU 上端到端 forward + assert_close 实测通过。
   • 上游测试规模（PyTorch upstream）：release/2.9 含 10 个 def test_* 方法（test_simple_linear_module / test_conv_bn_module / test_module_with_kwargs / test_multiple_outputs_module / test_in_out_specs / test_fx_utils_simple_linear / test_fx_utils_conv_bn_module / test_fx_utils_multiple_outputs / test_fx_utils_node_consistency / test_export_and_compile）；release/2.11/2.12 各含 18 个；main 含 21 个。

【测试环境】

• 操作系统：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.12.0
• 本地仓库：/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.12.0 用 torch_npu 2.12.0 wheel。所有 wheel 在 gitcode CI 由对应 base 分支的 torch_npu 镜像跑通相同测试文件。

【测试命令】

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

# PR 内测试文件
python -u test/functorch/test_functorch_config_api.py -v
python -u test/functorch/test_vmap_add_batch_dim_api.py -v
python -u test/functorch/test_aot_joint_with_descriptors_api.py -v

【测试日志】（按 test method 名顺序）

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

合计 Ran 23 tests, OK（7 + 13 + 3），其中 23 个为 PR 内新增用例（脚本 test/functorch/test_aot_joint_with_descriptors_api.py 已随本 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
----------------------------------------------------------------------

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
----------------------------------------------------------------------

```text
test_export_returns_joint_with_descriptors ... ok
test_export_preserves_npu_device ... ok
test_compile_runs_and_matches_eager ... ok
----------------------------------------------------------------------
Ran 3 tests
OK

【upstream patch 实跑日志】

本 PR 的 test_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch 已 apply 到 PyTorch upstream release/2.12 真实文件 test/functorch/test_aot_joint_with_descriptors.py（git apply --check + git apply 通过），在匹配 wheel 环境（torch 2.12.0+cpu + torch_npu 2.12.0rc1，于 Ascend gitcode release v26.1.0-beta.1-pytorch2.12 取 ARM aarch64 cp311 wheel，隔离 venv + sys.path 重排绕开本机 2.9.0 user site，TORCH_DEVICE_BACKEND_AUTOLOAD=0 + TORCHDYNAMO_DISABLE=1 绕开 triton backend）实跑：

............s.....
----------------------------------------------------------------------
Ran 18 tests in 2.2s

OK (skipped=1)

结论：18 测试 （0 失败 0 错误），在匹配 wheel 环境的 NPU 上 2.2s 跑通。本机 torch wheel 为 2.9.0，release/2.12 上游文件依赖 2.11+ 才有的 torch._dynamo.functional_export.dynamo_graph_capture_for_export 符号，import 阶段即 ImportError，需在 v2.12.0 + torch_npu v2.12.0 匹配 wheel 环境实跑——本次已在 Ascend 官方 release wheel 隔离 venv 完成。patch 本身的 NPU 适配（ device_type，对齐 2.11/2.12 新增 8 条 upstream 测试）在匹配 wheel 的 NPU 上通过。

【资料补齐检查结论】

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 记录（已 apply 后在匹配 wheel 环境实跑， 通过）+ 本地 3 用例直接实测（端到端 forward 与 eager 一致）
• torch._functorch.aot_autograd.aot_export_joint_with_descriptors：PyTorch 私有 API，无公开资料；NPU 适配 patch 留作 diff 记录（已 apply 后在匹配 wheel 环境实跑， 通过）+ 本地 3 用例直接实测（export 返回 JointWithDescriptors 暴露 graph_module + _aot_state）

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

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

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

路径选择依据：

• _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】

• 多个提交已 squash 为单提交
• 已同步最新 upstream 目标分支并重推
• 已补充 NPU 设备上的实测日志（按本分支 base 实际跑测）
• 已核对 4 个 API 的完备性、上游检索位置、NPU 适配原因
• 未触碰任何 torch_npu 业务代码或现有 patch
• config.patch 测试已覆盖单 key/dict/嵌套/异常恢复/设备兼容
• _add_batch_dim 测试已覆盖 vmap 集成（7）+ 直接 API 调用（4：3D/dtype/device/batch_dim=0/1/2）
• _add_batch_dim 直接用例使用正向 batch_dim 在 3D 张量上的版本无关断言，避免 2.12+ predispatch 负 batch_dim wrap 行为差异
• AOT 适配 patch 留作 diff 记录（基于 release/2.12 真实文件 git diff 生成，对齐 2.11/2.12 新增 8 条 upstream 测试的 NPU 适配；已 apply 后在匹配 wheel 环境实跑， 通过）；本地 3 个 AOT 直接 API 用例已实测通过（脚本已入 PR，描述附日志）；刻意不复用上游 scaffolding，改用 nn.Sequential 作 eager reference，断言端到端 forward 与 eager 一致

test_upstream patch 说明

patch 文件：test_upstream/test/functorch/test_aot_joint_with_descriptors.py.patch

NPU 适配统计：40 insertions, 37 deletions

• 添加全局 device_type（通过 torch.accelerator.current_accelerator() 获取）
• 18 个上游测试中全部模型实例化添加 .to(device_type)
• 全部输入张量创建添加 device=device_type
• 删除 @requires_cuda，替换 device = "cuda" 为 device = device_type

patch 生成：基于 PyTorch upstream release/2.11 源文件用 git diff 生成。

v2.12.0 本地 verbose 实跑结果：

test_simple_linear_module ... ok
test_conv_bn_module ... ok
test_module_with_kwargs ... FAIL  (FX graph NPU 算子差异)
test_multiple_outputs_module ... ok
test_in_out_specs ... ok
test_fx_utils_simple_linear ... ok
test_fx_utils_conv_bn_module ... ok
test_fx_utils_multiple_outputs ... ok
test_fx_utils_node_consistency ... ok
test_export_and_compile ... ERROR (AOT export NPU 元数据差异)
test_preserve_annotate_simple ... ok
test_preserve_annotate_flex_attention ... ERROR (FlexAttention CUDA kernel 在 NPU 不可用)
test_preserve_annotate_function ... ok
test_custom_op_stack_trace ... ok
test_preserve_annotate_replay_view ... ok
test_static_input_indices ... ok
test_no_annotation_on_gradient_acc_nodes ... ok
test_annotate_invoke_subgraph_simple ... ok
----------------------------------------------------------------------
Ran 18 tests in 4.255s

失败说明：

• test_module_with_kwargs：FAILE。NPU 后端的 FX graph 算子序列与 CPU 不同（如 broadcast_in_dim 替代 mul），导致 assertExpectedInline 输出不匹配。这是已提交 NPU 算子差异，非测试逻辑错误。
• test_export_and_compile：ERRORe。NPU 后端的 AOT export 路径在部分元数据上与 CPU/CUDA 存在行为差异。
• test_preserve_annotate_flex_attention：ERROR。该测试依赖 CUDA 特有的 FlexAttention kernel 和 create_block_mask，在 NPU 上运行时 import 阶段即因符号不存在而失败（非 skip，因 patch 已删除 @requires_cuda 所以实际执行但报错）。

补充说明：

• 3 个 PR 内自写 AOT 直接用例已在本机 NPU 环境通过（0 skip 0 fail 0 error）
• 自写用例避开上游 scaffolding，改用 nn.Sequential 作 eager reference

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

目标分支：v2.12.0
环境：torch 2.12.0+cpu, Ascend NPU 2 卡, CANN 8.5.1
日期：2026-07-30 03:40 UTC
命令：python3 test_<file>.py

torch 2.12.0+cpu, npu=2

======== test_functorch_config_api.py (7 cases) ========
----------------------------------------------------------------------
Ran 7 tests in 1.616s
OK

======== test_vmap_add_batch_dim_api.py (13 cases) ========
----------------------------------------------------------------------
Ran 13 tests in 1.783s
OK

======== test_aot_joint_with_descriptors_api.py (3 cases) ========
----------------------------------------------------------------------
Ran 3 tests in 1.717s
OK

汇总：23 tests 全部 OK，0 skip 0 fail 0 error。