Pyrefly Tensor Shape 注解的运行时基石：深入解析 pyrefly-shape-extensions 包-二趣网

Pyrefly Tensor Shape 注解的运行时基石：深入解析 pyrefly-shape-extensions 包

【免费下载链接】pyreflyA fast type checker and language server for Python项目地址: https://gitcode.com/GitHub_Trending/py/pyrefly

导读：本文以仓库 tensor-shapes/pyrefly-shape-extensions/README.md 为核心，讲解 Pyrefly（一个面向 Python 的快速类型检查器与语言服务器）如何通过pyrefly-shape-extensions这个轻量运行时包，让Tensor[B, T]、IntVar("B")、assert_shape(x.shape, (2, 3))等静态形状注解既能在 Python 解释器中"平安无事"地求值，又能在 Pyrefly 静态检查时发挥完整的形状推导能力。读完本文，你将掌握 shape typing 的"静态桩（.pyi）与运行时空操作（no-op）"双轨设计、各形状原语的运行时语义、TorchScript 兼容方案，以及如何在真实模型（Torch/JAX/NumPy）中使用这套注解。

一、为什么需要这个包：静态形状注解的运行时难题

Pyrefly 的 Tensor Shape 系统允许开发者写出如下注解：

Tensor[B, T] # B、T 是符号维度变量 nn.Linear[In, Out] # 模块的输入/输出维度 IntVar("B") # 声明一个符号整数维度 assert_shape(x.shape, (2, 3)) # 运行时与静态双重形状断言

问题在于：这些写法在 Python 运行时并不天然合法。torch.Tensor不是泛型类，直接写Tensor[B, T]会触发TypeError: type 'torch.Tensor' is not subscriptable；IntVar("B")如果在运行时参与N + 1这类运算也会因算术运算符缺失而崩溃。

pyrefly-shape-extensions的定位正是解决这一鸿沟。README 明确说明：

This package provides the lightweightshape_extensionsmodule used by Pyrefly's tensor shape stubs. It defines runtime no-op versions of the shape typing primitives...

即：静态检查由.pyi桩承担全部类型推导，运行时由本包提供"最小可执行、绝不崩溃"的空操作版本，二者通过同名同 API 对齐。

这一"双轨"设计的代码证据在 shape_extensions/init.py 的模块 docstring 中写得非常直白：

The .pyi stub provides full type information to pyrefly. This .py file provides minimal runtime classes so that annotations using these types don't crash when evaluated by Python.

二、安装与项目结构

包使用 hatchling 构建，定义在 pyproject.toml：

[build-system] requires = ["hatchling"] build-backend = "hatchling.build" [project] name = "pyrefly-shape-extensions" requires-python = ">=3.12" version = "0.0.0" # 与 Pyrefly 主版本 lockstep（同步版本） classifiers = [ "Development Status :: 4 - Beta", "License :: OSI Approved :: MIT License", "Programming Language :: Python", ] [tool.hatch.build.targets.wheel] packages = ["shape_extensions"]

仓库结构非常精简，仅四个文件加一个 LICENSE：

tensor-shapes/pyrefly-shape-extensions/ ├── shape_extensions/ │ ├── __init__.py # 公共原语与运行时 no-op 实现 │ ├── dsl.py # 类型级 Shape DSL 的内部实现 │ ├── py.typed # PEP 561 内联类型标记 │ └── torchscript.py # TorchScript 兼容层（擦除形状注解） ├── LICENSE ├── README.md └── pyproject.toml

需要注意两点：

requires-python = ">=3.12"：包依赖 Python 3.12+，因为源码中大量使用 PEP 695 的泛型语法（如class Int[T]:）和type语句。
版本 lockstep：README 声明 "The package is versioned in lockstep with Pyrefly."，pyproject.toml中version = "0.0.0"是开发占位版本，正式发布时与 Pyrefly 主版本同步。

三、核心原语：运行时 no-op 的完整语义

shape_extensions/init.py 通过__all__导出了全部公共 API，共 18 个名字。下面按功能分组逐一解读（每条都对应源码中的真实实现）。

3.1 形状容器：IntTuple / IntTuples / Elements

class IntTuple: def __new__(cls, iterable=()): return tuple(iterable) def __class_getitem__(cls, params): return cls

IntTuple：元组型形状注解的表面类型。静态上 Pyrefly 将其视为整体形状（运行时表示是tuple[int, ...]）；运行时调用它会强制把任意可迭代对象转成普通 tuple（见 dsl.py 中 DSL 内部对IntTuple的使用）。__class_getitem__返回cls本身，保证IntTuple[2, 3]这类下标写法求值时不会崩溃。
IntTuples：元素为IntTuple的元组，用于描述"多个形状"的场景（如批量操作的形状列表）。__new__同样做 tuple 化。
Elements：tuple[Unpack[S]]的逆运算——从一个IntTuple中抽取各维度。源码 docstring 给出了关键示例：

