docs: fix tensor shapes syntax to match shipped stubs

website/docs/tensor-shapes-ai-porting.mdx

@@ -84,7 +84,7 @@ The skill produces a `verify_port.sh` report with these key metrics:
 | **`ig` (ignore)** | `type: ignore` comments. Lower is better. Each should have a category (algebraic gap, conditional equality, stub gap). |
 | **`bs` (bare sig)** | Bare `Tensor` in function signatures. Should be 0 for well-typed modules. |
 | **`bv` (bare var)** | Bare `Tensor` in local variable annotations. Lower is better. |
-| **`sh` (shaped)** | `assert_type` calls with full shapes (e.g., `Tensor[B, D]`). Higher is better. |
+| **`sh` (shaped)** | `assert_type` calls with full shapes (e.g., `Tensor[[B, D]]`). Higher is better. |
 | **`ba` (bare assert)** | `assert_type(x, Tensor)` — tracking gaps. Lower is better. Each should have a root-cause comment. |
 | **`sm` (smoke)** | Smoke test functions. At least 1-2 per model. |
 

website/docs/tensor-shapes-reference.mdx

@@ -23,7 +23,11 @@ description: Experiment with Int operators, type-level arithmetic in shapes, and
 ## `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.
+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
 
@@ -61,11 +65,11 @@ For optional dimensions — parameters that may or may not be present — use
 `if param is not None:` to recover `Int[X]` inside the branch:
 
 ```python
-class Attention[D, RK](nn.Module):
+class Attention[D: IntVar, RK: IntVar](nn.Module):
     def __init__(self, dim: Int[D], rank_k: Int[RK] | None = None):
         ...
 
-    def forward[B, T](self, x: Tensor[B, T, D]) -> Tensor[B, T, D]:
+    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
             ...
@@ -76,11 +80,11 @@ class Attention[D, RK](nn.Module):
 | Pattern | Purpose |
 |---------|---------|
 | `def __init__(self, dim: Int[D])` | Accept a dimension as a constructor parameter |
-| `class Model[D](nn.Module)` | Make a dimension a class-level type parameter |
-| `def forward[B](self, x: Tensor[B, D])` | Bind a per-call dimension |
+| `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]`) |
 
-## `Tensor[D1, D2, ...]`
+## `Tensor[[D1, D2, ...]]`
 
 `Tensor` with type arguments represents a tensor with a known shape. The
 type arguments are the dimensions, in order.
@@ -89,31 +93,33 @@ type arguments are the dimensions, in order.
 
 | 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 |
+| `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 `*Bs`
+### Variadic dimensions with `Elements`
 
-Use `TypeVarTuple` (or PEP 646 `*Bs` syntax) for dimensions that should be
-propagated without being enumerated:
+Use a type parameter bound to `IntTuple`, spliced into the shape list with
+`*Elements[...]`, for dimensions that should be propagated without being
+enumerated:
 
 ```python
-def forward[*Bs](self, x: Tensor[*Bs, InDim]) -> Tensor[*Bs, OutDim]:
+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[*Bs, D]` not `Tensor[*S]`.
+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]`:
+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]`
@@ -130,7 +136,7 @@ reports an error.
 
 ```python
 h = self.fc1(x)
-assert_type(h, Tensor[B, 512])  # checked by pyrefly
+assert_type(h, Tensor[[B, 512]])  # checked by pyrefly
 ```
 
 Use `assert_type` during development to verify inferred shapes as you port
@@ -161,7 +167,7 @@ Use it to understand what pyrefly infers before writing `assert_type`:
 
 ```python
 h = self.fc1(x)
-reveal_type(h)  # Revealed type: Tensor[B, 512]
+reveal_type(h)  # Revealed type: Tensor[[B, 512]]
 ```
 
 Replace `reveal_type` with `assert_type` once you know the expected type.
@@ -172,11 +178,11 @@ 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]` |
+| 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
 
@@ -201,7 +207,7 @@ 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(...)`.
+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).
@@ -216,8 +222,8 @@ annotations as an alternative front-end:
 
 | Pyrefly native | Jaxtyping equivalent |
 |---------------|---------------------|
-| `Tensor[M, 2, M // 2]` | `Shaped[Tensor, "M 2 M//2"]` |
-| `Tensor[B, C, H, W]` | `Shaped[Tensor, "B C H W"]` |
+| `Tensor[[M, 2, M // 2]]` | `Shaped[Tensor, "M 2 M//2"]` |
+| `Tensor[[B, C, H, W]]` | `Shaped[Tensor, "B C H W"]` |
 
 Jaxtyping annotations are translated internally to generics and display
 back in jaxtyping syntax. Note that jaxtyping cannot share symbolic

website/docs/tensor-shapes-tutorial-advanced.mdx

