【合入来源】

issue / 工单

Torch-NPU ShapeEnv API 补齐任务。

PyTorch v2.10.0 官方社区在 test/test_proxy_tensor.py 中已有相关测试用例，覆盖以下 ShapeEnv API：

• torch.fx.experimental.symbolic_shapes.ShapeEnv.produce_guards
• torch.fx.experimental.symbolic_shapes.ShapeEnv.get_nontrivial_guards

Torch-NPU 当前 v2.10.0 分支下 test_upstream/test/ 目录中没有对应的 test_proxy_tensor.py.patch，因此本 PR 新增 upstream patch 文件，将 PyTorch 官方社区已有用例适配到 Torch-NPU 的 NPU 测试场景。

本 PR 属于：

1.1 pytorch 官方社区有用例，并且需要做 NPU 适配
1.1.1 torch-npu 社区 test_upstream 无 patch

【修改方案】

本 PR 新增文件：

test_upstream/test/test_proxy_tensor.py.patch

具体修改内容如下：

1. 在 PyTorch 官方测试文件 test/test_proxy_tensor.py 中补充 import torch_npu。
2. 将 _trace helper 中构造的输入 Tensor 从 CPU Tensor 改为 NPU Tensor：
   • torch.randn(arg) -> torch.randn(arg).npu()
3. 将 TestSymbolicTracing.test_mega_guard 中的输入 Tensor 改为 NPU Tensor：
   • torch.randn(16) -> torch.randn(16).npu()
   • torch.randn(8) -> torch.randn(8).npu()
4. 通过官方已有测试覆盖 ShapeEnv.produce_guards 和 ShapeEnv.get_nontrivial_guards 在 NPU Tensor 输入场景下的行为。

一、API 功能

1. torch.fx.experimental.symbolic_shapes.ShapeEnv.produce_guards

该 API 用于根据 ShapeEnv 中记录的符号形状约束生成 guards 表达式，用于保证后续运行时输入形状满足 tracing 阶段建立的符号约束。

在 symbolic tracing 场景下，当输入 Tensor 的 shape 被符号化后，produce_guards 会根据符号变量、输入来源以及约束关系生成运行时检查条件。

2. torch.fx.experimental.symbolic_shapes.ShapeEnv.get_nontrivial_guards

该 API 用于获取 ShapeEnv 中非平凡的 guards，用于判断当前符号形状环境中是否存在需要保留的动态形状约束。

在部分 shape 等价或可化简场景下，如果所有 guards 都可以被消除或证明为平凡条件，则 get_nontrivial_guards 应返回空结果。

二、官方测试用例完整性分析

PyTorch v2.10.0 官方测试文件：

test/test_proxy_tensor.py

已经包含对上述 API 的测试覆盖。

1. ShapeEnv.produce_guards

官方用例：

TestSymbolicTracing.test_mega_guard

该用例通过 symbolic tracing 构造动态 shape 约束：

assert a.shape[0] == b.shape[0] * 2

随后调用：

fx_g.shape_env.produce_guards(...)

并通过 assertExpectedInline 验证生成的 guards 内容符合预期。

因此，该用例可以覆盖 ShapeEnv.produce_guards 的核心功能，即根据 symbolic shape 约束生成运行时 guards。

2. ShapeEnv.get_nontrivial_guards

官方 helper：

_assert_no_guards

其中调用：

fx_g.shape_env.get_nontrivial_guards()

并断言：

len(fx_g.shape_env.get_nontrivial_guards()) == 0

官方用例：

TestSymbolicTracing.test_guards_equal

该用例通过多个 symbolic tracing 场景调用 _assert_no_guards，包括：

1. 多个输入 Tensor shape 完全相等的场景。
2. 广播相关 shape 场景。
3. torch.cat 组合 shape 场景。
4. view / sum 等 shape 传播场景。

这些场景验证了在 shape 约束可以被化简或无需保留非平凡 guards 时，ShapeEnv.get_nontrivial_guards 能正确返回空 guards。

3. 完整性结论

综上，PyTorch v2.10.0 官方用例已经覆盖：

API	官方测试用例
ShapeEnv.produce_guards	TestSymbolicTracing.test_mega_guard
ShapeEnv.get_nontrivial_guards	_assert_no_guards / TestSymbolicTracing.test_guards_equal

因此，本 PR 不重复新增自写测试，而是按照 Torch-NPU 的 test_upstream 机制，对 PyTorch 官方已有用例进行 NPU 适配。

三、NPU 适配

3.1 API 适配

不需要修改 API 实现。

ShapeEnv.produce_guards 和 ShapeEnv.get_nontrivial_guards 属于 symbolic shape / guard 管理相关的 Python 层逻辑，主要负责符号形状约束生成、查询和化简。

这两个 API 不涉及：

1. NPU kernel 实现。
2. NPU 显存管理。
3. NPU stream / event。
4. 分布式通信算子。
5. 设备侧数值计算逻辑。

因此，本次不需要修改 API 源码实现。

3.2 测试用例适配

需要适配测试输入 Tensor 的设备。

PyTorch 官方用例默认使用 CPU Tensor，例如：

torch.randn(arg)
torch.randn(16)
torch.randn(8)

为了在 Torch-NPU 场景下验证 symbolic tracing 和 ShapeEnv guards 行为，本 PR 将相关输入 Tensor 改为 NPU Tensor：

torch.randn(arg).npu()
torch.randn(16).npu()
torch.randn(8).npu()

同时补充：

import torch_npu

以确保 NPU 后端可用。

【资料变更】

不涉及。

本 PR 仅新增 upstream 测试适配 patch，不修改文档资料。

【接口变更】

不涉及。

本 PR 不新增、不删除、不修改 Python API 或 C++ 接口，仅新增测试适配 patch。

【功能验证】

1. patch dry-run 验证

已执行以下命令验证 patch 可以正确应用到 PyTorch v2.10.0 官方测试文件：

patch --dry-run -p1 -d /workspace/user_data/pytorch_official_v2.10.0 < test_upstream/test/test_proxy_tensor.py.patch

验证结果：

checking file test/test_proxy_tensor.py
Hunk #2 succeeded at 970 (offset 3 lines).
Hunk #3 succeeded at 1833 (offset 1 line).

验证结论：

patch dry-run 执行成功，未出现 hunk failed，说明 test_upstream/test/test_proxy_tensor.py.patch 可以正确应用到 PyTorch v2.10.0 官方 test/test_proxy_tensor.py 文件。

2. 测试执行说明

当前 Web IDE 环境为 Torch-NPU 源码目录环境，未加载完整编译后的 torch_npu._C 扩展。因此本次主要完成 patch 生成与 dry-run 验证。

后续 PR 创建后，可通过社区 CI / 门禁环境执行完整 upstream 测试流程。

【关联 PR / Issue】

v2.10.0 自写测试 PR：

https://gitcode.com/Ascend/pytorch/pull/36393

v2.9.0 同类 upstream patch PR：

https://gitcode.com/Ascend/pytorch/pull/36369

v2.9.0 ignore_fresh_unbacked_symbols 无需 NPU 适配说明 issue：

https://gitcode.com/Ascend/pytorch/issues/2046