Shape Types

This page documents the types, functions, and type-level constructs that make up pyrefly's tensor shape type system.

Declaring shape variables and runtime evaluation

PEP 695 syntax (recommended)

Use native PEP 695 type parameters and add from __future__ import annotations at the top of every implementation file that contains shape annotations:

from __future__ import annotations

from torch import Tensor
from shape_extensions import IntTuple, IntVar


def project[Batch: IntTuple, M: IntVar, N: IntVar](
    x: Tensor[[*Batch, M, N]],
) -> Tensor[[*Batch, M * N]]:
    ...

The future import and the bounds serve different purposes and belong together. N: IntVar declares a symbolic integer parameter; the future import prevents Python from evaluating shape arithmetic and bare splats when the function is defined. Batch: IntTuple similarly declares a whole-shape parameter.

Type expressions evaluated at runtime

The future import defers annotations, not ordinary expressions. Python still evaluates a type passed to cast or assert_type, and tools such as typing.get_type_hints or frameworks that inspect annotations may explicitly evaluate deferred annotations. Use runtime wrappers around PEP 695 parameters where those expressions would otherwise fail:

from __future__ import annotations

from typing import cast

from shape_extensions import Elements, IntTuple, IntVar
from torch import Tensor


def double_last[Batch: IntTuple, N: IntVar](value: object) -> Tensor[[*Batch, N * 2]]:
    return cast(Tensor[[*Elements[Batch], IntVar[N] * 2]], value)

Pyrefly checks the wrapped and unwrapped spellings identically:

• IntVar[N] makes arithmetic on a PEP 695 dimension parameter safe to evaluate. At runtime, a bare N * 2 raises TypeError.
• Elements[Batch] makes a whole-shape parameter safe to unpack. At runtime, a bare *Batch raises TypeError.

Importing shape_extensions also patches supported PyTorch classes, including Tensor, nn.Linear, and nn.Conv2d, to accept shape subscripts at runtime. It applies the patch whether PyTorch is imported before or after it, without importing PyTorch itself.

The second argument to assert_type is another ordinary expression. Wrap its parameters as above, or quote the type to avoid runtime evaluation:

assert_type(result, "Tensor[[*Batch, N * 2]]")

Imports under TYPE_CHECKING are appropriate only when no runtime expression or annotation resolver needs those names.

Legacy type variables

For code that cannot use PEP 695 syntax, declare module-level legacy variables. Use IntVar for individual dimensions and the runtime-safe shape_extensions.TypeVar for whole shapes:

from typing import Generic

from shape_extensions import Int, IntTuple, IntVar, TypeVar
from torch import Tensor

N = IntVar("N")
Batch = TypeVar("Batch", bound=IntTuple)


class Doubler(Generic[N]):
    def forward(self, x: Tensor[[*Batch, N]]) -> Tensor[[*Batch, N * 2]]:
        ...

shape_extensions.TypeVar has ordinary legacy TypeVar semantics in Pyrefly and can be unpacked directly when annotations are evaluated. If you instead use typing.TypeVar, write *Elements[Batch] for runtime evaluation. Legacy IntVar objects already support arithmetic such as N * 2.

Int[X]

Int[X] is a type constructor that bridges runtime integer values to type-level symbols. It is defined in the shape_extensions package, alongside the other names used throughout this page: IntVar (the bound for shape type parameters), IntTuple (the bound for whole-shape type parameters), and Elements (splices a whole-shape type parameter into a Tensor argument list).

Basics

Int[X] denotes the type of an integer value whose type-level identity is X. For example:

• Int[5] is the type of the literal 5
• Int[N] (where N is a type variable) is the type of an integer whose value is bound to N at the type level

Int is a subtype of int, so Int values can be used anywhere int is expected. However, the reverse is not true — passing a plain int where Int[X] is expected loses tracking.

Arithmetic

Arithmetic on Int values produces Int results with the corresponding type-level expression:

Expression	Type
a + b where a: Int[A], b: Int[B]	Int[A + B]
a - b	Int[A - B]
a * b	Int[A * B]
a // b	Int[A // B]
a ** n where n is statically known to be non-negative	Int[A ** n]

A symbolic exponent such as b: Int[B] may be negative, so a ** b is currently gradual (Unknown) rather than Int[A ** B].

Caution: int * Int produces a plain int because the int side has no type-level identity, so the result is no longer tracked. Use Int * Int or literal * Int instead.

Int[X] | None

For optional dimensions — parameters that may or may not be present — use Int[X] | None. In the forward method, narrow with if param is not None: to recover Int[X] inside the branch:

class Attention[D: IntVar, RK: IntVar](nn.Module):
    def __init__(self, dim: Int[D], rank_k: Int[RK] | None = None):
        ...

    def forward[B: IntVar, T: IntVar](self, x: Tensor[[B, T, D]]) -> Tensor[[B, T, D]]:
        if self.rank_k is not None:
            # rank_k is Int[RK] here
            ...

Usage patterns

Pattern	Purpose
def __init__(self, dim: Int[D])	Accept a dimension as a constructor parameter
class Model[D: IntVar](nn.Module)	Make a dimension a class-level type parameter
def forward[B: IntVar](self, x: Tensor[[B, D]])	Bind a per-call dimension
self.head_dim = dim // n_head	Compute a derived dimension (Int[D // NHead])

