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. Shipped stubs use the type-level DSL exclusively. Special handlers require Pyrefly implementation changes and should be treated as kernel work.

Fixture Stubs

Where They Live

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 the Torch stub package root and shape_extensions as Pyrefly search paths, plus the shared virtualenv as a site-package fallback. The partial overlay therefore supplies shape-aware symbols while undeclared Torch modules resolve from the installed library.

How Stubs Work

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

class Linear[IN: IntVar, OUT: IntVar](Module):
    def __init__(
        self,
        in_features: Int[IN],
        out_features: Int[OUT],
        bias: bool = True,
    ) -> None: ...

    def forward[Bs: IntTuple](
        self, input: Tensor[[*Elements[Bs], IN]]
    ) -> Tensor[[*Elements[Bs], OUT]]: ...

The constructor captures input and output dimensions as type parameters bound by IntVar. The forward method uses those parameters plus an IntTuple-bound parameter, unpacked with Elements[...], for the batch dimensions.

Writing a New Stub

Identify the shape signature: input dimensions, output dimensions, and how they relate.
Use Int[X], with X bound by IntVar, 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 an IntTuple-bound parameter with a bare *Bs splat for batch dimensions under deferred evaluation; use *Elements[Bs] only when annotations evaluate eagerly at runtime.
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:

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

    def forward[Shape: IntTuple](self, input: Tensor[Shape]) -> Tensor[Shape]: ...

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

Shape DSL Functions

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

Where They Live

DSL functions live in:

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

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

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. For einops-like APIs, NamedInts with CaptureNamedInts[...] carries named axis lengths from **kwargs into the DSL. 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, and dsl.Invalid
dsl.Int.gradual() for a gradual dimension expression, and a direct return dsl.IntTuple.gradual() for a gradual shape

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

Integer Helper Arguments

Declare a helper parameter as Int when it consumes one dimension, or as Int | None when None has a distinct meaning. For a public runtime integer that passes through to a helper, prefer a type parameter bound exactly by shape Int and annotate the runtime parameter with that type parameter:

@type_shape_dsl_function
def resize_shape(size: Int) -> IntTuple:
    return dsl.IntTuple((size,))

def resize[N: Int](size: N) -> Tensor[resize_shape(N)]: ...

This form requires the bound to be exactly Int. Use IntVar instead when a type parameter names a symbolic dimension in direct IntTuple or list shape syntax. If that symbolic dimension is also passed to a DSL helper, wrap it with Int[...] at the call boundary:

def zeros[N: IntVar](n: Int[N]) -> Tensor[[N]]: ...
def resize_symbolic[N: IntVar](size: Int[N]) -> Tensor[resize_shape(Int[N])]: ...

Raw IntVar arguments and arithmetic such as N + 1 are rejected in helper calls; write Int[N] and Int[N] + 1 instead. Int[N] | None written directly as an argument is type-union syntax, not a runtime DSL value. Pass Int[N], None, or a type parameter whose bound is exactly Int | None, as appropriate. An argument that still resolves to Int | None is accepted as gradual until control flow narrows it; the non-None branch can then use it as an Int.

A broad runtime int used as a dimension becomes a gradual dimension, which preserves known rank and other dimensions. Any remains unknown rather than being treated as a gradual integer. dsl.Int.gradual() is itself an Int expression, so it can participate in arithmetic and dsl.IntTuple(...) construction. dsl.IntTuple.gradual() currently represents a whole gradual shape only as a direct DSL-function return; it cannot be assigned to a local or embedded in a larger expression.

IntVar[...] is the compatibility wrapper for dimension arithmetic that Python would otherwise evaluate eagerly: IntVar[N] + 1 checks as N + 1 while shielding the arithmetic from runtime evaluation. It does not replace Int[...]. The call form IntVar(...) is the legacy TypeVar constructor, not a wrapper, and is rejected in dimension position.

Use dsl.is_concrete_int(value) with an Int or Int | None value when a branch requires an integer literal known during shape evaluation; it is false for None, symbolic dimensions, and gradual Int values. Use dsl.is_int_value(value) to narrow the integer member of a compatible Flag[int | tuple[int, ...] | None] value. That predicate does not prove the integer is concrete.

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

@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 Int, 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.

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

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

Each file is an explicitly scoped, shape-annotated port of real-world PyTorch code with assert_type checkpoints and, where useful, smoke tests. Its header should identify the upstream revision, included files/configuration/modes, and any deliberately omitted wrapper or runtime path.

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 line-oriented script is an advisory check for common issues:

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

