已合并
docs:更新 torch.fx.experimental.symbolic_shapes API 文档 #36879
cuiyunhao-2026创建于 5月27日
docs:更新 torch.fx.experimental.symbolic_shapes API 文档 #36879
已合并
cuiyunhao-2026创建于 5月27日
cuiyunhao-2026
5月27日

【合入来源】

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

关联issue:https://gitcode.com/Ascend/pytorch/issues/1625

【修改方案】

完善 Torch-NPU 对 torch.fx.experimental.symbolic_shapes.ShapeEnv 类 5 个 API 在 NPU 环境下的兼容性验证,覆盖以下 API:

API 说明
torch.fx.experimental.symbolic_shapes.ShapeEnv.format_guards 将 ShapeEnv 中的 guard 表达式格式化为字符串,verbose=True 时附加源码位置
torch.fx.experimental.symbolic_shapes.ShapeEnv.freeze 冻结 ShapeEnv 停止累积 guards,设置 frozen=True
torch.fx.experimental.symbolic_shapes.ShapeEnv.freeze_runtime_asserts 冻结 ShapeEnv 停止添加延迟 runtime asserts,设置 runtime_asserts_frozen=True
torch.fx.experimental.symbolic_shapes.ShapeEnv.get_axioms 返回 runtime asserts 和 guards 的组合,支持 symbols 筛选和 compute_hint 参数
torch.fx.experimental.symbolic_shapes.ShapeEnv.get_implications 给定布尔表达式,返回所有蕴含谓词及其真值

1. 上游社区用例情况分析

在 PyTorch 上游仓库(github.com/pytorch/pytorch)中搜索这 5 个 API:

  • format_guards:仅在 test/functorch/test_aotdispatch.py(第10150行)和 test/test_proxy_tensor.py(第1963行)中间接调用,无直接 API 验证用例
  • freeze:上游 test/ 目录下无任何 ShapeEnv.freeze() 的直接调用
  • freeze_runtime_asserts:上游 test/ 目录下无任何调用
  • get_axioms:上游 test/ 目录下无任何直接调用(仅在 torch/_inductor/sizevars.py 等内部使用)
  • get_implications:上游 test/ 目录下无任何直接调用(仅在 symbolic_shapes.py 内部被其他方法调用)

上游 test/fx/ 目录下不存在 test_symbolic_shapes.py 文件,该文件为 torch-npu 社区自建。

结论:5 个 API 均属于上游"无直接用例"场景,按 1.3 规范自写用例,提交到 torch-npu 的 test 目录。

2. NPU 适配方案

5 个 API 均为 ShapeEnv 类的纯 Python 方法,操作 sympy 符号表达式与内部状态,不涉及张量计算和 NPU 算子调用,在 NPU 上行为与 CPU 完全一致,无需做任何 NPU 适配修改。

3. 新增测试用例

在已有 test/fx/test_symbolic_shapes.pyTestShapeEnvNPU 类中追加 5 个测试方法:

测试方法 覆盖场景
test_torch_fx_experimental_symbolic_shapes_ShapeEnv_format_guards 空 guards 返回空字符串;多条 guards 格式化输出;verbose 模式含源码位置信息
test_torch_fx_experimental_symbolic_shapes_ShapeEnv_freeze 冻结前后 frozen 状态切换
test_torch_fx_experimental_symbolic_shapes_ShapeEnv_freeze_runtime_asserts 冻结前后 runtime_asserts_frozen 状态切换
test_torch_fx_experimental_symbolic_shapes_ShapeEnv_get_axioms 验证返回类型为 tuple;symbols 参数筛选
test_torch_fx_experimental_symbolic_shapes_ShapeEnv_get_implications Eq/Lt/Ne/Le 表达式蕴含推导

4. API 补齐

不需要。5 个 API 均为纯 Python 方法,在 NPU 上运行结果与 CPU 完全一致。

5. 多版本分支合入

目标分支 PR
v2.7.1 https://gitcode.com/Ascend/pytorch/pull/36867
v2.9.0 https://gitcode.com/Ascend/pytorch/pull/36871
v2.10.0 https://gitcode.com/Ascend/pytorch/pull/36872
v2.11.0 https://gitcode.com/Ascend/pytorch/pull/36873
v2.12.0 https://gitcode.com/Ascend/pytorch/pull/36874
master https://gitcode.com/Ascend/pytorch/pull/36866