General-purpose IntVar arguments and defaults accept zero and negative integers, which can represent physical-unit exponents. Literal dimensions in tensor shapes still require non-negative values.

Tensor[[D1, D2, ...]]

Tensor with type arguments represents a tensor with a known shape. The type arguments are the dimensions, in order.

Forms

Form	Meaning
Tensor[[3, 4]]	Concrete 2D tensor with shape (3, 4)
Tensor[[B, C, H, W]]	Generic 4D tensor with symbolic dimensions
Tensor[[B, 3 * C, H // 2]]	Dimensions can contain arithmetic expressions
Tensor[[*Bs, D]]	Variadic: any number of leading batch dimensions (Bs: IntTuple)
Tensor (bare)	Shape unknown — tracking gap

Variadic dimensions

Use a type parameter bound to IntTuple, spliced into the shape list with a bare *Bs splat, for dimensions that should be propagated without being enumerated:

def forward[Bs: IntTuple](self, x: Tensor[[*Bs, InDim]]) -> Tensor[[*Bs, OutDim]]:
    ...

Try it in the sandbox →

This accepts any number of leading dimensions (batch, sequence, etc.) and preserves them in the output.

*Elements[Bs] is the explicit spelling of a bare *Bs splat and checks identically. It is required whenever the annotation is evaluated at runtime, since unpacking a bare TypeVar raises TypeError. In stubs, under from __future__ import annotations, or on Python 3.14+, evaluation is deferred, so the bare spelling is fine unless something forces evaluation (such as typing.get_type_hints or a framework that resolves annotations); use *Elements[Bs] anywhere evaluation is possible.

Don't hide known class dims inside variadic params. If the module has a class-level Int D, use Tensor[[*Bs, D]] not folding D into the variadic carrier itself.

.shape and .size()

When x: Tensor[[B, C, H, W]]:

• x.shape has type tuple[Int[B], Int[C], Int[H], Int[W]], which Pyrefly displays as IntTuple[B, C, H, W]
• x.size(0) has type Int[B]
• x.size() has the same type as x.shape

This means you can extract dimensions from tensors and use them to construct new tensors with matching shapes.

assert_type

assert_type(expr, Type) is checked by the type checker: it verifies that expr has exactly the stated type. If the types don't match, the checker reports an error.

h = self.fc1(x)
assert_type(h, Tensor[[B, 512]])  # checked by pyrefly

Use assert_type during development to verify inferred shapes as you port a model. Once the port is complete, remove temporary checkpoints from the model body; each one corresponds to an inlay type hint that your IDE shows permanently. Keep selected assertions in type-checking tests when you want permanent regression guards. Pyrefly catches shape errors through function signatures and return types regardless.

The second argument of assert_type is an ordinary expression that Python evaluates when the call runs, even under from __future__ import annotations. Arithmetic on an unwrapped typing.TypeVar, or a name imported only under TYPE_CHECKING, can therefore fail at runtime. Arithmetic on shape_extensions.IntVar and expressions wrapped as IntVar[N] are designed to evaluate safely. To avoid evaluating the type expression at all, quote it, as in assert_type(h, "Tensor[[B, 4 * D]]").

When to use

• During porting, after key shape-changing operations (reshapes, convolutions, matmuls)
• As regression guards for complex shape computations

In practice, pyrefly shows inferred shapes as inlay type hints in your editor, so you can verify shapes visually. Use assert_type at key checkpoints where you want a permanent regression guard.

reveal_type

reveal_type(expr) prints the inferred type of expr during type checking. Use it to understand what pyrefly infers before writing assert_type:

h = self.fc1(x)
reveal_type(h)  # Revealed type: Tensor[[B, 512]]

Replace reveal_type with assert_type once you know the expected type.

Type-level arithmetic

Annotations can contain arithmetic on type parameters and literals:

Expression	Example
Addition	Tensor[[B, C1 + C2, H, W]] — concatenation
Subtraction	Tensor[[B, T, D - 1]]
Multiplication	Tensor[[B, NHead * DK]] — multi-head reshape
Floor division	Tensor[[B, NHead, T, D // NHead]]
Exponentiation	Tensor[[B, C * 2 ** I, H // 2 ** I]]

Simplification rules

The type checker automatically simplifies expressions:

• 2 * C // 2 → C
• (H - 1) * 2 + 2 → H * 2
• (a * b) // b → a (sound for all positive integers)

Known limitations

N * (X // N) does not simplify to X — floor division loses the remainder, so the equivalence only holds when X is divisible by N. The checker can't assume this. Common instances:

• Multi-head reassembly: NHead * (D // NHead) — use type: ignore
• BiLSTM output: 2 * (D // 2) — use type: ignore

Annotation hierarchy

When annotating local variables, choose from most to least desirable:

1. assert_type — verifies the checker's inference. Proves the system works, not just that you annotated correctly.
2. Annotation fallback — x: Tensor[[B, C, H, W]] = untracked_op(...). The checker can't infer the shape, but the annotation is compatible. Document WHY.
3. type: ignore — the checker produces a WRONG type (algebraic gap). Last resort. Always include a comment explaining the specific gap.
4. Bare Tensor — shape genuinely unknowable (data-dependent token counts, conditional accumulation). Document the specific reason.