Its counts are hints, not authoritative coverage metrics: multiline calls and signatures can evade or confuse the shell patterns. Use the explicit port ledger and Pyrefly results as the source of truth.

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 can also type check stub files directly, but the Torch package currently opts out via check_stubs=False in run_pyrefly.py because of the concrete issues listed there. Search-path imports do not report errors from the stub files unless they are direct check targets.

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

The runner builds Pyrefly itself, with cargo build by default and with Buck under --buck, so it always checks against your working copy. Building separately first is unnecessary, and skipping the build is what makes a run report results from an older Pyrefly.

If your build uses a custom target directory, run_pyrefly.py respects CARGO_TARGET_DIR. Passing a binary explicitly is the one mode that does not build, since a bare path says nothing about how to rebuild it:

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

Suite names are generated from test/test_*.py; discover the current list with run_pyrefly.py --help. Examples include:

python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-joining
python3 tensor-shapes/pyrefly-torch-stubs/run_pyrefly.py --suite torch-einsum
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:

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:

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

The Torch and NumPy shape stubs fall back to definitions from the installed libraries, so even --static-only needs the shared virtualenv. Pass --python to select a different virtualenv interpreter with the required libraries installed; the root runner uses it for runtime tests and forwards it to Torch and NumPy static checking.

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:

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.

Every test must call assert_shape at least once, so that a test cannot pass vacuously. A bare assert x.shape == (...) does not count, because the runner cannot see it, and a test that asserts no shapes fails rather than passing.

assert_shape(x.shape, shape) verifies the runtime shape and the statically inferred shape together. The positional shape is always the shape Pyrefly is expected to infer. Where the library produces a different one, pass it as runtime=; only the runtime check uses it. Two situations need it:

An expression Pyrefly infers gradually. Write the expected shape as a bare IntTuple when it has no shape at all, or as a tuple such as (int,) when the rank is known and only a dimension is not. A bare IntTuple holds only when nothing was inferred, so it cannot quietly paper over a known shape.
A known bug, where Pyrefly infers a shape the library does not produce. Recording it makes the discrepancy visible and makes the test fail once the inferred shape changes, instead of leaving it undocumented. Unlike a shape annotation, the expected shape accepts a degenerate dimension, because the point is to record what Pyrefly currently infers.

Reach for runtime= only when the shapes really differ. Without it one call pins both behaviors, which is what most tests want. Add a TODO next to it when the gradual result is expected to become exact; some cannot, and saying which is which is the useful part.

The tests live in:

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

Runtime tests and static fallback checks 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:

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. Torch and NumPy static checking use installed library definitions; JAX static checking does not need the virtualenv.

Run one suite while iterating:

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
special handlers in Pyrefly's Rust source

The focused Pyrefly unit tests live in:

pyrefly/lib/test/shape_dsl.rs

Run them with Cargo:

cargo test shape_dsl

In an internal Buck checkout:

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

Python files in the tensor-shape packages use Ruff's formatter, rather than Black. Format them from the repository root with the same Ruff version as CI:

uv tool run --from ruff==0.16.5 ruff format \
  tensor-shapes

The skills directory is documentation rather than corpus source and is excluded because Ruff also formats Python snippets embedded in Markdown.

Before handing off changes, also run the repository formatting and linting:

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

Navigation Menu

Repository navigation

TENSOR_SHAPES_CONTRIBUTING.md

TENSOR_SHAPES_CONTRIBUTING.md

Contributing to Tensor Shape Support

Architecture Overview

Fixture Stubs

Where They Live

How Stubs Work

Writing a New Stub

Example: Adding a New Module

Shape DSL Functions

Where They Live

The DSL Subset

Integer Helper Arguments

Example: reduction

Adding a New DSL Function

Current Limitations

Ported Models

Where They Live

Adding a New Model

`verify_port.sh`

Testing Stub and Example Changes

Runtime Tests

Kernel Tests

Pre-Commit Checks

FilesExpand file tree

TENSOR_SHAPES_CONTRIBUTING.md

Latest commit

History

TENSOR_SHAPES_CONTRIBUTING.md

File metadata and controls

Contributing to Tensor Shape Support

Architecture Overview

Fixture Stubs

Where They Live

How Stubs Work

Writing a New Stub

Example: Adding a New Module

Shape DSL Functions

Where They Live

The DSL Subset

Integer Helper Arguments

Example: reduction

Adding a New DSL Function

Current Limitations

Ported Models

Where They Live

Adding a New Model

verify_port.sh

Testing Stub and Example Changes

Runtime Tests

Kernel Tests

Pre-Commit Checks

`verify_port.sh`