@@ -24,7 +24,7 @@ modules, make it generic so dimensions propagate through constructors:
 
 ```python
 @dataclass
-class GPTConfig[VocabSize, BlockSize, NEmbedding, NHead, NLayer]:
+class GPTConfig[VocabSize: IntVar, BlockSize: IntVar, NEmbedding: IntVar, NHead: IntVar, NLayer: IntVar]:
     block_size: Int[BlockSize]
     vocab_size: Int[VocabSize]
     n_layer: Int[NLayer]
@@ -37,17 +37,17 @@ class GPTConfig[VocabSize, BlockSize, NEmbedding, NHead, NLayer]:
 Modules extract only the type parameters they need, using `Any` for the rest:
 
 ```python
-class MLP[NEmbedding](nn.Module):
+class MLP[NEmbedding: IntVar](nn.Module):
     def __init__(self, config: GPTConfig[Any, Any, NEmbedding, Any, Any]):
         super().__init__()
         self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
         self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
 
-    def forward[B, T](
-        self, x: Tensor[B, T, NEmbedding]
-    ) -> Tensor[B, T, NEmbedding]:
+    def forward[B: IntVar, T: IntVar](
+        self, x: Tensor[[B, T, NEmbedding]]
+    ) -> Tensor[[B, T, NEmbedding]]:
         h = F.gelu(self.c_fc(x))
-        assert_type(h, Tensor[B, T, 4 * NEmbedding])
+        assert_type(h, Tensor[[B, T, 4 * NEmbedding]])
         return self.c_proj(h)
 ```
 
@@ -101,7 +101,7 @@ remain grouped:
 self.fc1 = nn.Linear(dim, hidden)
 self.fc2 = nn.Linear(hidden, dim)
 
-def forward[B](self, x: Tensor[B, D]) -> Tensor[B, D]:
+def forward[B: IntVar](self, x: Tensor[[B, D]]) -> Tensor[[B, D]]:
     h = F.relu(self.fc1(x))
     return self.fc2(h)
 ```
@@ -124,16 +124,16 @@ Fix: use a class with a typed `forward` method:
 
 ```python
 # Good: class preserves shape contract
-class Block[InC, OutC](nn.Module):
+class Block[InC: IntVar, OutC: IntVar](nn.Module):
     def __init__(self, in_c: Int[InC], out_c: Int[OutC]) -> None:
         super().__init__()
         self.net = nn.Sequential(
             nn.Conv2d(in_c, 128, ...), nn.Conv2d(128, out_c, ...)
         )
 
-    def forward[B, H, W](
-        self, x: Tensor[B, InC, H, W]
-    ) -> Tensor[B, OutC, H, W]:
+    def forward[B: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, InC, H, W]]
+    ) -> Tensor[[B, OutC, H, W]]:
         return self.net(x)
 ```
 
@@ -150,7 +150,7 @@ Fix: add an explicit `Int` field to the config:
 
 ```python
 @dataclass
-class Config[K, MlpOut]:
+class Config[K: IntVar, MlpOut: IntVar]:
     num_output_features: Int[K]
     mlp_output_dim: Int[MlpOut]        # explicit — was hidden_units[-1]
     mlp_hidden_units: list[int] = field(default_factory=lambda: [512, 256])
@@ -163,7 +163,7 @@ resort is a **typed interface**: the module's `forward` signature provides
 the shape contract, and `type: ignore` narrows the internal result:
 
 ```python
-class DynamicMLP[InDim, OutDim](nn.Module):
+class DynamicMLP[InDim: IntVar, OutDim: IntVar](nn.Module):
     def __init__(self, in_dim: Int[InDim], out_dim: Int[OutDim],
                  hidden: list[int]) -> None:
         super().__init__()
@@ -175,15 +175,15 @@ class DynamicMLP[InDim, OutDim](nn.Module):
         layers.append(nn.Linear(prev, out_dim))
         self.layers = nn.ModuleList(layers)
 
-    def forward[B](self, x: Tensor[B, InDim]) -> Tensor[B, OutDim]:
+    def forward[B: IntVar](self, x: Tensor[[B, InDim]]) -> Tensor[[B, OutDim]]:
         h = x
         for layer in self.layers:
             h = layer(h)
-        result: Tensor[B, OutDim] = h  # type: ignore[bad-assignment]
+        result: Tensor[[B, OutDim]] = h  # type: ignore[bad-assignment]
         return result
 ```
 
-The caller sees a clean `Tensor[B, InDim] -> Tensor[B, OutDim]` contract.
+The caller sees a clean `Tensor[[B, InDim]] -> Tensor[[B, OutDim]]` contract.
 The `type: ignore` is localized to the module's internals.
 
 Use typed interfaces only after exhausting all restructuring options —

website/docs/tensor-shapes-tutorial-architectures.mdx

@@ -42,27 +42,27 @@ This gives a recursive signature where `recurse` takes and returns the same
 shape:
 
 ```python