Array[[*Elements[S], OUT], DType]

它填补了当前 typing 规范的一个空白：没有TypeVarTuple就无法标准地分解一个可变形状。运行时实现很巧妙：__class_getitem__返回cls(shape)实例，且__iter__产出self，从而支持*Elements[S]的解包语法而不崩溃。

3.2 符号维度：Int / IntVar / TypeVarTuple

class Int[T]: """Symbolic integer type for dimension values.""" pass

Int[T]：符号整数维度的泛型标记类，使用 PEP 695 语法。静态形状推导完全交给.pyi桩，运行时只是空壳类。

IntVar是本文档强调的另一个关键原语，README 示例IntVar("B")即指它。源码实现非常考究：

class IntVar: __class__ = typing.TypeVar def __init__(self, name: str, *, bound=None): self.__name__ = name self.name = name self.__bound__ = bound def __add__(self, other): return self def __sub__(self, other): return self def __mul__(self, other): return self def __floordiv__(self, other): return self ...

算术运算符全部返回self：N + 1、N * 2在运行时不会抛TypeError，静态上 Pyrefly 则做真正的符号运算。
__class__ = typing.TypeVar：通过伪装 class，isinstance(x, typing.TypeVar)返回True，从而Generic[N]与TypedDict + Generic[N]的组合都能正常工作。
__typing_subst__与has_default()用于兼容 typing 的替换协议。

配套的TypeVarTuple与此类似：__class__ = typing.TypeVarTuple，并提供__typing_is_unpacked_typevartuple__属性让Generic[*Ns]可用；__iter__产出self，使Tensor[*Ns, 3]中的*Ns星号解包在 Python 中合法。

3.3 运行时断言与装饰器

@defines_assert_shape def assert_shape(actual, shape): if not isinstance(actual, tuple) and hasattr(actual, "shape"): actual_tuple = tuple(actual.shape) # 兼容旧式传数组对象 else: actual_tuple = tuple(actual) expected = tuple(shape) if any(isinstance(dim, SymbolicArithExpr) for dim in expected): if len(actual_tuple) != len(expected): # 符号维度只校验 rank raise AssertionError(...) elif actual_tuple != expected: raise AssertionError(f"expected shape {expected}, got {actual_tuple}") return actual

assert_shape：运行时校验 tuple 形状与期望值一致，Pyrefly 会以类似assert_type的方式做静态形状匹配。
符号维度降级策略：源码 TODO 注释明确指出，当期望形状中含SymbolicArithExpr（符号表达式）时，运行时只校验秩（rank），完整校验交给静态分析。
向后兼容：第一个参数若传的是带.shape属性的数组对象（而非其形状），会被自动转成tuple(actual.shape)。
defines_assert_shape：装饰器，用于标记自定义assert_shape辅助函数，允许用户替换默认实现。

3.4 符号算术：D 与 SymbolicArithExpr

class D: def __new__(cls, value): return SymbolicArithExpr("var", (value,)) def __class_getitem__(cls, value): return cls(value)

D把形状类型变量包装成SymbolicArithExpr，让 Python 能"求值"维度算术。SymbolicArithExpr是 frozen dataclass，重载了__add__/__sub__/__mul__/__floordiv__/__pow__/__neg__及右操作数版本，构造形如SymbolicArithExpr("+", (self, other))的表达式树；__str__负责把表达式树渲染成人类可读形式（如N + 1、-N），复杂子表达式自动加括号（_format_symbolic_arg）。

3.5 广播与广义 ufunc：broadcast / gufunc_broadcast

@type_shape_dsl_function def gufunc_broadcast(spec: str, shapes: IntTuples) -> IntTuple: return _dsl._gufunc_broadcast(spec, shapes) @type_shape_dsl_function def broadcast(left: IntTuple, right: IntTuple) -> IntTuple: spec = "(),()->()" shapes = _dsl.IntTuples((left, right)) return gufunc_broadcast(spec, shapes)

broadcast以 gufunc 签名(),()->()实现两个形状的广播（等价于 NumPy/Torch 的广播语义）。
gufunc_broadcast接受任意广义 ufunc 签名（如(m,n),(n,p)->(m,p)），计算输出形状。
type_shape_dsl_function是运行时 no-op 装饰器，标记用户自定义的类型级 Shape DSL 函数。
index_shape是 Pyrefly 原生形状索引内建的运行时占位，直接返回空IntTuple。
MapIntTuples：将一元类型 lambda 映射到IntTuples值上，支持前向推导（MapIntTuples[lambda S: Tensor[S], tuple[IntTuple[2], IntTuple[3, 4]]]）与参数注解反向推断；运行时__class_getitem__返回tuple，不做参数检查（因为静态合法源如Any、Never无法作为 Python 值映射）。

