API Reference

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

tip

Try it yourself

Open this example in the Pyrefly sandbox — Experiment with Int operators, type-level arithmetic in shapes, and concatenation.

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 ** b	Int[A ** B]

Caution: int * Int produces Unknown because the int side has no type-level identity. 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 positive 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[[*Elements[Bs], D]]	Variadic: any number of leading batch dimensions (Bs: IntTuple)
Tensor (bare)	Shape unknown — tracking gap

Variadic dimensions with Elements

Use a type parameter bound to IntTuple, spliced into the shape list with *Elements[...], for dimensions that should be propagated without being enumerated:

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

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

Don't hide known class dims inside variadic params. If the module has a class-level Int D, use Tensor[[*Elements[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]]
• x.size(0) has type Int[B]
• x.size() has type tuple[Int[B], Int[C], Int[H], Int[W]]

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 the assert_type calls — each one corresponds to an inlay type hint that your IDE shows permanently. Pyrefly catches shape errors through function signatures and return types regardless.

assert_type forces evaluation of its type argument at runtime, so a file with assert_type calls will crash if executed. This is fine during development (you run pyrefly check, not the file itself) — just remove them when the port is done.

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.

Jaxtyping compatibility

Pyrefly has experimental static support for a subset of jaxtyping annotations. Support is disabled by default. Enable it for a project or sub-config with:

jaxtyping = true

Pyrefly recognizes jaxtyping's public array and dtype wrappers, including Shaped, Float, and Int. An annotation is shape-checked when its array type is either a Pyrefly shaped-array class or a generic class with one gradual IntTuple-bounded shape argument. The currently supported shape syntax is:

• named dimensions, integer dimensions, and scalar shapes;
• anonymous dimensions (_);
• one named variadic dimension (*shape) or ellipsis (...), with optional prefix and suffix dimensions;
• broadcast markers (#batch and *#batch), currently treated as their non-broadcast equivalents; and
• simple addition and subtraction such as n+1, n-1, and (1+n).

Within a function, repeated names represent the same dimension and participate in inference across parameters and returns. Supported annotations are translated internally to generics and display back in jaxtyping syntax:

Pyrefly native	Jaxtyping equivalent
Tensor[[M, 2, M + 1]]	Shaped[Tensor, "M 2 M+1"]
Tensor[[B, C, H, W]]	Shaped[Tensor, "B C H W"]

This is not yet full static jaxtyping compatibility. In particular:

• dtype wrappers are recognized as markers, but dtype refinements are not checked;
• broadcast markers do not yet implement jaxtyping's broadcast semantics;
• the full dimension-expression grammar and PyTree annotations are not supported;
• generic type aliases, dependent generic defaults, and array types without a unique gradual shape argument fall back to their ordinary stub types; and
• jaxtyping cannot share symbolic dimensions across methods or class boundaries.

Full static support would require implementing those semantics and extending the compatibility tests across the remaining jaxtyping grammar and container forms. The cross-method/class limitation also requires an upstream annotation model for shared dimensions; see the overview for details.