-class UNet[NChannels, NClasses](nn.Module):
-    def _encode[B, C, H, W](
-        self, x: Tensor[B, C, H, W], depth: int
-    ) -> Tensor[B, 2 * C, (H - 2) // 2 + 1, (W - 2) // 2 + 1]:
+class UNet[NChannels: IntVar, NClasses: IntVar](nn.Module):
+    def _encode[B: IntVar, C: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, C, H, W]], depth: int
+    ) -> Tensor[[B, 2 * C, (H - 2) // 2 + 1, (W - 2) // 2 + 1]]:
         idx = len(self.downs) - depth
         down: Down[C, 2 * C] = self.downs[idx]
         return down(x)
 
-    def _decode[B, C, H, W](
+    def _decode[B: IntVar, C: IntVar, H: IntVar, W: IntVar](
         self,
-        skip: Tensor[B, C, H, W],
-        deep: Tensor[B, 2 * C, (H - 2) // 2 + 1, (W - 2) // 2 + 1],
+        skip: Tensor[[B, C, H, W]],
+        deep: Tensor[[B, 2 * C, (H - 2) // 2 + 1, (W - 2) // 2 + 1]],
         depth: int,
-    ) -> Tensor[B, C, H, W]:
+    ) -> Tensor[[B, C, H, W]]:
         idx = len(self.ups) - depth
         up: Up[2 * C, C] = self.ups[idx]
         return up(deep, skip)
 
-    def recurse[I, B, C, H, W](
-        self, x: Tensor[B, C, H, W], depth: Int[I]
-    ) -> Tensor[B, C, H, W]:
+    def recurse[I: IntVar, B: IntVar, C: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, C, H, W]], depth: Int[I]
+    ) -> Tensor[[B, C, H, W]]:
         if depth == 0:
             return x
         skip = x
@@ -111,27 +111,27 @@ Use `@overload` to separate the base case from the recursive case:
 
 ```python
 class Generator(nn.Module):
-    def _apply_stage[B, C, H, W](
-        self, x: Tensor[B, C, H, W], depth: int
-    ) -> Tensor[B, C // 2, (H - 1) * 2 + 2, (W - 1) * 2 + 2]:
+    def _apply_stage[B: IntVar, C: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, C, H, W]], depth: int
+    ) -> Tensor[[B, C // 2, (H - 1) * 2 + 2, (W - 1) * 2 + 2]]:
         idx = len(self.up_stages) - depth
         stage: GenUpStage[C] = self.up_stages[idx]
         return stage(x)
 
     @overload
-    def _chain[B, C, H, W](
-        self, x: Tensor[B, C, H, W], depth: Int[1]
-    ) -> Tensor[B, C // 2, H * 2, W * 2]: ...
+    def _chain[B: IntVar, C: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, C, H, W]], depth: Int[1]
+    ) -> Tensor[[B, C // 2, H * 2, W * 2]]: ...
 
     @overload
-    def _chain[I, B, C, H, W](
-        self, x: Tensor[B, C, H, W], depth: Int[I]
-    ) -> Tensor[B, C // 2 ** I, H * 2 ** I, W * 2 ** I]: ...
-
-    def _chain[I, B, C, H, W](
-        self, x: Tensor[B, C, H, W], depth: Int[I]
-    ) -> (Tensor[B, C // 2, H * 2, W * 2]
-         | Tensor[B, C // 2 ** I, H * 2 ** I, W * 2 ** I]):
+    def _chain[I: IntVar, B: IntVar, C: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, C, H, W]], depth: Int[I]
+    ) -> Tensor[[B, C // 2 ** I, H * 2 ** I, W * 2 ** I]]: ...
+
+    def _chain[I: IntVar, B: IntVar, C: IntVar, H: IntVar, W: IntVar](
+        self, x: Tensor[[B, C, H, W]], depth: Int[I]
+    ) -> (Tensor[[B, C // 2, H * 2, W * 2]]
+         | Tensor[[B, C // 2 ** I, H * 2 ** I, W * 2 ** I]]):
         y = self._apply_stage(x, depth)
         if depth == 1:
             return y
@@ -154,11 +154,11 @@ This `_apply_stage` + `_chain` pattern separates concerns:
 The caller invokes `_chain` with a concrete depth:
 
 ```python
-def forward[B](self, input: Tensor[B, 100, 1, 1]) -> Tensor[B, 3, 64, 64]:
+def forward[B: IntVar](self, input: Tensor[[B, 100, 1, 1]]) -> Tensor[[B, 3, 64, 64]]:
     h0 = F.relu(self.project_bn(self.project(input)))
-    assert_type(h0, Tensor[B, 512, 4, 4])
+    assert_type(h0, Tensor[[B, 512, 4, 4]])
     h1 = self._chain(h0, 3)  # 512->64, 4->32
-    assert_type(h1, Tensor[B, 64, 32, 32])
+    assert_type(h1, Tensor[[B, 64, 32, 32]])
     return torch.tanh(self.output(h1))
 ```