3.6 静态标记类：Flag / Index / ProxyMethod

Flag（保留字面值的类型级求值标记）、Index（保留索引值用于类型级形状求值）、ProxyMethod（方法转发注解的类型检查标记）三者均为空壳类，纯粹作为静态标记存在，运行时无行为。

四、type-shape DSL：静态推导的内核实现

dsl.py 是"仅在 DSL 定义文件（如torch/_shapes.pyi）内使用，普通桩与用户代码不直接使用"的内部模块，它复用了公共原语并扩展了 DSL 领域：

名称	作用
`Invalid(message)`	返回一个无效的形状计算结果，用于 DSL 体中表示形状不合法
`Int.gradual()`/`IntTuple.gradual()`/`IntTuples.gradual()`	各整数域的"渐进（gradual）"值
`is_concrete_int(value)`	`TypeGuard`判断具体整数（注释特别说明不用`TypeIs`：false 结果不应收窄符号或渐进`Int`值）
`is_int_value(value)`	`TypeIs[int]`判断普通整数值
`concat(left, right, /)`	DSL 内拼接两个形状值
`prod(xs)`/`sum(xs)`	形状的乘积与求和
`einsum(spec, shapes, /)`	按显式 einsum 方程计算输出形状
`_gufunc_broadcast(spec, shapes, /)`	gufunc 广播的实际实现

这些函数共同支撑了三个 stub 仓库中的形状推导：_shapes.pyi分别在 pyrefly-torch-stubs/torch-stubs/_shapes.pyi、pyrefly-jax-stubs/jax-stubs/_shapes.pyi、pyrefly-numpy-stubs/numpy-stubs/_shapes.pyi 中定义形状原语，Pyrefly 的 DSL 编译器会将 DSL 定义文件中的这些内置名识别为 builtin（与 Python 层的定义无关）。

五、框架接入：自动给 torch.Tensor 与 jax.Array 打补丁

为了让Tensor[B, T]在不显式 import 任何额外东西的情况下就能运行，包在导入期主动探测并修补第三方框架：

def _patch_torch_if_available() -> None: try: import torch import torch.nn as nn except ImportError: return subscriptable_classes = [ torch.Tensor, nn.Embedding, nn.Linear, nn.ModuleList, nn.Conv1d, nn.Conv2d, nn.Conv3d, nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d, nn.MaxPool1d..3d, nn.AvgPool1d..3d, nn.AdaptiveAvgPool1d..3d, nn.AdaptiveMaxPool1d..3d, ] for cls in subscriptable_classes: if not hasattr(cls, "__class_getitem__"): cls.__class_getitem__ = classmethod(_return_class) def _patch_jax_if_available() -> None: ... if hasattr(jax, "Array") and not hasattr(jax.Array, "__class_getitem__"): jax.Array.__class_getitem__ = classmethod(_return_class)

关键点：

补丁范围：torch.Tensor、所有常见 Linear/Embedding/ModuleList、Conv、Pool 系列；JAX 侧只补jax.Array。被补丁的类获得__class_getitem__，且_return_class直接返回原类，因此Tensor[B, T]、nn.Linear[In, Out]求值结果就是类本身（no-op）。
幂等性：if not hasattr(cls, "__class_getitem__")保证已具备下标能力的类不被覆盖。
按需加载：try/except ImportError让包在没有 torch/jax 的环境（如纯 NumPy 场景）下也完全可用。

六、TorchScript 兼容：从源码擦除形状注解

这是该包最具工程复杂度的一部分。torchscript.py 的模块 docstring 说明了动机：TorchScript 会从__annotations__读回类属性注解，因此擦除必须发生在任何带注解的类体求值之前；同时它还会修补 TorchScript 的源码加载，把源码中写的形状注解一并擦除——"两个半场缺一不可，模型才能被 script"。

导入shape_extensions.torchscript后自动发生两件事：

6.1 让 Int 下标返回 int

_Int.__class_getitem__ = classmethod(_return_int)

TorchScript 不理解Int[3]，改为返回int，x: Int[3]即成为x: int。

6.2 源码级擦除（AST + 文本双重处理）

_replace_shape_types_in_source用ast解析源码，遍历ast.arg的 annotation、FunctionDef/AsyncFunctionDef的 returns、AnnAssign的 annotation 节点，对每个节点应用_replace_shape_types——后者按_SHAPE_TYPE_REPLACEMENTS表做文本替换：