文档 PR:https://gitcode.com/Ascend/pytorch/pull/36879

【资料变更】

涉及。torch.fx.experimental.symbolic_shapes 下的 5 个 API 实际已支持,但 docs/zh/native_apis 中缺少对应文档。

v2.7.1 分支资料变更内容:

  1. 新增 docs/zh/native_apis/pytorch_2-7-1/torch-fx-experimental-symbolic_shapes.mdpytorch_2-9-0/torch-fx-experimental-symbolic_shapes.mdpytorch_2-10-0/torch-fx-experimental-symbolic_shapes.mdpytorch_2-11-0/torch-fx-experimental-symbolic_shapes.mdpytorch_2-12-0/torch-fx-experimental-symbolic_shapes.md
  2. 修改对应版本 PyTorch-*.md,添加子模块链接

API 支持情况:

API名称 是否支持 限制与说明
torch.fx.experimental.symbolic_shapes.ShapeEnv -
torch.fx.experimental.symbolic_shapes.ShapeEnv.format_guards -
torch.fx.experimental.symbolic_shapes.ShapeEnv.freeze -
torch.fx.experimental.symbolic_shapes.ShapeEnv.freeze_runtime_asserts -
torch.fx.experimental.symbolic_shapes.ShapeEnv.get_axioms -
torch.fx.experimental.symbolic_shapes.ShapeEnv.get_implications -

资料 PR 仅合入 v2.7.1 分支,刷新 pytorch_2-7-1、pytorch_2-9-0、pytorch_2-10-0、pytorch_2-11-0、pytorch_2-12-0 五个版本的 API 支持情况。v2.9.0 及以后分支无 docs/zh/native_apis 目录,无需处理。

【接口变更】

不涉及(无新增/修改对外 C++ 或 Python 接口,仅新增测试与文档)

【功能验证】

1. 测试环境

  • 设备:Ascend NPU
  • 框架:对应分支 Torch-NPU + PyTorch
  • 说明:5 个 API 为 ShapeEnv 类的纯 Python 方法,操作符号表达式与内部状态,不涉及张量计算,在 NPU 上行为与 CPU 完全一致

2. 测试命令

cd <pytorch-repo-root>
python test/fx/test_symbolic_shapes.py

3.测试结果
测试结果.png

4.与上游用例的关系

验证维度 上游现有用例 本 PR 新增用例
ShapeEnv.format_guards 无直接用例(test_aotdispatch.py 间接调用) 已覆盖(空guards/有guards/verbose)
ShapeEnv.freeze 无任何调用 已覆盖(状态切换)
ShapeEnv.freeze_runtime_asserts 无任何调用 已覆盖(状态切换)
ShapeEnv.get_axioms 无直接用例 已覆盖(空/guards/symbols/hint)
ShapeEnv.get_implications 无直接用例 已覆盖(Eq/Lt/Ne/Le 蕴含推导)
NPU 侧验证闭环 不完整 本 PR 补齐

【CheckList】

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

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 cuiyunhao-2026 的贡献)
cuiyunhao-2026cuiyunhao-2026
5月27日 创建了 pull request,commit 24cbca75
cuiyunhao-2026cuiyunhao-2026
5月27日 关联了issue:【开源实习】Torch-NPU API一致性对齐:补齐测试用例、API功能对齐、补齐文档(19)
ascend-robotascend-robot成员
5月27日 添加了label:stat/needs-squash
ascend-robot
ascend-robot成员
5月27日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here
You can get sig-info at here


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
docs renyujin, huangjingwei, lyx324521 (3/2) renyujin (1/1)
test renyujin, huangjingwei (2/2) renyujin (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

cuiyunhao-2026, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了118条消息 查看更多
Jingwei Huang
Jingwei Huang成员
7月8日 评论:

/lgtm

likedislike
renyujin成员
7月8日 评论:

/approve

likedislike
ascend-robotascend-robot成员
7月8日 添加了label:approvedlgtm
ascend-robotascend-robot成员
7月8日 合入了pull request
ascend-robot
ascend-robot成员
7月8日 评论:

Pull Request 已合并或已关闭。

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

likedislike