Pyrefly 张量形状类型系统贡献指南Fixture Stubs、类型级 Shape DSL 与验证工作流【免费下载链接】pyreflyA fast type checker and language server for Python项目地址: https://gitcode.com/GitHub_Trending/py/pyreflyPyrefly 的张量形状tensor shape类型追踪系统被刻意设计为一个可扩展的三层架构绝大多数 PyTorch 形状覆盖可以通过编辑.pyistub 和测试完成而无需触碰 Pyrefly 的 Rust 内核。本文基于仓库根目录的 TENSOR_SHAPES_CONTRIBUTING.md 展开完整覆盖三类形状机制Fixture stubs、类型级形状 DSL、Special handlers的工作原理与编写规范并给出从run_pyrefly.py静态校验、run_tests.py全量测试到cargo test shape_dsl内核测试的完整验证工作流。读完本文你可以独立地为 Pyrefly 添加新的形状 stub、编写类型级 DSL 函数并为移植的 PyTorch 模型补充assert_type形状检查点。架构总览三种互补的形状追踪机制Pyrefly 的形状追踪由三种互补机制组成详见 TENSOR_SHAPES_CONTRIBUTING.mdFixture stubs夹具 stub带形状泛型签名的.pyi文件覆盖nn.Linear、nn.Conv2d这样的模块和torch.mm这样的函数。类型级形状 DSL 函数用一小套 Python 子集在 tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi 中书写的形状变换函数以type_shape_dsl_function装饰并被公共返回注解直接调用。用于覆盖 reduction、padding、pooling、convolution 这类需要计算的形状逻辑。Special handlers特殊处理器Pyrefly 实现层的逻辑服务于需要更深类型系统集成的一级语法例如nn.Sequential链式调用、.shape、.size()、assert_shape与装饰器解释。前两种机制都位于tensor-shapes/目录下是添加或改进形状覆盖的常规途径。发布版 stub 只使用类型级 DSLPyrefly 目前仍保留旧的shape_dsl_function与uses_shape_dsl(...)V1 机制的内核支持与隔离测试仅为了让已锁定的 V1 stub 在迁移期间保持兼容——贡献者不应再添加新的 V1 规则。Special handlers 则需要改动 Pyrefly 实现应按内核工作kernel work对待。从源码结构看运行时支撑来自 tensor-shapes/pyrefly-shape-extensions/shape_extensions/init.py这个 Python 包提供Int、IntTuple、IntVar、D、assert_shape、type_shape_dsl_function等原语的最小运行时类——正如文件头注释所述.pyistub 为 Pyrefly 提供完整类型信息而.py文件只保证这些注解在 Python 中求值时不会崩溃。它还包含一个实用的兼容性技巧当检测到 torch/jax 已安装时会为torch.Tensor、nn.Linear等类动态打上__class_getitem__使Tensor[B, T, N]、nn.Linear[In, Out]这样的下标注解在运行时成为无操作no-op而不是抛出type is not subscriptable异常。Fixture Stubs用形状泛型签名描述张量变换存放位置tensor-shapes/pyrefly-torch-stubs/torch-stubs/ |-- __init__.pyi |-- _shapes.pyi |-- nn/ | |-- __init__.pyi # nn.Linear, nn.Conv2d, nn.LSTM, etc. | -- functional.pyi # F.relu, F.softmax, F.conv2d, etc. |-- distributions/ | -- ... # torch.distributions -- ...张量形状测试运行器会把tensor-shapes/作为 Pyrefly 搜索路径传入因此这些 stub 在验证时会覆盖常规的torchstub。Stub 是如何工作的Fixture stub 提供形状泛型的类型签名。以nn.Linear为例class LinearN, M: def __init__( self, in_features: SymInt[N], out_features: SymInt[M], bias: bool True, ) - None: ... def forward*Xs - Tensor[*Xs, M]: ...构造函数把输入/输出维度捕获为类型参数forward方法再借助变长类型参数*Xs表示可任意透传的 batch 维度。仓库中的实际实现 tensor-shapes/pyrefly-torch-stubs/torch-stubs/nn/init.pyi 采用了当前推荐写法类型参数以IntVar为界forward用IntTupleElements解构批量维度class LinearIN: IntVar, OUT: IntVar: weight: Tensor[[OUT, IN]] bias: Tensor[[OUT]] | None def __init__( self, in_features: _Int[IN], out_features: _Int[OUT], bias: bool True, ... ) - None: ... def forwardBs: IntTuple - Tensor[[*Elements[Bs], OUT]]: ...可以看到构造参数通过_Int[IN]/_Int[OUT]即Int[...]包裹与类型参数绑定而forward中*Elements[Bs]把任意 rank 的批量维度解包后原样透传——这正是batch 维度不变、仅最后一维从IN变为OUT这一形状语义的声明式表达。同文件中的Dropout、GELU则展示了最简形态forwardShape: IntTuple - Tensor[Shape]即形状完全透传。编写新 Stub 的步骤识别形状签名输入维度、输出维度以及它们之间的关系。对决定张量维度的参数使用SymInt[X]当前实现中等价于用IntVar界类型参数 Int[...]注解bias、dropout这类非形状参数保持原始类型。写出表达形状变换的方法或函数签名。对原样透传的 batch 维度使用*Xs或*Bs。把 stub 加入 tensor-shapes/pyrefly-torch-stubs/torch-stubs 中对应的.pyi文件。在tensor-shapes/pyrefly-torch-stubs/test/下添加或更新聚焦测试。示例添加一个新模块假设要添加保持空间维度不变的nn.GroupNormclass GroupNormNumGroups, NumChannels: def __init__( self, num_groups: SymInt[NumGroups], num_channels: SymInt[NumChannels], eps: float 1e-5, affine: bool True, ) - None: ... def forward*S - Tensor[*S]: ...由于GroupNorm不改变形状forward 签名就是简单的Tensor[*S] - Tensor[*S]。类型级 Shape DSL 函数当签名不足以表达输出形状时当一条普通签名无法表达输出形状时例如 conv 的 stride/padding 计算、reshape 的-1推导就需要类型级 DSL。存放位置与调用方式DSL 函数统一放在tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi公共 stub 在返回注解中直接调用类型级 DSL 函数。例如from shape_extensions import IntTuple, type_shape_dsl_function import shape_extensions.dsl as dsl type_shape_dsl_function def repeat_shape(shape: IntTuple, repeats: IntTuple) - IntTuple: if len(repeats) len(shape): return dsl.Invalid(repeat dimensions cannot be shorter than the input rank) extra len(repeats) - len(shape) return dsl.IntTuple( repeats[i] if i extra else shape[i - extra] * repeats[i] for i in range(len(repeats)) ) def repeatShape: IntTuple, Repeats: IntTuple - Tensor[repeat_shape(Shape, Repeats)]: ...真实的 _shapes.pyi 已包含十余个这样的函数。以reduce_shape为例它对dim is None全归约、is_int_value单维、负数维度、空元组、0 维标量的0/-1别名、重复维度等边界都做了显式处理越界或重复维度统一返回dsl.Invalid(dimension out of range)/dsl.Invalid(duplicate dimension)reshape_shape则用dsl.is_concrete_int区分具体与符号量在只允许一个-1的前提下推导其大小并对元素总数不匹配给出dsl.Invalid诊断。这些实现印证了下面DSL 子集一节的能力清单。DSL 子集有意保持很小的语言DSL 是有意的一个小而代数化的语言其主要值域是Int单个形状维度与IntTuple完整形状运行时配置值通过公共签名上的Flag[...]类型参数接入。函数体支持dsl.IntTuple(...)构造结果形状len、索引、切片以及有界生成器表达式、-、*、//、%等算术运算if/else单赋值局部变量直接return其他type_shape_dsl_function帮助函数的结果dsl.concat、dsl.prod、dsl.Invalid等 DSL 操作dsl.Int.gradual()表示一个渐进维度表达式return dsl.IntTuple.gradual()直接返回一个渐进形状。DSL 函数应保持简单、代数化。它们由 Pyrefly 分析而不是PyTorch 运算的常规运行时实现。整型参数与 IntVar 的传递边界这是 DSL 写作中最容易出错、也最无法靠猜的部分规范如下帮助函数参数消费单个维度时声明为IntNone有独立含义时声明为Int | None。公共签名中的运行时整数若需透传给帮助函数优先用一个界恰好为形状Int的类型参数并用该类型参数标注运行时参数type_shape_dsl_function def resize_shape(size: Int) - IntTuple: return dsl.IntTuple((size,)) def resizeN: Int - Tensor[resize_shape(N)]: ...这种形式要求 bound 恰好是Int。当类型参数只是在直接IntTuple或 list 形状语法中命名一个符号维度时改用IntVar。若这个符号维度还要传给 DSL 帮助函数必须在调用边界用Int[...]包裹def zerosN: IntVar - Tensor[[N]]: ... def resize_symbolicN: IntVar - Tensor[resize_shape(Int[N])]: ...裸IntVar参数以及N 1这类算术在帮助函数调用中会被拒绝应写成Int[N]与Int[N] 1。直接写成实参的Int[N] | None是类型联合语法不是运行时 DSL 值应按语义传Int[N]、None或 bound 恰好为Int | None的类型参数。仍解析为Int | None的实参在控制流收敛前按渐进gradual接受非None分支中可作为Int使用。宽泛的运行时int用作维度时会变成渐进维度保留已知 rank 与其他维度Any则保持未知不会被当作渐进整数。dsl.Int.gradual()本身是Int表达式可参与算术与dsl.IntTuple(...)构造dsl.IntTuple.gradual()目前只能作为 DSL 函数的直接返回表示整个渐进形状不能赋值给局部变量或嵌入更大的表达式。D[...]与D(...)是兼容包装器用于 Python 会急于求值的注解场景。它们不能替代Int[...]D[N]里仍是裸IntVar会被拒绝D[Int[N] 1]合法。当某个分支需要在形状求值期已知的整数字面量时用dsl.is_concrete_int(value)接受Int或Int | None它对None、符号维度、渐进Int均为False。用dsl.is_int_value(value)收窄兼容的Flag[int | tuple[int, ...] | None]值中的整型成员——注意它不能证明该整数是具体的。关于 NumPy 和 JAX stub 所用的另一套更小的类型级 DSL 子集仍在建设中有两条无法靠猜的经验值得记住一个使用了不支持语法的 DSL 函数会在所有调用点求值为Unknown而调用点本身不报告任何错误。真正的方法是直接类型检查 stub 文件——测试运行器会以stubssuite 的形式替你完成这一步。int | tuple[int, ...]类型的参数仅靠is_int_value收窄后不能迭代。以is None检查开头可以让收窄生效因此这类参数应声明为int | tuple[int, ...] | None函数体中拒绝None。Torch stub 的conv_shape与 JAX stub 的reshape_shape都采用这个写法。示例reductiontype_shape_dsl_function def reduce_shape(shape: IntTuple, dim: int, keepdim: bool) - IntTuple: axis dim % len(shape) return dsl.IntTuple( 1 if keepdim and i axis else shape[i] for i in range(len(shape)) if keepdim or i ! axis )公共 stub 会把输入形状与运行时选项绑定到类型参数然后在返回注解中调用reduce_shape(...)。添加新 DSL 函数的步骤在 tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi 中书写形状变换。用type_shape_dsl_function装饰。用Int、IntVar、IntTuple或Flag[...]类型参数绑定相关公共参数并在返回注解中调用 DSL 函数。添加使用assert_type检查计算后形状的正向测试。若 DSL 应当拒绝非法形状或报告形状错误添加带# E:期望的负向测试。旧的基于装饰器的 DSL 仅保留给尚未迁移的规则。避免在新规则中混合 V1 与 V2 逻辑如果 V2 目前无法表达某个操作记录该缺口document the gap而不是增加新的 V1 表面积。当前已知限制类型级 DSL 使用一个小的、可组合的语言并不会精确建模所有形状行为。当前已知限制包括符号arange的取整、unfold与diag_embed的符号配置值、结构化tensordot轴列表、符号 rank 形状或派生符号维度的乘积、以及 list 形式的 padding。遇到这些情况应尽可能保持渐进gradual补充聚焦测试并在受影响的规则处留下TODO(stroxler)使精度损失保持可见。移植模型Ported Models存放位置tensor-shapes/pyrefly-torch-stubs/examples/每个文件都是对某个真实 PyTorch 模型的完整标注移植带assert_type检查点与冒烟测试。添加新模型的步骤从 TorchBenchPyTorch 官方基准模型集或其他来源挑选一个模型。参照仓库内教程tensor-shapes-tutorial-basics 文档或 Agent 移植技能tensor-shapes/skills/add-shape-types-to-torch-model完成移植。在形状变化操作之后添加assert_type或assert_shape检查点。运行时执行有价值时在文件底部添加冒烟测试。运行verify_port.sh检查常见质量问题。verify_port.sh检查项该脚本检查移植模型中的常见问题tensor-shapes/skills/add-shape-types-to-torch-model/verify_port.sh tensor-shapes/pyrefly-torch-stubs/examples/model.py它报告如下指标指标含义igtype: ignore计数bs签名中的裸Tensor计数bv变量注解中的裸Tensor计数sh带形状的assert_type计数ba裸assert_type计数sm冒烟测试计数测试 Stub 与示例变更tensor-shape 专用 Pyrefly 运行器对多数贡献而言最重要的验证是 tensor-shape Pyrefly 运行器。它使用形状感知的 stub 检查聚焦测试、负向期望、jaxtyping 示例与示例语料库并且还会类型检查 stub 文件本身以stubssuite 形式报告。这一点比听起来更重要Pyrefly 只对它被要求检查的文件报告错误因此通过--search-path触达的 stub 是沉默的。一个无法编译的 stub 不会自我暴露它只是停止贡献类型让所有调用点安静地推断出Unknown——这看起来像规则缺失而不是规则损坏。直接检查 stub 能把这种情况变成带行号的错误。Torch 包目前通过其 run_pyrefly.py 中的check_stubsFalse暂时退出该检查源码中对应的 TODO 指出包中仍存在未完成的内部导入、类型参数遮蔽等问题且 torch-stubs/_shapes.pyi 的 V1shape_dsl_function函数体不是合法 Python类型级 DSL 文件可以干净通过检查因此迁移这些规则到类型级 DSL 正是移除该退出项的方式。基本用法python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py运行器会自行构建 Pyrefly——默认用cargo build加--buck时用 Buck——因此始终针对你的工作副本检查。先单独构建一次是多余的跳过构建则意味着运行结果来自旧版 Pyrefly。从 run_pyrefly.py 的参数定义可以确认完整的命令行选项--pyrefly /path/to/pyrefly显式传入二进制的唯一不构建模式——因为裸路径无法表达如何重新构建它--buck改用 Buck 构建并运行--release用 Cargo release profile 而非 debug 构建--python指定提供 Torch fallback 模块的解释器默认共享 virtualenv--suite可重复只运行指定 suite默认全部--nocapture流式打印 Pyrefly 完整输出默认只在失败时转储检查器输出成功时打印紧凑的PASS ...行。构建使用自定义目标目录时run_pyrefly.py会遵守CARGO_TARGET_DIR。迭代时单跑某个 suitesuites.py 定义了 Torch 包的五个 suite--suite的可选值正来源于此Suite匹配文件说明torch-examplesexamples/*.py、examples/runtime/*.py示例语料库torch-positivetest/test_*.py正向测试torch-negativetest/negative_tests/test_*.py带# E:期望的负向测试jaxtyping-positivetest/jaxtyping/test_*.pyjaxtyping 集成Python 3.12 配置jaxtyping-negativetest/jaxtyping/negative_tests/test_*.pyjaxtyping 负向测试python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-positive python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-negative python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-examplesstub 没有对应的 Buck 测试目标内部 checkout 通过--buck运行同一运行器只是以不同方式获取 Pyreflypython3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --buck一次运行所有库与 CI 完全一致python3 tensor-shapes/run_tests.py # 内部 checkout 追加 --buck python3 tensor-shapes/run_tests.py --static-only由于 Torch 与 NumPy 的形状 stub 会回退到已安装库中的定义即使--static-only也需要共享 virtualenv。可用--python选择安装了所需库的另一个 virtualenv 解释器根运行器把它用于运行时测试并转发给 Torch 与 NumPy 的静态检查。项目的 test.py 运行器把张量形状验证与默认 Pyrefly 测试循环分开只跑这些验证时python3 test.py --no-fmt --no-lint --no-test --tensor-shapes --no-conformance --no-jsonschema运行时测试Runtime Tests运行时测试验证注解帮助函数与可运行示例模型在 Python 中行为正确而不只是通过 Pyrefly 静态检查。测试位于tensor-shapes/pyrefly-torch-stubs/test/runtime_tests/运行时测试与静态 fallback 检查都需要共享 virtualenv同一环境同时提供 torch、numpy、jax。Bootstrap 是唯一下载依赖的步骤因此也是唯一需要网络访问的步骤python3 tensor-shapes/bootstrap_venv.py # 内部机器经由 fwdproxy 追加 --fwdproxy python3 tensor-shapes/run_tests.py --runtime-onlyvirtualenv 默认位于~/.tensor-shapes-venv设置$TENSOR_SHAPES_VENV可放到别处。各运行器从不创建virtualenv也从不触网缺失时它们会明确说明并打印 bootstrap 命令。Torch 与 NumPy 的静态检查使用已安装库的定义JAX 静态检查不需要 virtualenv。迭代时单跑某个 suitepython tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py --suite annotation python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py --suite model运行时运行器会为shape_extensions与可运行示例模块设置 import 路径。运行时测试在内部 checkout 中完全相同它们针对 virtualenv 运行从不经过 Buck从而保证没有任何工作流会重新构建 torch、numpy 或 jax。内核测试Kernel Tests多数贡献者不需要本节。只有当改动的是 Pyrefly 的张量形状内核而非仅 stub 或示例时才使用这些测试。内核变更包括shape_extensions原语或装饰器assert_shape的类型检查器行为shape_dsl_function的解析、校验或求值uses_shape_dsl的处理Pyrefly Rust 源码中的 special handlers。聚焦的 Pyrefly 单元测试位于 pyrefly/lib/test/shape_dsl.rs其中保留的 V1 兼容路径测试隔离在该文件的legacy模块内并使用私有内存 stub。用 Cargo 运行cargo test shape_dsl内核测试刻意比 stub/示例 suite 小得多它覆盖核心原语与不变量而张量形状 stub 测试通过真实的 PyTorch 签名对 DSL 施压。提交前检查Pre-Commit Checkstensor-shape 包中的 Python 文件使用Ruff 格式化器而非 Black从仓库根目录以与 CI 相同的 Ruff 版本格式化uv tool run --from ruff0.16.5 ruff format \ tensor-shapesskills目录是文档而非语料源码之所以被排除是因为 Ruff 还会格式化 Markdown 中内嵌的 Python 片段。在移交变更之前还要运行仓库级格式化与 lint./test.py --no-test --no-tensor-shapes --no-conformance --no-jsonschema以及按所触达文件运行相应的张量形状检查Stub/测试/示例变更python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py运行时帮助函数或可运行模型变更python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py内核变更cargo test shape_dsl或上述 Buck 等价方式小结选择正确的贡献路径你想做的事机制主要改动位置验证命令覆盖nn.Xxx、torch.xxx的形状Fixture stubtorch-stubs/nn/、torch-stubs/*.pyirun_pyrefly.py --suite torch-positive表达需要计算的形状reduction、conv、padding…类型级 DSLtorch-stubs/_shapes.pyirun_pyrefly.py全量用真实模型验证覆盖质量移植模型pyrefly-torch-stubs/examples/verify_port.sh--suite torch-examples校验运行时注解行为运行时测试test/runtime_tests/run_runtime_tests.py改形状内核原语Special handlers / 内核pyrefly/lib/Rustcargo test shape_dsl这条stub 优先、DSL 次之、内核兜底的分工正是 TENSOR_SHAPES_CONTRIBUTING.md 的核心主张外部贡献应停留在 stub-only 或 example/test-only 层面内核变更属于更窄的工作流。遵循这一分层你可以不触碰 Rust 内部实现就把 Pyrefly 的 PyTorch 形状覆盖稳步扩大。【免费下载链接】pyreflyA fast type checker and language server for Python项目地址: https://gitcode.com/GitHub_Trending/py/pyrefly创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 Pyrefly 张量形状类型系统贡献指南：Fixture Stubs、类型级 Shape DSL 与验证工作流
Pyrefly 张量形状类型系统贡献指南Fixture Stubs、类型级 Shape DSL 与验证工作流【免费下载链接】pyreflyA fast type checker and language server for Python项目地址: https://gitcode.com/GitHub_Trending/py/pyreflyPyrefly 的张量形状tensor shape类型追踪系统被刻意设计为一个可扩展的三层架构绝大多数 PyTorch 形状覆盖可以通过编辑.pyistub 和测试完成而无需触碰 Pyrefly 的 Rust 内核。本文基于仓库根目录的 TENSOR_SHAPES_CONTRIBUTING.md 展开完整覆盖三类形状机制Fixture stubs、类型级形状 DSL、Special handlers的工作原理与编写规范并给出从run_pyrefly.py静态校验、run_tests.py全量测试到cargo test shape_dsl内核测试的完整验证工作流。读完本文你可以独立地为 Pyrefly 添加新的形状 stub、编写类型级 DSL 函数并为移植的 PyTorch 模型补充assert_type形状检查点。架构总览三种互补的形状追踪机制Pyrefly 的形状追踪由三种互补机制组成详见 TENSOR_SHAPES_CONTRIBUTING.mdFixture stubs夹具 stub带形状泛型签名的.pyi文件覆盖nn.Linear、nn.Conv2d这样的模块和torch.mm这样的函数。类型级形状 DSL 函数用一小套 Python 子集在 tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi 中书写的形状变换函数以type_shape_dsl_function装饰并被公共返回注解直接调用。用于覆盖 reduction、padding、pooling、convolution 这类需要计算的形状逻辑。Special handlers特殊处理器Pyrefly 实现层的逻辑服务于需要更深类型系统集成的一级语法例如nn.Sequential链式调用、.shape、.size()、assert_shape与装饰器解释。前两种机制都位于tensor-shapes/目录下是添加或改进形状覆盖的常规途径。发布版 stub 只使用类型级 DSLPyrefly 目前仍保留旧的shape_dsl_function与uses_shape_dsl(...)V1 机制的内核支持与隔离测试仅为了让已锁定的 V1 stub 在迁移期间保持兼容——贡献者不应再添加新的 V1 规则。Special handlers 则需要改动 Pyrefly 实现应按内核工作kernel work对待。从源码结构看运行时支撑来自 tensor-shapes/pyrefly-shape-extensions/shape_extensions/init.py这个 Python 包提供Int、IntTuple、IntVar、D、assert_shape、type_shape_dsl_function等原语的最小运行时类——正如文件头注释所述.pyistub 为 Pyrefly 提供完整类型信息而.py文件只保证这些注解在 Python 中求值时不会崩溃。它还包含一个实用的兼容性技巧当检测到 torch/jax 已安装时会为torch.Tensor、nn.Linear等类动态打上__class_getitem__使Tensor[B, T, N]、nn.Linear[In, Out]这样的下标注解在运行时成为无操作no-op而不是抛出type is not subscriptable异常。Fixture Stubs用形状泛型签名描述张量变换存放位置tensor-shapes/pyrefly-torch-stubs/torch-stubs/ |-- __init__.pyi |-- _shapes.pyi |-- nn/ | |-- __init__.pyi # nn.Linear, nn.Conv2d, nn.LSTM, etc. | -- functional.pyi # F.relu, F.softmax, F.conv2d, etc. |-- distributions/ | -- ... # torch.distributions -- ...张量形状测试运行器会把tensor-shapes/作为 Pyrefly 搜索路径传入因此这些 stub 在验证时会覆盖常规的torchstub。Stub 是如何工作的Fixture stub 提供形状泛型的类型签名。以nn.Linear为例class LinearN, M: def __init__( self, in_features: SymInt[N], out_features: SymInt[M], bias: bool True, ) - None: ... def forward*Xs - Tensor[*Xs, M]: ...构造函数把输入/输出维度捕获为类型参数forward方法再借助变长类型参数*Xs表示可任意透传的 batch 维度。仓库中的实际实现 tensor-shapes/pyrefly-torch-stubs/torch-stubs/nn/init.pyi 采用了当前推荐写法类型参数以IntVar为界forward用IntTupleElements解构批量维度class LinearIN: IntVar, OUT: IntVar: weight: Tensor[[OUT, IN]] bias: Tensor[[OUT]] | None def __init__( self, in_features: _Int[IN], out_features: _Int[OUT], bias: bool True, ... ) - None: ... def forwardBs: IntTuple - Tensor[[*Elements[Bs], OUT]]: ...可以看到构造参数通过_Int[IN]/_Int[OUT]即Int[...]包裹与类型参数绑定而forward中*Elements[Bs]把任意 rank 的批量维度解包后原样透传——这正是batch 维度不变、仅最后一维从IN变为OUT这一形状语义的声明式表达。同文件中的Dropout、GELU则展示了最简形态forwardShape: IntTuple - Tensor[Shape]即形状完全透传。编写新 Stub 的步骤识别形状签名输入维度、输出维度以及它们之间的关系。对决定张量维度的参数使用SymInt[X]当前实现中等价于用IntVar界类型参数 Int[...]注解bias、dropout这类非形状参数保持原始类型。写出表达形状变换的方法或函数签名。对原样透传的 batch 维度使用*Xs或*Bs。把 stub 加入 tensor-shapes/pyrefly-torch-stubs/torch-stubs 中对应的.pyi文件。在tensor-shapes/pyrefly-torch-stubs/test/下添加或更新聚焦测试。示例添加一个新模块假设要添加保持空间维度不变的nn.GroupNormclass GroupNormNumGroups, NumChannels: def __init__( self, num_groups: SymInt[NumGroups], num_channels: SymInt[NumChannels], eps: float 1e-5, affine: bool True, ) - None: ... def forward*S - Tensor[*S]: ...由于GroupNorm不改变形状forward 签名就是简单的Tensor[*S] - Tensor[*S]。类型级 Shape DSL 函数当签名不足以表达输出形状时当一条普通签名无法表达输出形状时例如 conv 的 stride/padding 计算、reshape 的-1推导就需要类型级 DSL。存放位置与调用方式DSL 函数统一放在tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi公共 stub 在返回注解中直接调用类型级 DSL 函数。例如from shape_extensions import IntTuple, type_shape_dsl_function import shape_extensions.dsl as dsl type_shape_dsl_function def repeat_shape(shape: IntTuple, repeats: IntTuple) - IntTuple: if len(repeats) len(shape): return dsl.Invalid(repeat dimensions cannot be shorter than the input rank) extra len(repeats) - len(shape) return dsl.IntTuple( repeats[i] if i extra else shape[i - extra] * repeats[i] for i in range(len(repeats)) ) def repeatShape: IntTuple, Repeats: IntTuple - Tensor[repeat_shape(Shape, Repeats)]: ...真实的 _shapes.pyi 已包含十余个这样的函数。以reduce_shape为例它对dim is None全归约、is_int_value单维、负数维度、空元组、0 维标量的0/-1别名、重复维度等边界都做了显式处理越界或重复维度统一返回dsl.Invalid(dimension out of range)/dsl.Invalid(duplicate dimension)reshape_shape则用dsl.is_concrete_int区分具体与符号量在只允许一个-1的前提下推导其大小并对元素总数不匹配给出dsl.Invalid诊断。这些实现印证了下面DSL 子集一节的能力清单。DSL 子集有意保持很小的语言DSL 是有意的一个小而代数化的语言其主要值域是Int单个形状维度与IntTuple完整形状运行时配置值通过公共签名上的Flag[...]类型参数接入。函数体支持dsl.IntTuple(...)构造结果形状len、索引、切片以及有界生成器表达式、-、*、//、%等算术运算if/else单赋值局部变量直接return其他type_shape_dsl_function帮助函数的结果dsl.concat、dsl.prod、dsl.Invalid等 DSL 操作dsl.Int.gradual()表示一个渐进维度表达式return dsl.IntTuple.gradual()直接返回一个渐进形状。DSL 函数应保持简单、代数化。它们由 Pyrefly 分析而不是PyTorch 运算的常规运行时实现。整型参数与 IntVar 的传递边界这是 DSL 写作中最容易出错、也最无法靠猜的部分规范如下帮助函数参数消费单个维度时声明为IntNone有独立含义时声明为Int | None。公共签名中的运行时整数若需透传给帮助函数优先用一个界恰好为形状Int的类型参数并用该类型参数标注运行时参数type_shape_dsl_function def resize_shape(size: Int) - IntTuple: return dsl.IntTuple((size,)) def resizeN: Int - Tensor[resize_shape(N)]: ...这种形式要求 bound 恰好是Int。当类型参数只是在直接IntTuple或 list 形状语法中命名一个符号维度时改用IntVar。若这个符号维度还要传给 DSL 帮助函数必须在调用边界用Int[...]包裹def zerosN: IntVar - Tensor[[N]]: ... def resize_symbolicN: IntVar - Tensor[resize_shape(Int[N])]: ...裸IntVar参数以及N 1这类算术在帮助函数调用中会被拒绝应写成Int[N]与Int[N] 1。直接写成实参的Int[N] | None是类型联合语法不是运行时 DSL 值应按语义传Int[N]、None或 bound 恰好为Int | None的类型参数。仍解析为Int | None的实参在控制流收敛前按渐进gradual接受非None分支中可作为Int使用。宽泛的运行时int用作维度时会变成渐进维度保留已知 rank 与其他维度Any则保持未知不会被当作渐进整数。dsl.Int.gradual()本身是Int表达式可参与算术与dsl.IntTuple(...)构造dsl.IntTuple.gradual()目前只能作为 DSL 函数的直接返回表示整个渐进形状不能赋值给局部变量或嵌入更大的表达式。D[...]与D(...)是兼容包装器用于 Python 会急于求值的注解场景。它们不能替代Int[...]D[N]里仍是裸IntVar会被拒绝D[Int[N] 1]合法。当某个分支需要在形状求值期已知的整数字面量时用dsl.is_concrete_int(value)接受Int或Int | None它对None、符号维度、渐进Int均为False。用dsl.is_int_value(value)收窄兼容的Flag[int | tuple[int, ...] | None]值中的整型成员——注意它不能证明该整数是具体的。关于 NumPy 和 JAX stub 所用的另一套更小的类型级 DSL 子集仍在建设中有两条无法靠猜的经验值得记住一个使用了不支持语法的 DSL 函数会在所有调用点求值为Unknown而调用点本身不报告任何错误。真正的方法是直接类型检查 stub 文件——测试运行器会以stubssuite 的形式替你完成这一步。int | tuple[int, ...]类型的参数仅靠is_int_value收窄后不能迭代。以is None检查开头可以让收窄生效因此这类参数应声明为int | tuple[int, ...] | None函数体中拒绝None。Torch stub 的conv_shape与 JAX stub 的reshape_shape都采用这个写法。示例reductiontype_shape_dsl_function def reduce_shape(shape: IntTuple, dim: int, keepdim: bool) - IntTuple: axis dim % len(shape) return dsl.IntTuple( 1 if keepdim and i axis else shape[i] for i in range(len(shape)) if keepdim or i ! axis )公共 stub 会把输入形状与运行时选项绑定到类型参数然后在返回注解中调用reduce_shape(...)。添加新 DSL 函数的步骤在 tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi 中书写形状变换。用type_shape_dsl_function装饰。用Int、IntVar、IntTuple或Flag[...]类型参数绑定相关公共参数并在返回注解中调用 DSL 函数。添加使用assert_type检查计算后形状的正向测试。若 DSL 应当拒绝非法形状或报告形状错误添加带# E:期望的负向测试。旧的基于装饰器的 DSL 仅保留给尚未迁移的规则。避免在新规则中混合 V1 与 V2 逻辑如果 V2 目前无法表达某个操作记录该缺口document the gap而不是增加新的 V1 表面积。当前已知限制类型级 DSL 使用一个小的、可组合的语言并不会精确建模所有形状行为。当前已知限制包括符号arange的取整、unfold与diag_embed的符号配置值、结构化tensordot轴列表、符号 rank 形状或派生符号维度的乘积、以及 list 形式的 padding。遇到这些情况应尽可能保持渐进gradual补充聚焦测试并在受影响的规则处留下TODO(stroxler)使精度损失保持可见。移植模型Ported Models存放位置tensor-shapes/pyrefly-torch-stubs/examples/每个文件都是对某个真实 PyTorch 模型的完整标注移植带assert_type检查点与冒烟测试。添加新模型的步骤从 TorchBenchPyTorch 官方基准模型集或其他来源挑选一个模型。参照仓库内教程tensor-shapes-tutorial-basics 文档或 Agent 移植技能tensor-shapes/skills/add-shape-types-to-torch-model完成移植。在形状变化操作之后添加assert_type或assert_shape检查点。运行时执行有价值时在文件底部添加冒烟测试。运行verify_port.sh检查常见质量问题。verify_port.sh检查项该脚本检查移植模型中的常见问题tensor-shapes/skills/add-shape-types-to-torch-model/verify_port.sh tensor-shapes/pyrefly-torch-stubs/examples/model.py它报告如下指标指标含义igtype: ignore计数bs签名中的裸Tensor计数bv变量注解中的裸Tensor计数sh带形状的assert_type计数ba裸assert_type计数sm冒烟测试计数测试 Stub 与示例变更tensor-shape 专用 Pyrefly 运行器对多数贡献而言最重要的验证是 tensor-shape Pyrefly 运行器。它使用形状感知的 stub 检查聚焦测试、负向期望、jaxtyping 示例与示例语料库并且还会类型检查 stub 文件本身以stubssuite 形式报告。这一点比听起来更重要Pyrefly 只对它被要求检查的文件报告错误因此通过--search-path触达的 stub 是沉默的。一个无法编译的 stub 不会自我暴露它只是停止贡献类型让所有调用点安静地推断出Unknown——这看起来像规则缺失而不是规则损坏。直接检查 stub 能把这种情况变成带行号的错误。Torch 包目前通过其 run_pyrefly.py 中的check_stubsFalse暂时退出该检查源码中对应的 TODO 指出包中仍存在未完成的内部导入、类型参数遮蔽等问题且 torch-stubs/_shapes.pyi 的 V1shape_dsl_function函数体不是合法 Python类型级 DSL 文件可以干净通过检查因此迁移这些规则到类型级 DSL 正是移除该退出项的方式。基本用法python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py运行器会自行构建 Pyrefly——默认用cargo build加--buck时用 Buck——因此始终针对你的工作副本检查。先单独构建一次是多余的跳过构建则意味着运行结果来自旧版 Pyrefly。从 run_pyrefly.py 的参数定义可以确认完整的命令行选项--pyrefly /path/to/pyrefly显式传入二进制的唯一不构建模式——因为裸路径无法表达如何重新构建它--buck改用 Buck 构建并运行--release用 Cargo release profile 而非 debug 构建--python指定提供 Torch fallback 模块的解释器默认共享 virtualenv--suite可重复只运行指定 suite默认全部--nocapture流式打印 Pyrefly 完整输出默认只在失败时转储检查器输出成功时打印紧凑的PASS ...行。构建使用自定义目标目录时run_pyrefly.py会遵守CARGO_TARGET_DIR。迭代时单跑某个 suitesuites.py 定义了 Torch 包的五个 suite--suite的可选值正来源于此Suite匹配文件说明torch-examplesexamples/*.py、examples/runtime/*.py示例语料库torch-positivetest/test_*.py正向测试torch-negativetest/negative_tests/test_*.py带# E:期望的负向测试jaxtyping-positivetest/jaxtyping/test_*.pyjaxtyping 集成Python 3.12 配置jaxtyping-negativetest/jaxtyping/negative_tests/test_*.pyjaxtyping 负向测试python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-positive python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-negative python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-examplesstub 没有对应的 Buck 测试目标内部 checkout 通过--buck运行同一运行器只是以不同方式获取 Pyreflypython3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --buck一次运行所有库与 CI 完全一致python3 tensor-shapes/run_tests.py # 内部 checkout 追加 --buck python3 tensor-shapes/run_tests.py --static-only由于 Torch 与 NumPy 的形状 stub 会回退到已安装库中的定义即使--static-only也需要共享 virtualenv。可用--python选择安装了所需库的另一个 virtualenv 解释器根运行器把它用于运行时测试并转发给 Torch 与 NumPy 的静态检查。项目的 test.py 运行器把张量形状验证与默认 Pyrefly 测试循环分开只跑这些验证时python3 test.py --no-fmt --no-lint --no-test --tensor-shapes --no-conformance --no-jsonschema运行时测试Runtime Tests运行时测试验证注解帮助函数与可运行示例模型在 Python 中行为正确而不只是通过 Pyrefly 静态检查。测试位于tensor-shapes/pyrefly-torch-stubs/test/runtime_tests/运行时测试与静态 fallback 检查都需要共享 virtualenv同一环境同时提供 torch、numpy、jax。Bootstrap 是唯一下载依赖的步骤因此也是唯一需要网络访问的步骤python3 tensor-shapes/bootstrap_venv.py # 内部机器经由 fwdproxy 追加 --fwdproxy python3 tensor-shapes/run_tests.py --runtime-onlyvirtualenv 默认位于~/.tensor-shapes-venv设置$TENSOR_SHAPES_VENV可放到别处。各运行器从不创建virtualenv也从不触网缺失时它们会明确说明并打印 bootstrap 命令。Torch 与 NumPy 的静态检查使用已安装库的定义JAX 静态检查不需要 virtualenv。迭代时单跑某个 suitepython tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py --suite annotation python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py --suite model运行时运行器会为shape_extensions与可运行示例模块设置 import 路径。运行时测试在内部 checkout 中完全相同它们针对 virtualenv 运行从不经过 Buck从而保证没有任何工作流会重新构建 torch、numpy 或 jax。内核测试Kernel Tests多数贡献者不需要本节。只有当改动的是 Pyrefly 的张量形状内核而非仅 stub 或示例时才使用这些测试。内核变更包括shape_extensions原语或装饰器assert_shape的类型检查器行为shape_dsl_function的解析、校验或求值uses_shape_dsl的处理Pyrefly Rust 源码中的 special handlers。聚焦的 Pyrefly 单元测试位于 pyrefly/lib/test/shape_dsl.rs其中保留的 V1 兼容路径测试隔离在该文件的legacy模块内并使用私有内存 stub。用 Cargo 运行cargo test shape_dsl内核测试刻意比 stub/示例 suite 小得多它覆盖核心原语与不变量而张量形状 stub 测试通过真实的 PyTorch 签名对 DSL 施压。提交前检查Pre-Commit Checkstensor-shape 包中的 Python 文件使用Ruff 格式化器而非 Black从仓库根目录以与 CI 相同的 Ruff 版本格式化uv tool run --from ruff0.16.5 ruff format \ tensor-shapesskills目录是文档而非语料源码之所以被排除是因为 Ruff 还会格式化 Markdown 中内嵌的 Python 片段。在移交变更之前还要运行仓库级格式化与 lint./test.py --no-test --no-tensor-shapes --no-conformance --no-jsonschema以及按所触达文件运行相应的张量形状检查Stub/测试/示例变更python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py运行时帮助函数或可运行模型变更python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py内核变更cargo test shape_dsl或上述 Buck 等价方式小结选择正确的贡献路径你想做的事机制主要改动位置验证命令覆盖nn.Xxx、torch.xxx的形状Fixture stubtorch-stubs/nn/、torch-stubs/*.pyirun_pyrefly.py --suite torch-positive表达需要计算的形状reduction、conv、padding…类型级 DSLtorch-stubs/_shapes.pyirun_pyrefly.py全量用真实模型验证覆盖质量移植模型pyrefly-torch-stubs/examples/verify_port.sh--suite torch-examples校验运行时注解行为运行时测试test/runtime_tests/run_runtime_tests.py改形状内核原语Special handlers / 内核pyrefly/lib/Rustcargo test shape_dsl这条stub 优先、DSL 次之、内核兜底的分工正是 TENSOR_SHAPES_CONTRIBUTING.md 的核心主张外部贡献应停留在 stub-only 或 example/test-only 层面内核变更属于更窄的工作流。遵循这一分层你可以不触碰 Rust 内部实现就把 Pyrefly 的 PyTorch 形状覆盖稳步扩大。【免费下载链接】pyreflyA fast type checker and language server for Python项目地址: https://gitcode.com/GitHub_Trending/py/pyrefly创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考