用途:修改 Pyrefly 的张量形状 DSL,即根据 torch 算子的输入形状计算输出形状的逻辑。适用于 Pyrefly 计算出错误张量形状,或缺少无法在 stub 签名中表达的形状时。
主要能力与结构:形状规则是 tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi 中带 @type_shape_dsl_function 装饰器的 Python 函数,用受限 Python 子集计算类型级值;公共 stub 在返回注解中直接调用,如 Tensor[reshape(Shape, Target)]。改动分两类:仅 stub 改动编辑 _shapes.pyi 与公共返回注解,组合已有操作;DSL 内核改动编辑 Rust 校验器或求值器以新增操作,仅在无法组合现有 DSL 时使用。类型级 DSL 实现主要在 crates/pyrefly_types/src/type_level_dsl.rs,符号维度代数在 crates/pyrefly_types/src/dimension.rs。
用法要点:接受整数元组或列表的 stub 参数用 IntTupleOrList[Values] 配合 Values: IntTuple。数值公式中若结果需保持张量,应显式使用浮点,如指数乘 1.0 或用 2.0 而非 2。渐进形状规范写法:int 表示渐进维度,IntTuple 表示渐进整体形状,裸 Int 表示渐进形状整数。
测试与验证:仅 stub 改动应在该库静态形状语料加聚焦测试并尽量做运行时交叉检查,不要复制到 pyrefly/lib/test/shape_dsl.rs;DSL 内核改动应在 pyrefly/lib/test/shape_dsl.rs 加针对性测试,覆盖相关代数与边界情况,可用 assert_type 与 # E: 标记。运行内核测试:buck test fbcode//pyrefly:test-library -- <test_name> 或 cargo test <test_name>。Rust 改动后需重建:buck build fbcode//pyrefly:pyrefly 或 cargo build;仅 stub 改动无需重建。修改形状语义的默认验证门为 tensor-shapes/run_all_shape_tests.py,可用 --mode buck/cargo 固定后端,--include-runtime-tests 仅在需要运行时覆盖时添加。
限制:本技能指向代码而非复制代码,细节需阅读上述文件;必须遵守不变量并添加单元测试。贡献方式:fbsource 以 diff 落地,clone 则就地向上游 stub/Rust 源码提 PR。