diff --git a/website/docs/tensor-shapes-setup.mdx b/website/docs/tensor-shapes-setup.mdx index 79c0d40496..150244d9bb 100644 --- a/website/docs/tensor-shapes-setup.mdx +++ b/website/docs/tensor-shapes-setup.mdx @@ -15,37 +15,73 @@ description: How to configure Pyrefly for tensor shape checking and set up your This page walks you through configuring Pyrefly for tensor shape checking and getting your first shape-annotated code running. -## Configuration +## Installation Tensor shape checking requires shape-aware stubs to be available through normal -import resolution. For a local setup, add the fixture stub directory to your -`pyrefly.toml` (or under `[tool.pyrefly]` in `pyproject.toml`): +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 = [ - "path/to/fixtures", +search-path = [ + "tensor-shapes/pyrefly-torch-stubs", + "tensor-shapes/pyrefly-shape-extensions", ] ``` -Tensor shape support is enabled automatically when Pyrefly can resolve the -`shape_extensions` package. The fixture directory provides `shape_extensions` -and shape-aware versions of library stubs. +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: -**`search_path`** points to a directory of *fixture stubs* — `.pyi` files that -provide shape-generic type signatures for PyTorch modules and functions. The -real `torch` library's type stubs don't carry shape information, so the fixtures -replace them with 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). +``` +ERROR Cannot set default type for a type parameter on Python 3.12 (syntax was added in Python 3.13) [invalid-syntax] +``` -The stubs live in Pyrefly's source tree under -[`tensor-shapes/`](https://github.com/facebook/pyrefly/tree/main/tensor-shapes). -To use them in your project, copy the `tensor-shapes/` directory into your -project and set `search_path` to point to it. The path is relative to the -location of your `pyrefly.toml`. +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: -The fixtures also provide the `shape_extensions` package, which exports `Int` — the -bridge between runtime integer values and type-level symbols. +```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 @@ -148,10 +184,10 @@ from __future__ import annotations import torch import torch.nn as nn from torch import Tensor -from shape_extensions import Int +from shape_extensions import Int, IntVar -class TwoLayerNet[InDim, HidDim, OutDim](nn.Module): +class TwoLayerNet[InDim: IntVar, HidDim: IntVar, OutDim: IntVar](nn.Module): def __init__( self, in_dim: Int[InDim], @@ -162,8 +198,8 @@ class TwoLayerNet[InDim, HidDim, OutDim](nn.Module): self.fc1 = nn.Linear(in_dim, hid_dim) self.fc2 = nn.Linear(hid_dim, out_dim) - def forward[B](self, x: Tensor[B, InDim]) -> Tensor[B, OutDim]: - h = self.fc1(x) # pyrefly infers: Tensor[B, HidDim] + 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)) ``` @@ -171,7 +207,7 @@ 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]` +type hints showing the inferred shape of `h` as `Tensor[[B, HidDim]]` without needing any `assert_type` calls. ### Inlay hints in action