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

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

Torch-NPU ShapeEnv API 补齐任务。

PyTorch v2.11.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.11.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

【修改方案】

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

本 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.11.0 官方测试文件：

test/test_proxy_tensor.py

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

1. ShapeEnv.produce_guards

官方用例：

TestSymbolicTracing.test_mega_guard

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

if a.shape[0] != b.shape[0] * 2:
    raise AssertionError("a.shape[0] should equal 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()

并断言当前场景下不存在非平凡 guards。

官方用例：

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.11.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中体现，并简要说明修改内容。如不涉及，需填写“不涉及”

不涉及。

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

【接口变更】

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

不涉及。

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

【功能验证】

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

1. patch dry-run 验证

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

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

验证结果：

checking file test/test_proxy_tensor.py
Hunk #2 succeeded at 997 (offset -1 lines).
Hunk #3 succeeded at 1868 (offset -2 lines).

验证结论：

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

2. 关联测试覆盖

本 PR 适配 PyTorch 官方已有测试：

test/test_proxy_tensor.py::TestSymbolicTracing::test_mega_guard
test/test_proxy_tensor.py::TestSymbolicTracing::test_guards_equal

【CheckList】

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

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

关联 PR

v2.11.0 自写测试 PR：

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

v2.10.0 同类 upstream patch PR：

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

v2.9.0 同类 upstream patch PR：

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