Skip to main content

Tensor Shapes

Experimental

Tensor shape support is experimental. We aim to provide a mostly stable experience for users who consume the shape-aware stubs, but behavior may still change between releases. Library and stub authors interact more directly with the shape DSL and should expect that authoring API to evolve.

Pyrefly can track the shapes of arrays and tensors through NumPy, JAX, and PyTorch code, giving you end-to-end static type checking of shape transformations.

Pyrefly publishes shape-aware stub libraries for all three: pyrefly-numpy-stubs, jax-pyrefly-stubs, and pyrefly-torch-stubs. A stub library is a package of type declarations that describes another library's API to a type checker without changing the library at runtime. These stubs teach Pyrefly how array operations transform shapes.

This page focuses on the benefits and high-level approach. Want to jump straight in? Use Getting Started to try it locally, explore the sandbox-backed numerical example below, or open the sandbox-backed first tutorial.

Why track shapes?​

When you write array code, the hardest thing to keep track of is often shapes. Every operator transforms shapes in non-trivial ways, and shape mistakes don't always crash — they can silently produce wrong results. With Pyrefly's shape support, you get automatic inlay type hints showing shapes as you type, so you can see exactly what shape every intermediate value has without running your code or adding print statements.

Consider a small linear-model workflow: compute predictions for 1,000 observations, form a Gram matrix, and join the predictions back onto the features. Without running the code, Pyrefly tracks every array shape:

features: Tensor[[1000, 3]] = torch.randn((1000, 3))
coefficients: Tensor[[3, 1]] = torch.randn((3, 1))


predictions: Tensor[[1000, 1]] = features @ coefficients
gram_matrix: Tensor[[3, 3]] = features.transpose() @ features
results: Tensor[[1000, 4]] = torch.concatenate((features, predictions), axis=1)

Try it in the sandbox →

Pyrefly infers a (1000, 1) shape for predictions, a (3, 3) shape for gram_matrix, and a (1000, 4) shape for the combined results. It also rejects a coefficient matrix whose leading dimension does not match the three input features.

How it works​

Pyrefly's shape support is built on two extensions that work together:

  1. Symbolic integer arithmetic in the core type system — lets you write Tensor[[B, C, H, W]] and have arithmetic like D // NHead work at the type level.
  2. Shape transform specifications for array operators — libraries of shape rules that tell Pyrefly how NumPy, JAX, and PyTorch operations transform shapes.

With these two extensions, you can use Pyrefly to typecheck real-world array code, where the shapes of local variables are inferred from just a few annotations at class and function boundaries.

Symbolic integer arithmetic​

You can write array types with integer dimensions — jax.Array[[3, 4]] and np.ndarray[[3, 4]] are 2D arrays with shape (3, 4).

Int[X] bridges runtime integer values to the type level. When x: np.ndarray[[3, 4]], then x.shape has type tuple[Int[3], Int[4]] — you can extract dimensions from arrays and use them to construct new ones. Arithmetic works too: if a: Int[3] and b: Int[4], then a * b: Int[12].

Generic type parameters let you write shape-polymorphic JAX functions:

import jax

def project[Batch: IntVar, In: IntVar, Out: IntVar](
    features: jax.Array[[Batch, In]],
    weights: jax.Array[[In, Out]],
) -> jax.Array[[Batch, Out]]:
    return features @ weights

features: jax.Array[[2, 3]] = ...
weights: jax.Array[[3, 4]] = ...
projected: jax.Array[[2, 4]] = project(features, weights)

Arithmetic on type variables works in NumPy shapes too:

import numpy as np

