Contributing to Tensor Shape Support

Pyrefly's tensor shape tracking is designed so most PyTorch coverage can be extended by editing stubs and tests, without changing Pyrefly's Rust internals. This page explains the main mechanisms and how to validate changes.

Most external contributions should be stub-only or example/test-only changes. Kernel changes are possible, but they are a narrower workflow for changes to Pyrefly's shape machinery or the shape_extensions runtime package.

Architecture Overview

Shape tracking uses three complementary mechanisms:

Fixture stubs: .pyi files with shape-generic type signatures. These cover modules like nn.Linear, nn.Conv2d, and functions like torch.mm.
Type-level shape DSL functions: shape transforms written in a small Python subset in tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi, decorated with @type_shape_dsl_function, and called directly from public return annotations. These cover operations with computed shape logic like reductions, padding, pooling, and convolution.
Special handlers: Pyrefly implementation logic for constructs that need deeper type system integration, like nn.Sequential chaining, .shape, .size(), assert_shape, and decorator interpretation.

The first two mechanisms live in tensor-shapes/ and are the normal way to add or improve shape coverage. Some stubs still use the older @shape_dsl_function and @uses_shape_dsl(...) mechanism while they are being migrated. Do not add new V1 rules. Special handlers require Pyrefly implementation changes and should be treated as kernel work.

Fixture Stubs

Where They Live

