Getting Started with Pyrefly Tensor Shapes

This page walks you through configuring Pyrefly for tensor shape checking and getting your first shape-annotated code running.

Installation

Tensor shape checking requires shape-aware stubs to be available through normal import resolution. Install them from PyPI:

pip install pyrefly-torch-stubs

pyrefly-torch-stubs is a PEP 561 stub-only package: it carries type information for PyTorch, and leaves the runtime torch package alone. PyTorch's own type stubs don't carry shape information, so these stubs take precedence over them and provide shape-aware versions (e.g., nn.Conv2d.__init__ that captures kernel size, stride, and padding as type-level values, and a forward that computes the output spatial dimensions).

Installing it also pulls in pyrefly-shape-extensions, which provides the shape_extensions package. shape_extensions exports Int — the bridge between runtime integer values and type-level symbols. Both packages are versioned in lockstep with Pyrefly.

Tensor shape support is enabled automatically when Pyrefly can resolve the shape_extensions package, so nothing else needs configuring, as long as Pyrefly resolves imports against the Python environment you installed into.

Checking against a local copy of the stubs

The stubs also live in Pyrefly's source tree under tensor-shapes/, which is where to work if you want to read or modify them. Copy that directory into your project and point search-path at the two package directories inside it. The paths are relative to the location of your pyrefly.toml (or of your pyproject.toml, if you configure Pyrefly under [tool.pyrefly]):

search-path = [
    "tensor-shapes/pyrefly-torch-stubs",
    "tensor-shapes/pyrefly-shape-extensions",
]

A copy that lives inside your project is also checked as project code. The stubs use PEP 696 type parameter defaults (class Conv1d[..., S: IntVar = 1]), which Pyrefly parses only under python-version 3.13 or later, so checking them under an earlier version reports:

ERROR Cannot set default type for a type parameter on Python 3.12 (syntax was added in Python 3.13) [invalid-syntax]

Either raise the version Pyrefly checks against. python-version sets the version your code is checked for; it does not change the interpreter your code runs on:

python-version = "3.13"

Or, if your project targets an earlier version and you don't want to change how your own code is checked, exclude the copy from checking instead. Imports still resolve through search-path:

project-excludes = ["tensor-shapes/**"]

Imports and runtime considerations

Python evaluates type annotations at runtime by default. This is a problem for tensor shape annotations because Python's built-in typing.TypeVar doesn't support arithmetic — expressions like D // NHead in an annotation will raise TypeError when the annotation is evaluated. There are two ways to avoid this:

Option 1: from __future__ import annotations (recommended)

Adding this import at the top of the file defers evaluation of all annotations, so shape arithmetic never executes at runtime:

from __future__ import annotations

import torch
import torch.nn as nn
from torch import Tensor
from shape_extensions import Int

This works with both old-style and new-style generics (PEP 695 class Foo[T] syntax).

assert_type during development: You can use assert_type while porting to verify shapes via pyrefly check. Once you're done, remove the assert_type calls — each one corresponds to an IDE inlay type hint that shows the same information permanently. Pyrefly catches shape errors through your function signatures and return types regardless.

Note that assert_type forces evaluation of its type argument, so the file will crash if you try to run it with assert_type calls still present. This is fine — just remove them when the port is complete.

You can also guard Tensor and Int under TYPE_CHECKING if you prefer to keep shape imports invisible at runtime:

from __future__ import annotations
from typing import TYPE_CHECKING

import torch
import torch.nn as nn

if TYPE_CHECKING:
    from torch import Tensor
    from shape_extensions import Int

Option 2: shape_extensions.IntVar (runtime-compatible)

If you need annotations to evaluate at runtime (e.g., for runtime shape validation or keeping assert_type in production code), import shape_extensions directly. The package patches torch.Tensor, nn.Conv2d, and other torch classes to accept subscript syntax at runtime without crashing. It also provides a IntVar that supports arithmetic (N + 1 returns self instead of raising TypeError).

Use old-style generics with shape_extensions.IntVar:

from typing import assert_type

import torch
import torch.nn as nn
from torch import Tensor
from shape_extensions import Int, IntVar

N = IntVar("N")
M = IntVar("M")

class Linear(nn.Module):
    def __init__(self, n: Int[N], m: Int[M]):
        ...

PEP 695 new-style generics (class Foo[T]) automatically use typing.TypeVar internally, which doesn't support arithmetic — so this option requires old-style generics.

tip

Try it yourself

Open this example in the Pyrefly sandbox — Experiment with Int arithmetic, TYPE_CHECKING imports, and building typed tensors — no setup required.

Hello world

Here's a minimal example to verify everything works. This example uses Option 1 (from __future__ import annotations) since it's the simplest setup. We skip assert_type — instead, run pyrefly check and use your IDE's inlay type hints to verify shapes.

Create a file hello_shapes.py:

from __future__ import annotations

import torch
import torch.nn as nn
from torch import Tensor
from shape_extensions import Int, IntVar


class TwoLayerNet[InDim: IntVar, HidDim: IntVar, OutDim: IntVar](nn.Module):
    def __init__(
        self,
        in_dim: Int[InDim],
        hid_dim: Int[HidDim],
        out_dim: Int[OutDim],
    ):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hid_dim)
        self.fc2 = nn.Linear(hid_dim, out_dim)

    def forward[B: IntVar](self, x: Tensor[[B, InDim]]) -> Tensor[[B, OutDim]]:
        h = self.fc1(x)          # pyrefly infers: Tensor[[B, HidDim]]
        return self.fc2(torch.relu(h))

Run pyrefly check hello_shapes.py. You should see no errors — pyrefly infers the shapes through the nn.Linear calls.

If you're using an IDE with Pyrefly's language server, you'll see inlay type hints showing the inferred shape of h as Tensor[[B, HidDim]] without needing any assert_type calls.

Inlay hints in action

Here's what inlay hints look like on a real model (NanoGPT). The MLP module shows shapes flowing through linear layers and activations:

The forward method signature shows how x.size() unpacks into typed dimensions:

And the attention module, where view/transpose reshapes for multi-head attention are fully tracked:

The full attention body, including both flash and manual paths: