website/docs/tensor-shapes-setup.mdx - third_party/github.com/facebook/pyrefly - Git at Google

 ---
 title: Getting Started
 description: How to configure Pyrefly for tensor shape checking and set up your project.
 ---

 {/*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  *
  * This source code is licensed under the MIT license found in the
  * LICENSE file in the root directory of this source tree.
  */}

 # 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.

 ## Installation

 Tensor shape checking requires shape-aware stubs to be available through normal
 import resolution. Install them from PyPI:

 ```bash
 pip install pyrefly-torch-stubs
 ```

 `pyrefly-torch-stubs` is a PEP 561 stub-only package: it carries type
 information for PyTorch, and leaves the runtime `torch` package alone. PyTorch's
 own type stubs don't carry shape information, so these stubs take precedence
 over them and provide shape-aware versions (e.g., `nn.Conv2d.__init__` that
 captures kernel size, stride, and padding as type-level values, and a `forward`
 that computes the output spatial dimensions).

 Installing it also pulls in `pyrefly-shape-extensions`, which provides the
 `shape_extensions` package. `shape_extensions` exports `Int` — the bridge
 between runtime integer values and type-level symbols. Both packages are
 versioned in lockstep with Pyrefly.

 Tensor shape support is enabled automatically when Pyrefly can resolve the
 `shape_extensions` package, so nothing else needs configuring, as long as
 Pyrefly resolves imports against the
 [Python environment](../configuration#python-interpreter-path) you installed
 into.

 ### Checking against a local copy of the stubs

 The stubs also live in Pyrefly's source tree under
 [`tensor-shapes/`](https://github.com/facebook/pyrefly/tree/main/tensor-shapes),
 which is where to work if you want to read or modify them. Copy that directory
 into your project and point `search-path` at the two package directories inside
 it. The paths are relative to the location of your `pyrefly.toml` (or of your
 `pyproject.toml`, if you configure Pyrefly under `[tool.pyrefly]`):

 ```toml
 search-path = [
     "tensor-shapes/pyrefly-torch-stubs",
     "tensor-shapes/pyrefly-shape-extensions",
 ]
 ```

 A copy that lives inside your project is also checked as project code. The stubs
 use PEP 696 type parameter defaults (`class Conv1d[..., S: IntVar = 1]`), which
 Pyrefly parses only under `python-version` 3.13 or later, so checking them under
 an earlier version reports:

 ```
 ERROR Cannot set default type for a type parameter on Python 3.12 (syntax was added in Python 3.13) [invalid-syntax]
 ```

 Either raise the version Pyrefly checks against. `python-version` sets the
 version your code is checked for; it does not change the interpreter your code
 runs on:

 ```toml
 python-version = "3.13"
 ```

 Or, if your project targets an earlier version and you don't want to change how
 your own code is checked, exclude the copy from checking instead. Imports still
 resolve through `search-path`:

 ```toml
 project-excludes = ["tensor-shapes/**"]
 ```

 ## Imports and runtime considerations

 Python evaluates type annotations at runtime by default. This is a problem
 for tensor shape annotations because Python's built-in `typing.TypeVar`
 doesn't support arithmetic — expressions like `D // NHead` in an annotation
 will raise `TypeError` when the annotation is evaluated. There are two ways
 to avoid this:

 ### Option 1: `from __future__ import annotations` (recommended)

 Adding this import at the top of the file defers evaluation of all
 annotations, so shape arithmetic never executes at runtime:

 ```python
 from __future__ import annotations

 import torch
 import torch.nn as nn
 from torch import Tensor
 from shape_extensions import Int
 ```

 This works with both old-style and new-style generics (PEP 695
 `class Foo[T]` syntax).

 **`assert_type` during development:** You can use `assert_type` while
 porting to verify shapes via `pyrefly check`. Once you're done, remove the
 `assert_type` calls — each one corresponds to an IDE inlay type hint that
 shows the same information permanently. Pyrefly catches shape errors
 through your function signatures and return types regardless.

 Note that `assert_type` forces evaluation of its type argument, so the
 file will crash if you try to *run* it with `assert_type` calls still
 present. This is fine — just remove them when the port is complete.

 You can also guard `Tensor` and `Int` under `TYPE_CHECKING` if you prefer
 to keep shape imports invisible at runtime:

 ```python
 from __future__ import annotations
 from typing import TYPE_CHECKING

 import torch
 import torch.nn as nn

 if TYPE_CHECKING:
     from torch import Tensor
     from shape_extensions import Int
 ```

 ### Option 2: `shape_extensions.IntVar` (runtime-compatible)

 If you need annotations to evaluate at runtime (e.g., for runtime shape
 validation or keeping `assert_type` in production code), import
 `shape_extensions` directly. The package patches
 `torch.Tensor`, `nn.Conv2d`, and other torch classes to accept subscript
 syntax at runtime without crashing. It also provides a `IntVar` that
 supports arithmetic (`N + 1` returns `self` instead of raising `TypeError`).

 Use old-style generics with `shape_extensions.IntVar`:

 ```python
 from typing import assert_type

 import torch
 import torch.nn as nn
 from torch import Tensor
 from shape_extensions import Int, IntVar

 N = IntVar("N")
 M = IntVar("M")

 class Linear(nn.Module):
     def __init__(self, n: Int[N], m: Int[M]):
         ...
 ```

 PEP 695 new-style generics (`class Foo[T]`) automatically use
 `typing.TypeVar` internally, which doesn't support arithmetic — so this
 option requires old-style generics.

 ```sandbox
 dir: tensor-shapes-setup
 description: Experiment with Int arithmetic, TYPE_CHECKING imports, and building typed tensors — no setup required.
 ```

 ## Hello world

 Here's a minimal example to verify everything works. This example uses
 Option 1 (`from __future__ import annotations`) since it's the simplest
 setup. We skip `assert_type` — instead, run `pyrefly check` and use your
 IDE's inlay type hints to verify shapes.

 Create a file `hello_shapes.py`:

 ```python
 from __future__ import annotations

 import torch
 import torch.nn as nn
 from torch import Tensor
 from shape_extensions import Int, IntVar


 class TwoLayerNet[InDim: IntVar, HidDim: IntVar, OutDim: IntVar](nn.Module):
     def __init__(
         self,
         in_dim: Int[InDim],
         hid_dim: Int[HidDim],
         out_dim: Int[OutDim],
     ):
         super().__init__()
         self.fc1 = nn.Linear(in_dim, hid_dim)
         self.fc2 = nn.Linear(hid_dim, out_dim)

     def forward[B: IntVar](self, x: Tensor[[B, InDim]]) -> Tensor[[B, OutDim]]:
         h = self.fc1(x)          # pyrefly infers: Tensor[[B, HidDim]]
         return self.fc2(torch.relu(h))
 ```

 Run `pyrefly check hello_shapes.py`. You should see no errors — pyrefly
 infers the shapes through the `nn.Linear` calls.

 If you're using an IDE with Pyrefly's language server, you'll see inlay
 type hints showing the inferred shape of `h` as `Tensor[[B, HidDim]]`
 without needing any `assert_type` calls.

 ### Inlay hints in action

 Here's what inlay hints look like on a real model (NanoGPT). The MLP
 module shows shapes flowing through linear layers and activations:

 <img src="/img/tensor-shapes/nanogpt-mlp.png" alt="NanoGPT MLP module with inlay type hints showing Tensor[B, T, 4 * NEmbedding] after the expansion layer" />

 The forward method signature shows how `x.size()` unpacks into typed
 dimensions:

 <img src="/img/tensor-shapes/nanogpt-forward-signature.png" alt="NanoGPT forward signature with x.size() unpacking to b: Int[B], t: Int[T], c: Int[NEmbedding]" />

 And the attention module, where view/transpose reshapes for multi-head
 attention are fully tracked:

 <img src="/img/tensor-shapes/nanogpt-attention-qkv.png" alt="NanoGPT attention QKV computation with inlay type hints showing Tensor[B, NHead, T, NEmbedding // NHead]" />

 The full attention body, including both flash and manual paths:

 <img src="/img/tensor-shapes/nanogpt-attention-body.png" alt="NanoGPT attention body with flash and manual attention paths fully shape-tracked" />
	---
	title: Getting Started
	description: How to configure Pyrefly for tensor shape checking and set up your project.
	---

	{/*
	* Copyright (c) Meta Platforms, Inc. and affiliates.
	*
	* This source code is licensed under the MIT license found in the
	* LICENSE file in the root directory of this source tree.
	*/}

	# 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.

	## Installation

	Tensor shape checking requires shape-aware stubs to be available through normal
	import resolution. Install them from PyPI:

	```bash
	pip install pyrefly-torch-stubs
	```

	`pyrefly-torch-stubs` is a PEP 561 stub-only package: it carries type
	information for PyTorch, and leaves the runtime `torch` package alone. PyTorch's
	own type stubs don't carry shape information, so these stubs take precedence
	over them and provide shape-aware versions (e.g., `nn.Conv2d.__init__` that
	captures kernel size, stride, and padding as type-level values, and a `forward`
	that computes the output spatial dimensions).

	Installing it also pulls in `pyrefly-shape-extensions`, which provides the
	`shape_extensions` package. `shape_extensions` exports `Int` — the bridge
	between runtime integer values and type-level symbols. Both packages are
	versioned in lockstep with Pyrefly.

	Tensor shape support is enabled automatically when Pyrefly can resolve the
	`shape_extensions` package, so nothing else needs configuring, as long as
	Pyrefly resolves imports against the
	[Python environment](../configuration#python-interpreter-path) you installed
	into.

	### Checking against a local copy of the stubs

	The stubs also live in Pyrefly's source tree under
	[`tensor-shapes/`](https://github.com/facebook/pyrefly/tree/main/tensor-shapes),
	which is where to work if you want to read or modify them. Copy that directory
	into your project and point `search-path` at the two package directories inside
	it. The paths are relative to the location of your `pyrefly.toml` (or of your
	`pyproject.toml`, if you configure Pyrefly under `[tool.pyrefly]`):

	```toml
	search-path = [
	"tensor-shapes/pyrefly-torch-stubs",
	"tensor-shapes/pyrefly-shape-extensions",
	]
	```

	A copy that lives inside your project is also checked as project code. The stubs
	use PEP 696 type parameter defaults (`class Conv1d[..., S: IntVar = 1]`), which
	Pyrefly parses only under `python-version` 3.13 or later, so checking them under
	an earlier version reports:

	```
	ERROR Cannot set default type for a type parameter on Python 3.12 (syntax was added in Python 3.13) [invalid-syntax]
	```

	Either raise the version Pyrefly checks against. `python-version` sets the
	version your code is checked for; it does not change the interpreter your code
	runs on:

	```toml
	python-version = "3.13"
	```

	Or, if your project targets an earlier version and you don't want to change how
	your own code is checked, exclude the copy from checking instead. Imports still
	resolve through `search-path`:

	```toml
	project-excludes = ["tensor-shapes/**"]
	```

	## Imports and runtime considerations

	Python evaluates type annotations at runtime by default. This is a problem
	for tensor shape annotations because Python's built-in `typing.TypeVar`
	doesn't support arithmetic — expressions like `D // NHead` in an annotation
	will raise `TypeError` when the annotation is evaluated. There are two ways
	to avoid this:

	### Option 1: `from __future__ import annotations` (recommended)

	Adding this import at the top of the file defers evaluation of all
	annotations, so shape arithmetic never executes at runtime:

	```python
	from __future__ import annotations

	import torch
	import torch.nn as nn
	from torch import Tensor
	from shape_extensions import Int
	```

	This works with both old-style and new-style generics (PEP 695
	`class Foo[T]` syntax).

	`assert_type` during development: You can use `assert_type` while
	porting to verify shapes via `pyrefly check`. Once you're done, remove the
	`assert_type` calls — each one corresponds to an IDE inlay type hint that
	shows the same information permanently. Pyrefly catches shape errors
	through your function signatures and return types regardless.

	Note that `assert_type` forces evaluation of its type argument, so the
	file will crash if you try to run it with `assert_type` calls still
	present. This is fine — just remove them when the port is complete.

	You can also guard `Tensor` and `Int` under `TYPE_CHECKING` if you prefer
	to keep shape imports invisible at runtime:

	```python
	from __future__ import annotations
	from typing import TYPE_CHECKING

	import torch
	import torch.nn as nn

	if TYPE_CHECKING:
	from torch import Tensor
	from shape_extensions import Int
	```

	### Option 2: `shape_extensions.IntVar` (runtime-compatible)

	If you need annotations to evaluate at runtime (e.g., for runtime shape
	validation or keeping `assert_type` in production code), import
	`shape_extensions` directly. The package patches
	`torch.Tensor`, `nn.Conv2d`, and other torch classes to accept subscript
	syntax at runtime without crashing. It also provides a `IntVar` that
	supports arithmetic (`N + 1` returns `self` instead of raising `TypeError`).

	Use old-style generics with `shape_extensions.IntVar`:

	```python
	from typing import assert_type

	import torch
	import torch.nn as nn
	from torch import Tensor
	from shape_extensions import Int, IntVar

	N = IntVar("N")
	M = IntVar("M")

	class Linear(nn.Module):
	def __init__(self, n: Int[N], m: Int[M]):
	...
	```

	PEP 695 new-style generics (`class Foo[T]`) automatically use
	`typing.TypeVar` internally, which doesn't support arithmetic — so this
	option requires old-style generics.

	```sandbox
	dir: tensor-shapes-setup
	description: Experiment with Int arithmetic, TYPE_CHECKING imports, and building typed tensors — no setup required.
	```

	## Hello world

	Here's a minimal example to verify everything works. This example uses
	Option 1 (`from __future__ import annotations`) since it's the simplest
	setup. We skip `assert_type` — instead, run `pyrefly check` and use your
	IDE's inlay type hints to verify shapes.

	Create a file `hello_shapes.py`:

	```python
	from __future__ import annotations

	import torch
	import torch.nn as nn
	from torch import Tensor
	from shape_extensions import Int, IntVar


	class TwoLayerNet[InDim: IntVar, HidDim: IntVar, OutDim: IntVar](nn.Module):
	def __init__(
	self,
	in_dim: Int[InDim],
	hid_dim: Int[HidDim],
	out_dim: Int[OutDim],
	):
	super().__init__()
	self.fc1 = nn.Linear(in_dim, hid_dim)
	self.fc2 = nn.Linear(hid_dim, out_dim)

	def forward[B: IntVar](self, x: Tensor[[B, InDim]]) -> Tensor[[B, OutDim]]:
	h = self.fc1(x) # pyrefly infers: Tensor[[B, HidDim]]
	return self.fc2(torch.relu(h))
	```

	Run `pyrefly check hello_shapes.py`. You should see no errors — pyrefly
	infers the shapes through the `nn.Linear` calls.

	If you're using an IDE with Pyrefly's language server, you'll see inlay
	type hints showing the inferred shape of `h` as `Tensor[[B, HidDim]]`
	without needing any `assert_type` calls.

	### Inlay hints in action

	Here's what inlay hints look like on a real model (NanoGPT). The MLP
	module shows shapes flowing through linear layers and activations:

	<img src="/img/tensor-shapes/nanogpt-mlp.png" alt="NanoGPT MLP module with inlay type hints showing Tensor[B, T, 4 * NEmbedding] after the expansion layer" />

	The forward method signature shows how `x.size()` unpacks into typed
	dimensions:

	<img src="/img/tensor-shapes/nanogpt-forward-signature.png" alt="NanoGPT forward signature with x.size() unpacking to b: Int[B], t: Int[T], c: Int[NEmbedding]" />

	And the attention module, where view/transpose reshapes for multi-head
	attention are fully tracked:

	<img src="/img/tensor-shapes/nanogpt-attention-qkv.png" alt="NanoGPT attention QKV computation with inlay type hints showing Tensor[B, NHead, T, NEmbedding // NHead]" />

	The full attention body, including both flash and manual paths:

	<img src="/img/tensor-shapes/nanogpt-attention-body.png" alt="NanoGPT attention body with flash and manual attention paths fully shape-tracked" />