_SHAPE_TYPE_REPLACEMENTS = ( ("shape_extensions.torchscript.Int[", "int"), ("shape_extensions.Int[", "int"), ("torch.Tensor[", "torch.Tensor"), ("Tensor[", "torch.Tensor"), ("Int[", "int"), )

替换算法是括号配平扫描（_shape_type_end用深度计数找到匹配的]），并通过_has_name_boundary检查标识符边界（如Tensor前不能是字母或点），避免误伤子串。同时：

保留行数：_preserve_line_count用# shape annotation erased注释填充被删除的行，保证行号对齐、避免 TorchScript 的堆栈信息错位。
if TYPE_CHECKING:块处理：_erase_type_checking_blocks识别TYPE_CHECKING与typing.TYPE_CHECKING哨兵（其他名字的相似变量不被识别，避免误删真实运行时行为），将无else的块替换为pass并保留行数。原因在注释中写明：TorchScript 会编译if的两个分支，无法解析TYPE_CHECKING（报 "python value of type 'bool' cannot be used as a value"）。这也让 script 化模型能携带静态assert_type检查。
字符串注解：完全引号包裹的注解（如"Tensor[[B, D]] | None"）会改写引号内容并去掉引号——因为 TorchScript 拒绝字符串形式的 union，却接受不带引号的同一 union；代价是改写注解中的真正前向引用会被提前求值。
私有 torch API 适配：_ShapeAnnotationTypeRemover在ann_to_type侧用torch.jit.annotations._eval_no_call与EvalEnv尽力解析改写结果，失败则原样放行，让 TorchScript 以用户原始代码报出更友好的错误。

接口remove_shape_types_from_torch_sources()也对外公开，可显式调用；两个补丁类均有_is_shape_typing_remover/_is_shape_annotation_type_remover防重入检查。模块 docstring 特别强调：Int是进程级、单向的全局修改，同时影响直接import shape_extensions的代码，且有意不提供撤销途径。

该模块的运行时行为在 test/runtime_tests/test_torchscript_stripper_runtime.py 中有对应测试覆盖。

七、在真实仓库中的使用方式

shape_extensions被 torch、jax、numpy 三个 stub 仓库大量引用，是形状注解体系的公共依赖。以 torch 侧为例：

桩引用：分布在 pyrefly-torch-stubs/torch-stubs/ 下的__init__.pyi、_shapes.pyi、nn/__init__.pyi、nn/functional.pyi、fft.pyi、linalg.pyi等数十个文件。
真实模型验证：pyrefly-torch-stubs/examples/ 下有 bert.py、resnet.py、llama.py、unet.py、nanogpt.py 等 30+ 个真实架构示例，其中 examples/runtime/nanogpt_sym_int_var.py 等展示了IntVar的运行时用法。
运行时测试：test/runtime_tests/test_annotation_runtime.py、test_annotation_runtime_future.py、test_model_runtime.py 直接验证注解求值不崩溃。
集成方式：各 stub 仓库的pyrefly.toml与run_pyrefly.py（如 pyrefly-torch-stubs/run_pyrefly.py）负责把shape_extensions加入搜索路径，run_runtime_tests.py负责运行时测试。

一个典型的最小用法模式是：

from shape_extensions import IntVar, assert_shape from torch import Tensor B = IntVar("B") T = IntVar("T") x: Tensor[B, T] # 运行时 no-op，Pyrefly 静态追踪 B、T assert_shape(x.shape, (B, T)) # 符号维度只查 rank，静态完整校验

八、总结

pyrefly-shape-extensions用不到 500 行的纯 Python 代码，解决了 tensor shape 静态类型系统落地到真实运行时（Python 解释器、TorchScript）的全部摩擦：

双轨设计：.pyi桩承担全部静态推导，.py提供 no-op 运行时，二者 API 严格对齐（init.py）；
兼容性工程：通过__class__ = typing.TypeVar伪装、__iter__配合*解包、__class_getitem__兜底，让Generic[N]、Generic[*Ns]、Tensor[B, T]全部可在纯 Python 中求值；
框架自动接入：导入即修补 torch.Tensor 与 jax.Array 的可下标性（init.py#L47-L101）；
TorchScript 深度兼容：AST + 文本双通道擦除形状注解、处理TYPE_CHECKING块、保留行号（torchscript.py）；
版本同步：与 Pyrefly 主版本 lockstep 发布，Python 3.12+。

这套设计是 Pyrefly Tensor Shape 系统（见 tensor-shapes 目录及 TENSOR_SHAPES_CONTRIBUTING.md）能够在真实 PyTorch/JAX/NumPy 模型上既做精确静态形状检查、又不破坏运行时与 TorchScript 编译的关键支撑。

【免费下载链接】pyreflyA fast type checker and language server for Python项目地址: https://gitcode.com/GitHub_Trending/py/pyrefly

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析