Skip to main content

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.

Install​

Pyrefly's basic mode—the default when there is no type-checker configuration—does not show shape errors. Set python.pyrefly.typeCheckingMode to default or strict, or create a pyrefly.toml (for example, with pyrefly init) to see them.

PyPI. Install the package for your library into the Python environment Pyrefly uses: pyrefly-numpy-stubs, pyrefly-torch-stubs, jax-pyrefly-stubs, or pyrefly-einops-stubs (for PyTorch).

pip install pyrefly-numpy-stubs

The stub packages install pyrefly-shape-extensions and enable shape inference automatically.

Local stubs. To use unreleased stubs, copy tensor-shapes/ into your project, point search-path at the library and shape-extension packages, and exclude the source copy from checking. For PyTorch:

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

Start adding shape annotations​

Pick the example for your array library, save it as hello_shapes.py, and run pyrefly check hello_shapes.py. A clean check proves that Pyrefly found the shape-aware stubs and inferred the result shape expected by assert_type.

from __future__ import annotations

from typing import assert_type

import numpy as np


def multiply(
    left: np.ndarray[[2, 3], np.dtype[np.float64]],
    right: np.ndarray[[3, 4], np.dtype[np.float64]],
) -> None:
    result = left @ right
    assert_type(result, np.ndarray[[2, 4], np.dtype[np.float64]])

These examples use native PEP 695 type parameters with from __future__ import annotations. Pyrefly also supports runtime wrappers and legacy type variables; see Declaring shape variables and runtime evaluation for the available spellings.

Use an agent on real code​

For most real PyTorch projects, the next step is to have a coding agent annotate an existing model. The Agent Skill for Tensor Shape Porting gives the agent a workflow for adding and verifying those annotations.

See shapes with inlay hints​

Inferred shapes shown inline can make unfamiliar array code much easier to read. NanoGPT records tensor shapes in comments around its multi-head attention code; inlay hints provide the same context from checked types, so they update with the code instead of going stale.

In Visual Studio Code, open Settings and set Editor > Inlay Hints: Enabled to On. Here is the NanoGPT attention body with Pyrefly's inferred shapes displayed inline:

# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
if self.flash:
    # efficient attention using Flash Attention CUDA kernels
    y: Tensor[[B, NHead, T, NEmbedding // NHead]] = torch.nn.functional.scaled_dot_product_attention(
        q,
        k,
        v,
        attn_mask=None,
        dropout_p=self.dropout if self.training else 0,
        is_causal=True,
    )
else:
    # manual implementation of attention
    att: Tensor[[B, NHead, T, T]] = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
    mask: Tensor[[1, 1, T, T]] = self.bias[:, :, :t, :t] == 0
    att: Tensor[[B, NHead, T, T]] = att.masked_fill(mask, float("-inf"))
    att: Tensor[[B, NHead, T, T]] = F.softmax(att, dim=-1)
    att: Tensor[[B, NHead, T, T]] = self.attn_dropout(att)
    y: Tensor[[B, NHead, T, NEmbedding // NHead]] = att @ v  # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)


y: Tensor[[B, T, NEmbedding]] = (
    y.transpose(1, 2).contiguous().view(b, t, c)
)  # re-assemble all head outputs side by side


# output projection
y: Tensor[[B, T, NEmbedding]] = self.resid_dropout(self.c_proj(y))
return y