def midpoint_vector[A: IntVar, B: IntVar](
    a: Int[A], b: Int[B]
) -> np.ndarray[[(A + B) // 2]]:
    return np.zeros((a + b) // 2)

x: np.ndarray[[3]] = midpoint_vector(2, 4)

While these examples use type annotations for exposition, the types of local variables like projected and x are inferred automatically — Pyrefly shows them as inlay type hints in your editor.

Shape transform specifications​

Some array operators have simple shape signatures that can be expressed as standard type stubs. For example, JAX's outer:

def outer[M: IntVar, N: IntVar](
    a: jax.Array[[M]], b: jax.Array[[N]]
) -> jax.Array[[M, N]]:
    ...

For operators with shape logic that depends on runtime arguments, the stubs use a small DSL. For example, NumPy's expand_dims inserts a dimension at an axis chosen by the caller:

# Abridged from pyrefly-numpy-stubs/numpy-stubs/_shapes.pyi
@type_shape_dsl_function
def expand_dims_shape(shape: IntTuple, axis: int) -> IntTuple:
    output_rank = len(shape) + 1
    if axis < 0 - output_rank or axis >= output_rank:
        return dsl.Invalid("expand_dims axis out of bounds")
    if axis < 0:
        norm_axis = axis + output_rank
    else:
        norm_axis = axis + 0
    return dsl.IntTuple(
        (
            1
            if index == norm_axis
            else shape[index if index < norm_axis else index - 1]
            for index in range(output_rank)
        )
    )

This means you can extend Pyrefly's shape coverage for new array operators without touching Pyrefly's internals — see the contributing guide for details.

Comment-only shape annotations​

Shared stubs can keep their ordinary annotations for other type checkers and add Pyrefly shape information in comments. Enable this experimental syntax with # enable-shape-type-comments in the module header, before any statement other than the module docstring:

# enable-shape-type-comments
import numpy as np

# shapes: M, N
def transpose(
    n: int,  # shaped: M
    x: np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: [0] = [M, N]
) -> np.ndarray[tuple[int, int], np.dtype[np.float64]]:  # shaped: [0] = [N, M]
    ...

# shapes: M, N, *Batch declares scalar dimensions and variadic shapes on the following function or class. It uses the same declaration rules as @shape_vars; starred names have an IntTuple bound. A definition cannot combine the comment with a shape-declaration decorator. Class dimensions have gradual defaults.

A trailing # shaped: payload uses the compatibility Shaped semantics: on int, a scalar expression becomes Int[payload]; on a generic class, the payload supplies additional type arguments. For example, Tensor # shaped: [M, N] supplies a shape to the shape-aware tensor stubs. Classes without type parameters ignore the payload, preserving compatibility with ordinary library types.

Use # shaped: = payload to replace the attached type component, or # shaped: [0] = payload to replace its first written type argument. Selectors are zero-based unsigned decimal indices and can be chained, such as [0][1]. They require a following =: without it, [M, N] and [0] are ordinary shape payloads, not selectors. type =, .N, and .args[N] remain accepted aliases. Replacing an integer tuple carrier with a shape retains its rank and literal-dimension agreement checks. For non-carrier types, an explicit replacement supplies a complete type expression instead of appending arguments. For example, size: int # shaped: = Int[M] explicitly replaces the integer annotation. Use the concise size: int # shaped: M for the common scalar case.

Payloads support the full compatibility-mode expression language, including arithmetic, variadic shapes, and calls to imported shape DSL helpers. Shape names are scoped to the declaration; other names resolve through ordinary annotation scopes. Imports used only in comments may need an unused-import suppression for external linters.

The supported positions are parameter annotations, return annotations, and annotated assignments. A multiline annotation can carry separate comments on its type components. The comment applies to the outermost complete component ending on its line; a following comma is a separator, not part of the type. For example:

# shapes: M, N
def result(
    x: np.ndarray[
        tuple[int, int],  # shaped: = [M, N]
        np.dtype[np.float64],
    ],
) -> tuple[
    int,  # shaped: M
    np.ndarray[tuple[int, int], np.dtype[np.float64]],  # shaped: [0] = [N, M]
]:
    ...

A selector is relative to the component beside the comment, not the enclosing annotation. Put one component on each comment-bearing line. Multiple comments in an annotation must target disjoint components; overlapping parent/child or duplicate targets are rejected rather than applied in a particular order. Comments on union branches are supported by splitting the union across lines. A trailing comment on the line ending a return annotation belongs to that return, including a one-line stub definition. Selectors work in quoted components, but do not expand type aliases or select into union operators, literal values, annotation metadata, existing Shaped wrappers, or list-valued, starred, and ellipsis type arguments. Prefix comments for local annotations and IDE navigation inside comment payloads are not supported.

Without the header, shape comments have no effect and produce no diagnostics. With it, malformed or unattached directives are reported as invalid annotations, including directives in statically unreachable code. Diagnostics on standalone comment lines currently cannot be suppressed by a trailing ignore on that line. This syntax does not enable legacy # type: annotations.

Understand array code at a glance​

Shape information does more than catch bugs: it makes unfamiliar array code easier to read. In this NanoGPT forward pass, Pyrefly's inlay hints show the batch, sequence, and embedding dimensions as they flow through position broadcasting and every transformer block:

tok_emb: Tensor[[B, T, NEmbedding]] = self.transformer.wte(idx)
pos_emb: Tensor[[T, NEmbedding]] = self.transformer.wpe(pos)
x: Tensor[[B, T, NEmbedding]] = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
    x: Tensor[[B, T, NEmbedding]] = block(x)
x: Tensor[[B, T, NEmbedding]] = self.transformer.ln_f(x)

The hints provide information that would otherwise require shape comments, but they stay synchronized with the code as it changes. Only the shapes at the model boundary need annotations; Pyrefly infers the intermediate shapes from the shape-aware PyTorch stubs.

For comparisons with other shape-typing libraries and earlier type-checker experiments, see Related Work. Projects that already use jaxtyping can opt into Pyrefly's static jaxtyping support with @static_jaxtyping.

Get involved​

Tensor shapes are under active development. We welcome contributions — especially new fixture stubs and DSL specifications for NumPy, JAX, and PyTorch operators. See the contributing guide to get started.

Join the conversation: