Foundations of Pyrefly's Tensor Shape Types

[stroxler]

Introduction

We've rewired Pyrefly's tensor shape type system around composability. The new
system is not coupled to torch, can attach shape information to classes that
produce or operate on arrays without being arrays themselves, and could extend
to integer-valued properties beyond shapes. For example, the CuTe DSL treats
both shapes and strides as fundamental to a tensor layout. The system is also
built from relatively simple concepts that we think could eventually be
specified in a PEP.

This post covers:

• A discussion of the current type system, focusing on the primitives for int
  arithmetic, shape representation, and our transformation DSL
• A short history of how we got here, starting from the prototype that
  Avik shared
  at the PyCon Typing Summit in May 2026

Elements of the shape type system

Symbolic arithmetic

The basic building block of the shape type system is symbolic integer arithmetic.
An IntVar is a type variable whose values are integers rather than types, and
Int[N] is the type of an integer whose type-level value is N.

For example, we can write a function that adds two dimensions:

from __future__ import annotations

from shape_extensions import Int, IntTuple, IntVar

class Array[Shape: IntTuple]: ...

def sum_dims[N: IntVar, M: IntVar](
    a: Array[[N, M]],
) -> Array[[N + M]]:
    ...

If a has type Array[[2, 3]], Pyrefly binds N to 2 and M to
3, so the return type is Array[[5]]. The same function also works when
one or both dimensions are still symbolic:

from typing import assert_type

def apply_sum[N: IntVar, M: IntVar](
    x: Array[[N, M]],
    y: Array[[N, 3]],
    z: Array[[2, 3]],
) -> None:
    assert_type(sum_dims(x), Array[[N + M]])
    assert_type(sum_dims(y), Array[[N + 3]])
    assert_type(sum_dims(z), Array[[5]])

We can also capture integer values and use them in shapes:

def repeat[N: IntVar, Copies: IntVar](
    x: Array[[N]], copies: Int[Copies]
) -> Array[[N * Copies]]:
    ...

def concrete_repeat(x: Array[[10]]) -> None:
    assert_type(repeat(x, 3), Array[[30]])

The literal argument 3 is captured as Int[3], so N * Copies evaluates to
30.

Pyrefly normalizes symbolic expressions and knows that, for example:

• N + M == M + N
• N + 0 == N
• N + N == 2 * N

Pyrefly can substitute known values and compare normalized expressions, but it
is not a general integer constraint solver and cannot solve arbitrary
comparisons like an SMT solver.

In these examples, IntVar is a new kind of type variable, analogous to
ParamSpec and TypeVarTuple. It is allowed to appear in different places
(inside Int, which represents a symbolic integer, or IntTuple, which
represents a tuple of symbolic integers). It supports arithmetic which
TypeVar does not; if and when we make a PEP, we'll need a new runtime class
to enable this (similar to how TypeVarTuple allows the *Ts syntax whereas
TypeVar does not).

IntTuple and shapes

Symbolic arithmetic describes individual dimensions; to describe shapes, we
need to collect those dimensions. IntTuple is a tuple of symbolic integers,
commonly used to represent a shape. A fixed-rank form such as
IntTuple[M, int, 6] is equivalent to tuple[Int[M], Int, Literal[6]]: a
symbolic dimension becomes Int[M], an
unknown dimension becomes Int, and a concrete dimension becomes a literal.
Bare IntTuple has gradual rank. Unlike tuple[Int, ...], it is gradually
compatible with fixed-rank IntTuple types.

IntTuple is not tied to tensors, arrays, or any particular class. For example,
we can write:

class Array[Shape: IntTuple]: ...
class DeferredArrayOperation[InputShape: IntTuple, OutputShape: IntTuple]: ...

For readability, shape arguments have a compact list spelling, which is enabled
whenever a type parameter is bound by IntTuple. These two types mean the same
thing:

Array[[2, 3]]
Array[IntTuple[2, 3]]

The contents of an IntTuple[...] are dimensions rather than ordinary Python
types, so symbolic expressions can be written directly as its elements:

def add_border[Height: IntVar, Width: IntVar](
    image: Array[[Height, Width]],
) -> Array[[Height + 2, Width + 2]]:
    ...

The *Elements[...] form for unpacking IntTuple

IntTuple also supports an unpacked shape variable. Elements[Shape] means
"the dimensions contained in Shape", much like Unpack does for an ordinary
tuple type. For example, a linear layer preserves every leading dimension and
replaces only the last one:

from shape_extensions import Elements

class Linear[In: IntVar, Out: IntVar]:
    def __init__(self, in_features: Int[In], out_features: Int[Out]) -> None:
        ...

    def __call__[Batch: IntTuple](
        self, x: Array[[*Elements[Batch], In]]
    ) -> Array[[*Elements[Batch], Out]]:
        ...

def use_linear(x: Array[[2, 7, 3]]) -> None:
    layer = Linear(3, 5)
    assert_type(layer(x), Array[[2, 7, 5]])

In this call, Pyrefly infers Batch as IntTuple[2, 7], In as 3, and
Out as 5.

Bare IntTuple leaves both rank and dimensions unknown. Array[[int, 3]]
retains rank two and the second dimension while leaving the first dimension
gradual.

Shape transformation DSL functions

Arithmetic and unpacking can express direct relationships between input and
output shapes, but some transforms require control flow. Dropping an axis,
validating a rank, or computing every output shape of a split are easier to
express as functions. @type_shape_dsl_function marks a restricted Python
function that Pyrefly can call from an annotation.

For example:

import shape_extensions.dsl as dsl
from shape_extensions import type_shape_dsl_function

@type_shape_dsl_function
def flatten_last_two_shape(shape: IntTuple) -> IntTuple:
    if len(shape) < 2:
        return dsl.Invalid("flatten_last_two requires at least two dimensions")
    return dsl.concat(
        shape[:-2],
        dsl.IntTuple((shape[-2] * shape[-1],)),
    )

def flatten_last_two[Shape: IntTuple](
    x: Array[Shape],
) -> Array[flatten_last_two_shape(Shape)]:
    ...

def use_flatten(x: Array[[2, 3, 4]]) -> None:
    assert_type(flatten_last_two(x), Array[[2, 12]])

flatten_last_two_shape(Shape) runs after Pyrefly infers Shape. Although it
is not an IntTuple, dsl.Invalid(...) is an allowed DSL sentinel in any
result domain; returning it reports a type error at the call site.

The DSL supports local bindings, conditionals, integer arithmetic, indexing and
slicing, bounded generators, and helpers such as len, range, zip, any,
dsl.concat, dsl.sum, and dsl.prod. Its result domains are Int,
IntTuple, and IntTuples (a tuple of IntTuple values). Imprecise inputs
produce gradual results when evaluation cannot finish precisely.

The Flag bound

Output shapes often depend on ordinary control arguments: for example, a
keepdim argument determines whether a reduction removes an axis. A Flag[T]
bound preserves such arguments as literals and lets the captured value flow
into a DSL call. T may be composed from int, bool, str, None, and
their literal forms. This avoids a combinatorial set of overloads, reducing stub
maintenance and overload-resolution cost:

from shape_extensions import Flag

@type_shape_dsl_function
def reduce_last_shape(shape: IntTuple, keepdim: bool) -> IntTuple:
    if len(shape) == 0:
        return dsl.Invalid("reduce_last requires at least one dimension")
    if keepdim:
        return dsl.concat(shape[:-1], dsl.IntTuple((1,)))
    return shape[:-1]

def reduce_last[Shape: IntTuple, Keepdim: Flag[bool]](
    x: Array[Shape], keepdim: Keepdim = False
) -> Array[reduce_last_shape(Shape, Keepdim)]:
    ...

def use_reduce(x: Array[[2, 3, 4]]) -> None:
    assert_type(reduce_last(x), Array[[2, 3]])
    assert_type(reduce_last(x, keepdim=True), Array[[2, 3, 1]])

The default call binds Keepdim to Literal[False]; the second binds it to
Literal[True]. A Flag type variable directly annotates one function
parameter.

The argument does not have to be statically known. If the caller has only a
broad bool, the call is still valid:

def dynamic_reduce(x: Array[[2, 3, 4]], keepdim: bool) -> None:
    assert_type(reduce_last(x, keepdim), Array[IntTuple])

Because the DSL branches on keepdim, Pyrefly cannot choose a result shape and
falls back to a gradual one. Flag can also capture class parameters, which is
useful when an nn.Module constructor argument affects the shape returned by
forward.

The IntTuples type and MapIntTuples construct

The examples so far compute one output shape, but operations such as split
and chunk return multiple arrays with potentially different shapes.

The IntTuples type represents a tuple of shapes. To connect those shapes to
ordinary generic types such as Tensor or Array, we introduce a type-level
map operator:

MapIntTuples[lambda Shape: Array[Shape], Shapes]

This maps each shape to Array[Shape]. In a parameter position, it can reverse
that mapping to infer Shapes. This operator is restricted to cases where the
second argument is assignable to IntTuples (typically the second argument is
either a type variable or a DSL function call, as we'll describe below).

MapIntTuples in a parameter position

torch.cat uses the parameter form:

from shape_extensions import IntTuples, MapIntTuples

# An excerpt from the shape-aware torch stubs.
def cat[Shapes: IntTuples, Dim: Flag[int]](
    tensors: MapIntTuples[lambda Shape: Tensor[Shape], Shapes],
    dim: Dim = 0,
) -> Tensor[cat_shape(Shapes, Dim)]:
    ...

A call through the public torch API has a symbolic result:

import torch
from torch import Tensor


def concatenate[N: IntVar, M: IntVar](
    x: Tensor[[N, 3]], y: Tensor[[M, 3]]
) -> None:
    assert_type(torch.cat((x, y), dim=0), Tensor[[N + M, 3]])

From (x, y), Pyrefly infers
Shapes = tuple[IntTuple[N, 3], IntTuple[M, 3]]. It then evaluates
cat_shape(Shapes, 0), which checks the ranks and non-concatenated dimensions
and sums the concatenated dimensions.

MapIntTuples in a return position

Tensor.chunk uses the return form to map Tensor over the shapes computed by
its DSL function:

# An excerpt from the Tensor class in the shape-aware torch stubs.
def chunk[
    Shape: IntTuple,
    Chunks: Int,
    Dim: Flag[int],
](
    self: Tensor[Shape],
    chunks: Chunks,
    dim: Dim = 0,
) -> MapIntTuples[
    lambda ResultShape: Tensor[ResultShape],
    chunk_shapes(Shape, Chunks, Dim),
]:
    ...

For a concrete call, the mapped result is a fixed-length tuple with a distinct
shape type for each result:

def use_chunk(x: Tensor[[8, 3]]) -> None:
    assert_type(
        x.chunk(2, dim=0),
        tuple[Tensor[[4, 3]], Tensor[[4, 3]]],
    )

Potential applications to the general type system

Two of these operators may also be useful outside shape typing:

Elements for splatting tuples

The Elements operator is a missing combinator in our current tuple-type
algebra. We can capture a TypeVarTuple, for example:

def add_int_to_tuple[*Ts](x: tuple[*Ts]) -> tuple[*Ts, int]:
    return (*x, 5)

But TypeVarTuple ties the syntax for declaring variadic parameters to the
ability to unpack tuple contents. That prevents signatures which combine two
independently typed tuples.

With Elements, we could write:

def add_int_between_tuples[TX: tuple, TY: tuple](
    x: TX, y: TY
) -> tuple[*Elements[TX], int, *Elements[TY]]:
    ...

I'm not sure how often this comes up, but Elements supplies the missing
operation without requiring both tuples to originate as TypeVarTuples.

Map, which could generalize MapIntTuples to arbitrary tuples

MapIntTuples is currently restricted to IntTuples, a tuple of IntTuple
values. A general Map over tuples could replace fixed-length overload stacks
such as those used by asyncio.gather, which strips Awaitable from every
tuple element. This would simplify stubs, remove their maximum-arity cliff, and
avoid expensive overload resolution. I've seen asyncio.gather calls in
production with over 40 arguments; we are obviously not going to cover those
with overloads.

Some history for the curious: evolution of Pyrefly's shape type system

Since the May 2026 PyCon Typing Summit, Pyrefly's shape system has gone through
three versions as we tested it against array libraries and ML models at Meta.

V0 type system: the initial torch experiment

At the time of Avik's PyCon Typing Summit
talk,
Pyrefly tensor shape types were tied to torch.Tensor.

The shape was expressed as a TypeVarTuple and we used a registration-based DSL
where some functions would be analyzed twice: once using normal types, and then
again to evaluate the shape. We used normal TypeVar type variables to represent
symbolic integers, and we had a Dim[...] type used to capture a symbolic
integer from a normal integer parameter.

For example, we might have had a function like

def repeat2d[N, M](n: Dim[N], x: Tensor[M, K]) -> Tensor[M * N, K]:
    ...

where N and M here are normal type variables.

The DSL shape transforms were written in a Python-like language, but both the
transforms and their registration with torch functions were baked into the
Pyrefly binary so at that point only a Pyrefly developer could reasonably have
worked on developing shape type library stubs.

V1 type system: entirely stub-driven

In an initial sprint to make Pyrefly tensor shapes more useful after PyCon,
we made a series of improvements to the original prototype:

• We decoupled shape transformations from torch.Tensor by introducing a
  ShapedArray type to handle most special behaviors, including shape
  transformations.
• We created a @shape_dsl_function decorator to mark DSL transformations
  so that they could live in stubs rather than the Pyrefly binary, and we
  "registered" them using a @uses_shape_dsl decorator.
• We replaced shapes represented as TypeVarTuples of TypeVars with
  IntTuple, and introduced the IntVar type-variable kind for symbolic
  arithmetic. A distinct runtime class would be needed in any future PEP
  because ordinary TypeVars do not support expressions such as M + N.
• We renamed Dim to Int, which was mostly a matter of taste given
  that the variable dim is actually usually not a size integer but
  rather than axis number in libraries like numpy, pytorch, and jax.

Pyrefly 1.2 shipped this design, including enough support to prototype NumPy
stubs. Its cosine_similarity stub looked like this:

@shape_dsl_function
def cosine_similarity_ir(
    x1: ShapedArray,
    x2: ShapedArray,
    dim: int = 1,
) -> ShapedArray:
    shape = broadcast_ir(x1, x2).shape

    # Scalars accept both possible normalized dimensions.
    if len(shape) == 0:
        if dim == 0 or dim == -1:
            return ShapedArray(shape=shape)
        raise Error("cosine_similarity dimension out of range")

    if dim < -len(shape) or dim >= len(shape):
        raise Error("cosine_similarity dimension out of range")

    normalized_dim = dim + len(shape) if dim < 0 else dim
    return ShapedArray(
        shape=[
            shape[i]
            for i in range(len(shape))
            if i != normalized_dim
        ]
    )

@uses_shape_dsl(cosine_similarity_ir)
def cosine_similarity[S1: IntTuple, S2: IntTuple](
    x1: Tensor[S1], x2: Tensor[S2], dim: int = 1, eps: float = 1e-8
) -> Tensor:
    """Cosine similarity: dot product along dim, normalized."""
    ...

V2 type system: where we are today

V1 exposed two problems:

• The mapping from decorated-function parameters to transform arguments was
  hard to understand and unsound for some optional arguments, making the
  implementation difficult to maintain or specify.
• DSL functions accepted and returned only ShapedArray. That excluded
  non-array abstractions and layouts with multiple integer tuples, such as
  CuTe shapes and strides.

PEP 827 introduces expression-like
type-level syntax for transformations from input types to output types. That
suggested the replacement: make each shape transform an explicit type-level
function call. The V2 DSL accepts a restricted
subset of Python and returns types that can be used directly in annotations.

The cosine similarity example described above now looks like this:

from shape_extensions import broadcast

@type_shape_dsl_function
def cosine_similarity_shape(shape: IntTuple, dim: int) -> IntTuple:
    if dim == -1:
        if len(shape) == 0:
            return shape
        return shape[:-1]
    if len(shape) == 0:
        if dim == 0:
            return shape
        return dsl.Invalid("cosine_similarity dimension out of range")
    if dim < -len(shape) or dim >= len(shape):
        return dsl.Invalid("cosine_similarity dimension out of range")
    return dsl.IntTuple(
        (
            shape[index]
            for index in range(len(shape))
            if index != (dim + len(shape) if dim < 0 else dim)
        )
    )


def cosine_similarity[S1: IntTuple, S2: IntTuple, Dim: Flag[int]](
    x1: Tensor[S1], x2: Tensor[S2], dim: Dim = 1, eps: float = 1e-8
) -> Tensor[cosine_similarity_shape(broadcast(S1, S2), Dim)]:
    """Cosine similarity: dot product along dim, normalized."""
    ...

S1, S2, and Dim are inferred at the call site; the return annotation
passes broadcast(S1, S2) and Dim to the DSL function.

The result is a system whose core concepts are independent of torch and can
represent integer-valued properties on abstractions beyond arrays. There is
still work to do before these ideas are ready for standardization, but explicit
type-level functions give us a simpler foundation for experimenting with shape
typing and, potentially, other forms of integer-aware typing in Python.

Looking forward

We're still making some changes, for example adding an Index type so that we can eliminate some
baked-in logic for array indexing.

But we're fairly confident about the core components:

• IntVar and Int for symbolic arithmetic
• IntTuple and IntTuples for composing ints into shapes
• A type-level DSL for shape transformations that are too complex for closed-form notation

We'd love to hear thoughts about this or ideas about what could be improved