Tensor shapes docs use shape syntax that no longer checks against the shipped stubs

[ting-hong-shieh]

Describe the Bug

The tensor shapes documentation teaches a shape syntax that the shipped stubs no
longer accept. Code written the way the docs show it does not check.

The current stubs want IntVar-bounded type parameters and a single shape list
argument to Tensor:

class Linear[IN: IntVar, OUT: IntVar](Module):        # torch-stubs/nn/__init__.pyi
    weight: Tensor[[OUT, IN]]

The docs use unbounded parameters and a flat argument list — class Model[D],
Tensor[B, T, D]. Across website/docs/tensor-shapes*.mdx there are 81
occurrences of the flat Tensor[...] form, 31 declarations with unbounded shape
type parameters, and no IntVar bound on any of the pages listed below.

Reproduction. pyrefly 1.2.0 and pyrefly-torch-stubs 1.2.0, both from PyPI,
Python 3.12.

The step-by-step model from
Tutorial: the basics,
copied as the page builds it up:

class BaselineActor[S, A](nn.Module):
    def __init__(self, state_size: Int[S], action_size: Int[A]) -> None:
        super().__init__()
        self.fc1 = nn.Linear(state_size, 400)
        self.fc2 = nn.Linear(400, 400)
        self.out = nn.Linear(400, action_size)

    def forward[B](self, state: Tensor[B, S]) -> Tensor[B, A]:
        h1 = F.relu(self.fc1(state))
        assert_type(h1, Tensor[B, 400])
        h2 = F.relu(self.fc2(h1))
        return torch.tanh(self.out(h2))

ERROR `S` must be an `IntVar` to be used as a shape dimension. [invalid-annotation]
ERROR `A` must be an `IntVar` to be used as a shape dimension. [invalid-annotation]
ERROR `B` is not assignable to upper bound `IntTuple` of type variable `Shape` [bad-specialization]
ERROR Expected 1 type argument for `Tensor`, got 2 [bad-specialization]
ERROR assert_type(Tensor[[*tuple[int, ...], 400]], Tensor[B]) failed [assert-type]
 INFO 8 errors

The Int[X] | None example from
Reference fails the same
way, with Expected 1 type argument for Tensor, got 3. The same page's "Usage
patterns" table prescribes the old form directly:

Pattern	Purpose
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

Rewriting either example to match the stubs — [S: IntVar, A: IntVar],
Tensor[[B, S]] — checks clean, so this looks like documentation lagging an API
change rather than a checker bug. The repository's own examples, e.g.
tensor-shapes/pyrefly-torch-stubs/examples/nanogpt.py, already use the new
form.

Scope. Counts of the flat Tensor[...] form per page:

Page	Occurrences
tensor-shapes-tutorial-architectures.mdx	19
tensor-shapes-reference.mdx	17
tensor-shapes-tutorial-basics.mdx	12
tensor-shapes-tutorial-loops.mdx	12
tensor-shapes.mdx	11
tensor-shapes-tutorial-advanced.mdx	9
tensor-shapes-ai-porting.mdx	1

I verified two of these pages by running them; I have not run every example. The
examples on tensor-shapes.mdx are exposition rather than code to copy, but they
still show the old form.

The inlay-hint screenshots under website/static/img/tensor-shapes/ show the old
syntax too (def forward[B, T](self, x: Tensor[B, T, NEmbedding])), so they need
retaking rather than editing.

Question before anyone starts on it: is this the same underlying change as
#4153, and is a rewrite already underway? I'm happy to do the mechanical pass
across the tutorials and the reference page if that's useful, but not if it would
collide with work in progress.

Sandbox Link

No response

(Only applicable for extension issues) IDE Information

No response

Disclosure, per the AI Usage section of CONTRIBUTING.md: the counts, the reproductions, and this report were produced by an AI agent (Claude Code) working in my checkout. I reviewed them before posting.

[Vishwaspatel2401]

Hi @ting-hong-shieh I would like to take this on and work on it

[ting-hong-shieh]

@Vishwaspatel2401 go ahead, it's yours — I filed this to get it recorded rather than to reserve it. Here's what I already worked out, so you don't have to redo it.

Setup for checking the examples. No stub copying needed; the packages are on PyPI and version-matched to Pyrefly:

pip install pyrefly-torch-stubs

Then pyrefly check yourfile.py with the interpreter pointed at that environment. That's enough to run any example from the docs.

The syntax change to apply. Shape type parameters need an IntVar bound, and Tensor takes one shape list rather than a flat argument list:

# what the docs currently show
class Model[D](nn.Module):
    def forward[B](self, x: Tensor[B, D]) -> Tensor[B, D]: ...

# what the current stubs accept
class Model[D: IntVar](nn.Module):
    def forward[B: IntVar](self, x: Tensor[[B, D]]) -> Tensor[[B, D]]: ...

tensor-shapes/pyrefly-torch-stubs/examples/nanogpt.py is the best in-repo style reference — it's maintained against the current stubs and covers the trickier shapes (*Elements[Bs], arithmetic like NEmbedding // NHead).

Two things to be careful about.

tensor-shapes.mdx is different from the tutorials. Its examples are exposition — def mm[M, K, N](...) is described as a library-side stub signature, and one block is explicitly labelled "Internal library definition — not user-facing code". They still show the old form, but they aren't examples a reader copies, so they may want different treatment from the tutorial pages. Worth asking rather than assuming.

The inlay-hint screenshots under website/static/img/tensor-shapes/ also show the old syntax, so those need retaking rather than editing. That's probably out of scope for a docs pass.

Still unanswered: whether this overlaps #4153, which is assigned to @stroxler. I'd wait for a maintainer to weigh in before doing the full sweep across all seven pages — if a rewrite is already underway, the mechanical pass would collide with it.

One page is already handled: #4559 fixes the setup page's config section and its Hello world example, so you can skip tensor-shapes-setup.mdx.

Disclosure, per the AI Usage section of CONTRIBUTING.md: the investigation behind this issue and this comment were produced by an AI agent (Claude Code) working in my checkout. I reviewed them before posting.

[Vishwaspatel2401]

#4153 overlap: none — that's about the sandbox files (website/sandbox-examples/tensor-shapes-overview/), different error (TypeVarTuple), and stale (no activity since 2026-08-11). Proceeded independently.

tensor-shapes.mdx: fixed the exposition examples (Linear[N, M], custom_rand_tensor, the torch.mm stub example — checked that one against the real stub, it was stale too) and the prose syntax mentions. Left the repeat_ir DSL block (different sublanguage, never subject to this bug) and the Pyre/Pyright comparison (deliberately shows a different, older system) untouched.

One thing beyond the original scope: the variadic convention (TypeVarTuple/*Bs) doesn't just need brackets, it fails outright now. Replaced with the Bs: IntTuple + Tensor[[*Elements[Bs], ...]] form nn.Linear.forward actually uses — affects the reference page's variadic section and the Linear[N, M] example here.

Skipped tensor-shapes-setup.mdx — that's #4559.

Two unrelated things noticed but not fixed (out of scope): a stale Literal[512] claim in the DCGAN section of the advanced tutorial, and a couple of genuine algebraic-gap limitations in the recursive architecture examples (same class of gap that page's own "Algebraic gaps" section already documents).

Verified against pyrefly==1.2.0 / pyrefly-torch-stubs==1.2.0 from PyPI. Full test plan in the PR.

---

Disclosure per CONTRIBUTING.md: produced by an AI agent (Claude Code), reviewed before posting.