text

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
`-- ...

The tensor-shape test runner passes tensor-shapes/ as a Pyrefly search path, so these stubs override the normal torch stubs during validation.

How Stubs Work

A fixture stub provides a shape-generic type signature. For example, nn.Linear:

python

class Linear[N, M](Module):
    def __init__(
        self,
        in_features: SymInt[N],
        out_features: SymInt[M],
        bias: bool = True,
    ) -> None: ...

    def forward[*Xs](self, input: Tensor[*Xs, N]) -> Tensor[*Xs, M]: ...

The constructor captures input and output dimensions as type parameters. The forward method uses those parameters plus a variadic *Xs for batch dimensions.

Writing a New Stub

Identify the shape signature: input dimensions, output dimensions, and how they relate.
Use SymInt[X] for parameters that determine tensor dimensions. Non-shape parameters like bias and dropout stay as their original types.
Write the method or function signature expressing the shape transform. Use *Xs or *Bs for batch dimensions that pass through unchanged.
Add the stub to the appropriate .pyi file in tensor-shapes/pyrefly-torch-stubs/torch-stubs.
Add or update focused tests under tensor-shapes/pyrefly-torch-stubs/test/.

Example: Adding a New Module

Suppose you want to add nn.GroupNorm, which preserves spatial dimensions:

python

class GroupNorm[NumGroups, NumChannels](Module):
    def __init__(
        self,
        num_groups: SymInt[NumGroups],
        num_channels: SymInt[NumChannels],
        eps: float = 1e-5,
        affine: bool = True,
    ) -> None: ...

    def forward[*S](self, input: Tensor[*S]) -> Tensor[*S]: ...

Since GroupNorm does not change shape, the forward signature is simply Tensor[*S] -> Tensor[*S].

Shape DSL Functions

Use the DSL when a plain signature cannot express the output shape.

Where They Live

DSL functions live in:

text

tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi

Public stubs call a type-level DSL function directly in their return annotation. For example:

python

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 repeat[Shape: IntTuple, Repeats: IntTuple](
    self: Tensor[Shape], *sizes: *Repeats
) -> Tensor[repeat_shape(Shape, Repeats)]: ...

The DSL Subset

The DSL is intentionally small. Its main value domains are Int for one shape dimension and IntTuple for a complete shape. Runtime configuration values are connected through Flag[...] type parameters on public signatures. IntVar names a symbolic dimension, while Int[N] connects a runtime integer parameter to it; for example, def zeros[N: IntVar](n: Int[N]) -> Tensor[[N]]: .... The body language supports common shape computations, including:

dsl.IntTuple(...) to construct result shapes
len, indexing, slicing, and bounded generator expressions
Arithmetic such as +, -, *, //, and %
if / else
Single-assignment local variables
Direct returns from other @type_shape_dsl_function helpers
DSL operations such as dsl.concat, dsl.prod, dsl.Invalid, and gradual Int or IntTuple results

Keep DSL functions simple and algebraic. They are analyzed by Pyrefly; they are not normal runtime implementations of PyTorch operations.

The type-level DSL used by the NumPy and JAX stubs is a separate, smaller subset, and it is still being built out. Two things about it are worth knowing before writing one, because neither is guessable:

A DSL function that uses an unsupported construct evaluates to Unknown at every call site, and the call site itself reports nothing. Type check the stub files to see the real diagnostic; the runner does this for you as the stubs suite.
A parameter typed int | tuple[int, ...] cannot be iterated after narrowing with is_int_value alone. Leading with an is None check makes the narrowing work, so such parameters are declared int | tuple[int, ...] | None with a body that rejects None. Both conv_shape in the Torch stubs and reshape_shape in the JAX stubs do this.

Example: reduction

python

@type_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
    )

The public stub binds its input shape and runtime options to type parameters, then calls reduce_shape(...) in the return annotation.

Adding a New DSL Function

Write the shape transform in tensor-shapes/pyrefly-torch-stubs/torch-stubs/_shapes.pyi.
Decorate it with @type_shape_dsl_function.
Bind the relevant public arguments with IntVar, IntTuple, or Flag[...] type parameters and call the DSL function from the return annotation.
Add positive tests that use assert_type to check the computed shape.
Add negative tests with # E: expectations if the DSL should reject invalid shapes or report shape errors.

The older decorator-based DSL remains only for rules that have not yet been migrated. Avoid combining V1 and V2 logic in new rules; if V2 cannot yet express the operation, document the gap rather than adding new V1 surface area.

Current Limitations

The type-level DSL uses a small, composable language and does not model every shape behavior precisely. Current known limitations include symbolic arange rounding, symbolic configuration values for unfold and diag_embed, structured tensordot axis lists, products of symbolic-rank shapes or derived symbolic dimensions, and list-based padding. Keep these cases gradual where necessary, add focused tests, and leave a TODO(stroxler) at the affected rule so the loss of precision remains visible.

Ported Models

Where They Live

text

tensor-shapes/pyrefly-torch-stubs/examples/

Each file is a fully annotated port of a real-world PyTorch model with assert_type checkpoints and smoke tests.

Adding a New Model

Choose a model from TorchBench or another source.
Port it using the tutorials or the agent skill.
Add assert_type or assert_shape checkpoints after shape-changing operations.
Add smoke tests at the bottom of the file when runtime execution is useful.
Run verify_port.sh to check for common quality issues.

`verify_port.sh`

This script checks a ported model for common issues:

bash

tensor-shapes/skills/add-shape-types-to-torch-model/verify_port.sh tensor-shapes/pyrefly-torch-stubs/examples/<model>.py

It reports:

Metric	Description
`ig`	`type: ignore` count
`bs`	Bare `Tensor` in signatures
`bv`	Bare `Tensor` in variable annotations
`sh`	Shaped `assert_type` count
`ba`	Bare `assert_type` count
`sm`	Smoke test count

Testing Stub and Example Changes

For most contributions, the important validation is the tensor-shape Pyrefly runner. It checks the focused tests, negative expectations, jaxtyping examples, and the example corpus using the shape-aware stubs.

It also type checks the stub files themselves, reported as a stubs suite. This matters more than it sounds: Pyrefly reports errors only for the files it is asked to check, so a stub reached through --search-path is silent. A stub that fails to compile does not announce itself, it just stops contributing types, and every call site quietly infers Unknown -- which looks like a missing rule rather than a broken one. Checking the stubs directly turns that into an error with a line number.

The Torch package opts out for now, via check_stubs=False in its run_pyrefly.py. Most of its errors are in torch-stubs/_shapes.pyi, whose V1 @shape_dsl_function bodies are not valid Python. Type-level DSL files do check cleanly, so migrating those rules is what removes the opt-out.

Build Pyrefly first, then run:

bash

cargo build
python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py

If your build uses a custom target directory, run_pyrefly.py respects CARGO_TARGET_DIR. You can also pass the binary explicitly:

bash

python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --pyrefly /path/to/pyrefly

Run a single suite while iterating:

bash

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-examples

Use --nocapture when you want the full Pyrefly output on success. By default, the runner prints a compact PASS ... line and only dumps checker output on failure.

There are no Buck test targets for the stubs. An internal checkout runs the same runner and only sources Pyrefly differently, via --buck:

bash

python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --buck

To run every library at once, static and runtime, exactly as both CI systems do:

bash

python3 tensor-shapes/run_tests.py           # add --buck in an internal checkout
python3 tensor-shapes/run_tests.py --static-only   # no virtualenv needed

The project-level test.py runner keeps tensor-shape validation separate from the default Pyrefly test loop. To run just these validations through test.py:

bash

python3 test.py --no-fmt --no-lint --no-test --tensor-shapes --no-conformance --no-jsonschema

Runtime Tests

Runtime tests validate that the annotation helpers and runnable example models behave correctly in Python, not just in Pyrefly's static checker.

The tests live in:

text

tensor-shapes/pyrefly-torch-stubs/test/runtime_tests/

Runtime tests need the shared virtualenv, which serves torch, numpy and jax together. Bootstrapping is the only step that downloads anything, so it is also the only step that needs network access -- on a Meta machine, via fwdproxy:

bash

python3 tensor-shapes/bootstrap_venv.py            # add --fwdproxy internally
python3 tensor-shapes/run_tests.py --runtime-only

The virtualenv defaults to ~/.tensor-shapes-venv; set $TENSOR_SHAPES_VENV to put it elsewhere. The runners never create it, and never reach the network: if it is missing they say so and print the bootstrap command. Type checking does not need it at all.

Run one suite while iterating:

bash

python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py --suite annotation
python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py --suite model

The runtime runner sets up import paths for shape_extensions and the runnable example modules. Runtime tests are the same in an internal checkout: they run against the virtualenv, never through Buck, so that no workflow ever rebuilds torch, numpy or jax.

Kernel Tests

Most contributors should not need this section. Use these tests when you change Pyrefly's tensor-shape kernel rather than only stubs or examples. Kernel changes include:

shape_extensions primitives or decorators
assert_shape type-checker behavior
@shape_dsl_function parsing, validation, or evaluation
@uses_shape_dsl handling
special handlers in Pyrefly's Rust source

The focused Pyrefly unit tests live in:

text

pyrefly/lib/test/shape_dsl.rs

Run them with Cargo:

bash

cargo test shape_dsl

In an internal Buck checkout:

bash

buck test pyrefly:pyrefly_library -- shape_dsl

Kernel tests are intentionally much smaller than the stub/example suites. They cover the core primitives and invariants; the tensor-shape stub tests stress the DSL through realistic PyTorch signatures.

Pre-Commit Checks

Before handing off changes, run formatting and linting:

bash

./test.py --no-test --no-tensor-shapes --no-conformance --no-jsonschema

Also run the relevant tensor-shape checks for the files you touched:

Stub/test/example changes: python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py
Runtime helper or runnable model changes: python tensor-shapes/pyrefly-torch-stubs/run_runtime_tests.py
Kernel changes: cargo test shape_dsl or the Buck equivalent above