Add Shape Types To Torch Model

You are porting a PyTorch model to use pyrefly's tensor shape type system.

This usually works. The example corpus contains dozens of substantial, explicitly scoped ports of real open-source research models, including transformers, convolutional encoder-decoders, autoregressive architectures, and dynamic module stacks. A first checker pass with errors or gradual tensors is a diagnosis, not evidence that the model is unportable. Work module by module and expect a few short probe–fix–check iterations before accepting a boundary.

When subagents are available, use them for bounded, independent read-only work: map the source's stated shapes, audit the model's ops against the stubs, or cluster checker errors by root cause. The main agent still owns the port and integrates the evidence; do not have several agents edit the same file.

The gates below prevent the common failure mode of writing the whole port from reasoning and checking only at the end. Apply their rigor according to the deliverable:

• Annotating someone's own model (the common case): use the source evidence, op audit, probes, and verification below, but keep working notes private and checkpoint shape-changing locals and module boundaries rather than every trivial assignment. Hand back the annotated model plus a short report of what could not be tracked.
• Contributing a reference example (e.g. into a maintained example corpus): this is the exhaustive case. Record the audit table, probe every tensor local, keep per-local assert_type coverage, and provide the receipts and counts because other people use the port as reference material. A skill that invokes this one for corpus work will say so.
• Migrating an existing production codebase: preserve its structure and tests. Prioritize precise component inputs and outputs; internal inference is valuable evidence for those contracts but is secondary when the body is highly dynamic. Keep gradual regions small and recover precise shapes at their exit. Inventory existing shape annotations and public tensor boundaries, but do not add inventory comments, an assert_type after every local, or model-file smoke tests. Add focused static tests only for reusable stub behavior. Use temporary reveal_type probes as needed and remove them before handoff. Existing test and lint conventions take precedence.

Existing jaxtyping annotations

If the goal is native-syntax conversion, inventory every jaxtyping annotation and its intended shape, including local annotations. If the goal is shape checking with minimal production churn, consider declaring dimension names with @static_jaxtyping("...") and keeping the existing annotations instead. Choose from the request and repository context; do not convert syntax merely because a conversion is possible.

After a native conversion, compare the inventory against the new annotations: every removed annotation must have a native counterpart with equal or better precision, or a documented reason why that precision is not representable.

Treat each jaxtyping annotation as an information-preservation obligation, not merely syntax to remove. Preserve its rank, literal dimensions, named-dimension equalities, arithmetic relationships, and variadic-prefix semantics whenever native shape types can express them. A clean type check is not sufficient if a precise jaxtyping contract became gradual.

• Never delete a local annotation merely because current inference appears sufficient. Preserve it with native syntax and reuse any enclosing IntVar or IntTuple parameters so it continues to document and check the intended relationship.
• Convert a variadic prefix such as "*B D" using Bs: IntTuple and Tensor[[*Bs, D]]; do not collapse it to a bare Tensor. Use *Elements[Bs] only when annotations evaluate eagerly at runtime.
• When a dimension is known only from a runtime value and cannot be related to an input type parameter, preserve its rank and documentation with a named output-only IntVar, for example Tensor[[B, NewH, NewW]]. Pyrefly instantiates such unconstrained variables as gradual dimensions at call sites. Use int only when a useful name is unavailable, and bare Tensor only when even the rank is genuinely unknown. Record every loss of precision in the before/after audit.
• Do not weaken a public signature because an implementation detail is untracked. First try a more precise stub or a narrow cast immediately around the untracked operation. Keep all dimensions bound by parameters on both sides of that boundary. A runtime rewrite is separate work and requires explicit user direction.
• Do not assume an empty jaxtyping shape string accurately describes a scalar. Check how the value is used. If it accepts arbitrary ranks, bind its complete shape with Shape: IntTuple and Tensor[[*Shape]]; if the code requires trailing dimensions, spell those out after a variadic prefix.
• Treat code that reads annotations at runtime—including schema validators—as executable behavior. Adapt and test that consumer separately rather than assuming a syntactically equivalent annotation preserves validation.

Before you start: resolve two setup choices

Resolve these from the repository and invoking skill whenever possible. Do not interrupt the user when the checkout, config, or request already answers them; ask only when a real ambiguity remains.

1. Choose the check command and validate the environment. Default to pyrefly check and inspect the resolved paths with pyrefly dump-config. The Torch stubs are partial overlays, so Pyrefly must also see a real torch installation through the selected interpreter or --site-package-path. When the model uses einops, add the pyrefly-einops-stubs root too; without that overlay, an einops transform can incorrectly appear to preserve its input shape. Before porting, check one known-good nearby example and require 0 errors. If it fails, repair the environment rather than weakening the model annotations. Also record the target boundary's own pre-port error count before annotating; pre-existing errors are fixed or preserved as you go, but only the baseline tells them apart from errors you introduce. In a repository, prefer its documented check command; an invoking skill may provide all paths, as the corpus skill does.
2. Decide whether stub changes are in scope. For someone's model or a production migration, default to leaving shared stubs unchanged and report gaps. For a maintained reference-corpus contribution, stub improvements are in scope by default. If the user explicitly asks for or rules out stub work, follow that. Ask only when neither context nor request establishes the scope.

A no-stub-change port is a supported outcome, not a failure. If stub changes are in scope, port first and confirm each gap with reveal_type before editing a stub. Do not decide about deeper Pyrefly shape-logic changes up front; react only if a computed shape is provably wrong.

Pre-flight

Create tasks for each gate and the module loop/verification phases that follow. Update them as you complete each stage — this gives visibility during long-running ports.

Complete these gates before writing any code.

Gate 0: Understand the system

Read shape_tracking_capabilities.md and porting_principles.md (this skill dir). They explain the tracking mechanisms, current API, priority order, and stub philosophy.

Then skim the model index at the end of style_guide.md and open one to three closest examples under pyrefly-torch-stubs/examples/. Use them as evidence that substantial real models can be ported and as templates for architecture-level patterns such as attention, dynamic module stacks, encoder-decoders, and variadic batches. Existing prose can be stale: only a current assert_type, focused probe, or passing check proves the present behavior. Do not copy annotations blindly; the target source and checker remain authoritative. Save the full style-guide comparison for verification.

Gate 1: Audit ops

Inventory every tensor-producing or shape-changing expression in the selected model boundary: module construction and calls, Tensor methods, operators and indexing, torch/F. functions, module-instance methods, third-party APIs, and dynamic dispatch. For Torch APIs, inspect the relevant partial overlay module and read the return annotation; if it calls a type-level shape function, confirm that function in _shapes.pyi. Special handlers live separately in Pyrefly's Rust implementation, so confirm them with source or a minimal probe.

Classify each expression as tracked-stub, tracked-DSL, tracked-handler, gradual Tensor, Any, unavailable-overlay-symbol, or third-party-boundary. The distinction matters: these are partial PEP 561 overlays, and a public name omitted from an overlaid module is unavailable rather than guaranteed to fall back gracefully. A declared loose return may produce a gradual Tensor, while an undeclared Tensor member may flow through __getattr__ as Any.

This is a diagnostic pass, not a blocker. A subagent may collect the source operation inventory, but the primary agent must validate the final stub/DSL/handler classification against the local checkout. Do not use web search as authority for current support.

For a reference-corpus contribution, include the audit table in the work log before Gate 2. For ordinary and production ports, keep it as concise working notes unless the user asks for the full audit:

## Gate 1: Ops audit
| Op | Stub location | Shape function in the return annotation (or "none") | Status |
|----|---------------|------------------------------|--------|
| nn.Conv2d | tensor-shapes/pyrefly-torch-stubs/torch-stubs/nn/__init__.pyi — generic [InC,OutC,K,S,P,D] | none (generic signature) | tracked-stub |
| F.adaptive_avg_pool2d | tensor-shapes/pyrefly-torch-stubs/torch-stubs/nn/functional.pyi — `Tensor[adaptive_pool2d_shape(Shape, OH, OW)]` | `adaptive_pool2d_shape`, defined in `_shapes.pyi` | tracked-DSL |
| ...

Status is one of tracked-stub, tracked-DSL, tracked-handler, gradual (a declared bare return), Any, unavailable-overlay-symbol, or third-party-boundary.

For a stub or DSL result, read the declaration's return annotation; if it calls a shape function, confirm that function exists in _shapes.pyi. Write "none" only after confirming that the signature's own generics compute the shape.

A non-precise result is information, not a blocker, but its exact category determines the response. Preserve known rank and dimensions in the public contract, isolate a dynamic or third-party boundary with a precise cast when the shape is justified, and use bare Tensor only when rank itself is genuinely unknown. An unavailable overlay symbol may require a stub declaration just to type-check; it does not necessarily degrade to Tensor. If stub work is in scope, consider a general stub improvement.

Gate 2: Scope and inventory the original

Define the port boundary before inventorying it. For a single-file model, the boundary is usually that file. For a repository containing many model families (for example, a policy zoo or a framework that builds models from YAML), choose the smallest coherent requested model plus its transitive custom tensor modules; do not inventory unrelated training, data, CLI, or deployment code. State the boundary and proceed. Ask only if several incompatible model families are plausible and the request gives no way to choose. A corpus-invoking skill may require one complete upstream file or another explicit boundary. When the boundary is smaller than its package, state the exclusion rule alongside it: which files are out and why, so coverage is visible without opening the report.

Mine the source for shape evidence before inventing annotations:

• existing jaxtyping or other tensor annotations;
• docstrings and inline comments such as (B, C, H, W) or reshape arrows;
• shape destructuring, runtime shape assertions, view/reshape arguments, einsum or einops equations, config dimensions, and test fixtures;
• layer constructor dimensions and equalities implied by residuals, matmuls, concatenation, and split sizes.

Turn this evidence into a short shape map and use its names consistently. Comments are hypotheses, not proof: resolve conflicts by following the actual operations and tests, then preserve or correct useful comments in the port. This is especially valuable in research models, which often document shapes more precisely than their Python annotations do.

For a reference-corpus port, write the inventory as a comment block at the top of the port file using this format. For an ordinary port, keep the same list in working notes; production migrations should inventory only changed public boundaries and existing annotations.

# ## Inventory
# - [ ] ClassName.__init__ — Int: param1, param2; int: param3
# - [ ] ClassName.forward
# - [ ] function_name — utility, no tensors
# ...

Every class, function, and method inside the selected port boundary must appear, and every inventory item must be ported. Do not silently narrow the boundary after discovering a hard dependency. A dependency outside the boundary may remain behind a precise cast, typed interface, or gradual boundary. If stub work is in scope, add only general improvements that help more than this model.

For each class, list constructor parameters and whether each is Int or int. Also record configuration or mode flags that change tensor rank, layout, return arity, or container type (training/eval, cache, export, optional backends). Use Literal overloads when the existing API exposes a statically knowable mode; otherwise preserve the honest union or gradual boundary. Runtime assertions support preconditions but do not by themselves turn unrelated int values into symbolic dimensions.

Check off corpus inventory items as you port them. Do not proceed to verification with unchecked items inside the declared boundary.

Phased annotation (larger codebases)

For a single-file or few-module port, one commit is usually fine. When the boundary spans several files or many modules, consider annotating in phases on a branch, one commit per phase. Phases order the work; they do not shrink the boundary — every inventory item is still ported. Order phases by dependency: leaf modules and shared contracts first, then the modules that consume them. Run the module loop over one coherent subset per phase and require a 0 errors check before committing, so every commit is a working checkpoint and later phases never debug earlier ones. Commit messages should state the phase's coverage (which modules, public-contract coverage delta, remaining gradual boundaries); keep reveal_type probes and working notes out of the commits. For production migrations this also keeps each diff reviewable and landable on its own.

Transition to module loop

You have completed pre-flight. You have NOT written any model code yet.

Your analysis so far is a hypothesis. The module loop tests it empirically. If you catch yourself planning multiple modules' forwards in your head, STOP — you are substituting reasoning for testing, and reasoning is less reliable.

Write the file with imports, the inventory comment, and utility functions (no tensor shapes). Then start Step 1 for the FIRST module only.

Read porting_principles.md (this skill dir) for the mindset: why we port, priority order, and stub philosophy.

Module loop

Repeat the following for each module in dependency order. Each module's typing may inform the next — e.g., discovering that a submodule tracks shapes internally changes how the parent handles its loop.

ONE MODULE AT A TIME. Complete Steps 1–6 for module A before starting module B; in the corpus case, record the Step 6 checklist at that point. If you find yourself typing two modules' constructors before running the checker, you have already entered the primary failure mode: writing the entire file and validating at the end leads to hidden gradual boundaries.

Step 1: Inventory parameters

List every constructor parameter. For each, decide:

• Int: the value determines a tensor dimension (flows to nn.Linear, nn.Conv2d, tensor creation, or any typed function that uses the value as a shape dimension).
• int: iteration count (n_layers, n_res_block) or boolean-like flag.

If in doubt, make it Int. The cost is one more type param; the cost of int is permanent shape loss in everything downstream.

Critical rules:

• A dimension that enters new typed code should be Int[X], not int. In an existing public API, changing int to Int[X] narrows the static contract; get the user's agreement first. Otherwise preserve the signature and recover precision at the first boundary where the dimension is justified. When no user can be asked (a headless trial, eval, or bulk migration), default to preserve-and-recover: keep the signature, add a bridge dim or cast, and note each site in the report. Never narrow a public contract silently.

• len(tensor) returns plain int, while tensor.size(dim) can preserve a symbolic Int. In annotation-only work, keep the original expression and restore precision at a downstream boundary; suggest the equivalent size spelling separately if it would improve inference.

• Int is a subtype of int, so int(dim) erases tracking when dim is already an int at runtime. Do not remove conversions around floats, tensors, round(), or expressions such as seq**0.5; those conversions can affect runtime behavior and are expected gradual boundaries. Arithmetic mixing a plain int with an Int produces plain int.

• Derived dims use expressions (D // NHead, 4 * ES), not independent type params. Only independent degrees of freedom get type params. Keep preconditions such as divisibility, perfect-square sequence lengths, nonempty scans, and mask-cardinality equalities separate from shape types; runtime assertions document them but do not justify simplifying (H // P) * P to H.

• Structured inputs and mode-dependent outputs. Inventory heterogeneous tensor mappings and flags for training/eval, cache, export, or optional backends. Use an existing TypedDict/dataclass or honest union when available. Introducing a narrower public container type or overload requires the user's agreement that it matches the real API.

• Dimensions from list[int] or other untracked sources. Element access erases the value to int. Preserve the source and restore precision with a justified cast or typed interface when the value reaches a known component boundary. Adding an explicit Int config field or constructor parameter changes the API; suggest it as separate work unless the user requested it.

• Keep existing register_buffer, register_parameter, dynamic factories, and container construction. Replacing them with nn.Buffer, explicit attributes, or typed classes is a runtime refactor, not annotation work.

• Bridge dims. When part of the model is untracked (e.g., features built via nn.Sequential(*list)), look for dimensions that connect the untracked section to tracked downstream modules. For example, if features output feeds a Linear classifier, the Linear's in_features is a bridge dim — making it a class type param enables annotation fallback to recover a shaped type (e.g., Tensor[[B, LastC]]) that then flows naturally through downstream ops. Without it, annotation fallback can only recover bare Tensor or batch-only shapes.

  class Model[NC: IntVar, LC: IntVar](nn.Module):
      def __init__(self, num_classes: Int[NC] = 1000,
                   last_channel: Int[LC] = 1280):
          ...
          self.classifier = nn.Linear(last_channel, num_classes)

  Here LC bridges the untracked feature extractor to the typed classifier, recovering Tensor[[B, NC]] at the output.

• Int[X] | None for optional dimensions. In new code, a genuinely optional shape-bearing parameter can use Int[X] | None and narrow with if value is not None:. Replacing an existing public Optional[int] with Int[X] | None narrows its static contract; do so only with user agreement. Otherwise preserve the public type and recover a justified shape downstream.

• Parameterized config dataclasses. When multiple modules consume dimensions from the same @dataclass config, note it — Step 2 shows how to parameterize the config so dims propagate across module boundaries.

• Lazy-initialized buffer attributes. An attribute's type is fixed at its declaration. If self.x: Tensor | None = None is assigned a shaped tensor in a later setup hook, reads after narrowing retain only bare Tensor. Preserve the lifecycle and record this as a dynamic boundary. Eager initialization can recover the shape, but it changes runtime structure and requires separate user direction.

• Uppercase locals shadow type parameters. A local named B, C, or D can shadow a type parameter of the same name inside annotations evaluated in that scope (including assert_type shape strings), producing invalid-annotation. Rename the local behavior-neutrally (B_ssm, lowercase destructuring); call sites are unaffected.

Step 2: Type the constructor

Write __init__ with the Int params from Step 1. Construct sub-modules using those Int params — they get typed automatically.

Default values for Int params: a literal constructor default can be used with Int[X] directly. A PEP 696 default is optional: it makes the unspecialized class name carry that default at the type level.

# Works; callers bind NC from the argument or literal default:
def __init__(self, num_classes: Int[NC] = 1000): ...

# Also works; bare Model means Model[1000]:
class Model[NC: IntVar = 1000](nn.Module):
    def __init__(self, num_classes: Int[NC] = 1000): ...

Dataclass field defaults are different: dim: Int[D] = 768 currently needs a specific # type: ignore[pyrefly:bad-assignment].

Constructor patterns that create expected shape-typing boundaries:

• nn.Sequential(*list_var) and factory functions returning nn.Sequential erase the member module types at the dynamic boundary.
• getattr(nn, name)() and YAML/config-selected module factories return Any or a broad module type.
• Heterogeneous module containers cannot express a different shape transform for each runtime index.
• Method-level type params on class fields. A field cannot retain a type parameter scoped only to the method that assigned it.
• Non-generic subclass of a shape-generic base. The subclass binds the base's type arguments to their defaults, so explicit constructor values that conflict (for example, stride=2 against an inherited Int[1]) are correctly rejected — accepting them would be unsound, since the object is statically the default specialization. Respond by annotating attributes with the precise base type (self.conv: nn.Conv1d[...] = Subclass(...)), which recovers exact shapes at use sites. Specializing the base explicitly (class Sub(Base[2])) silences the constructor conflict but attribute reads still come back gradual, and a bare ignore is worse: the object keeps the default shapes, which then silently contradict the runtime values.

Preserve known dimensions in typed public interfaces around these boundaries. Do not extract modules, replace factories, or redesign containers solely for shape tracking unless the user separately requests that rewrite.

Parameterized config dataclasses. For a new/corpus model, or when the user agrees to the narrower static API, a @dataclass holding shared dimension hyperparameters can be generic so dimensions propagate through constructors. For an existing public production config, preserve its annotations by default and recover shapes at component boundaries instead:

@dataclass
class Config[D: IntVar, NHead: IntVar, VocabSize: IntVar]:
    dim: Int[D]
    n_head: Int[NHead]
    vocab_size: Int[VocabSize]
    dropout: float = 0.0

Modules extract only the params they need using Any for the rest:

class MLP[D: IntVar](nn.Module):
    def __init__(self, config: Config[D, Any, Any]):
        super().__init__()
        self.fc = nn.Linear(config.dim, 4 * config.dim)

Without this, each module must independently accept and thread every dim through its constructor — error-prone and verbose.

If the original config had default values (e.g., dim: int = 768), combine the two patterns above — give the dataclass type params PEP 696 defaults so callers can omit dims:

@dataclass
class Config[D: IntVar = 768, NHead: IntVar = 12, VocabSize: IntVar = 50257]:
    dim: Int[D] = 768  # type: ignore[pyrefly:bad-assignment]
    n_head: Int[NHead] = 12  # type: ignore[pyrefly:bad-assignment]
    vocab_size: Int[VocabSize] = 50257  # type: ignore[pyrefly:bad-assignment]
    dropout: float = 0.0

Two different defaults are at play here, and only one is clean:

• The PEP 696 defaults on the type params ([D: IntVar = 768, ...]) are what let callers omit dims — those need no ignore.
• The dataclass field literal defaults (dim: Int[D] = 768) still need # type: ignore[pyrefly:bad-assignment], because a plain int literal is not assignable to Int[D]. This is the accepted corpus pattern (see examples/finalmlp.py). Note that constructor-parameter defaults (def __init__(self, num_classes: Int[NC] = 1000)) do not need the ignore — only dataclass field defaults do.

Now Config() produces Config[768, 12, 50257] and Config(dim=1024) produces Config[1024, 12, 50257] — dims propagate even when callers don't pass every parameter.

DO NOT write the forward method yet. The forward signature and assert_type expressions depend on what the checker infers, which you don't know until Step 3.

Run the checker to verify the constructor compiles. In the corpus case, record that output before proceeding; otherwise keep iterating without dumping routine success output.

Step 3: Probe the forward

For a corpus contribution, count and probe every tensor local in the forward method. For an ordinary model, probe each shape-changing local and every point where the shape could be lost; include enough checkpoints to cover each operator chain and module boundary. For production code, use targeted temporary probes around changed boundaries and checker failures.

Add reveal_type at those points and run the checker. In the corpus case, record the results in this format; elsewhere keep concise working notes:

# reveal_type results for ClassName.forward:
# Locals: N (list them: var1, var2, var3)
# var1 (line N): Tensor[[B, C, H, W]]  → SHAPED
# var2 (line M): Tensor                 → BARE — investigate in Step 4
# var3 (line P): Tensor[[B, D]]         → SHAPED

For a corpus contribution, verify that the reveal count matches the tensor-local count. In every mode, do not continue while a shape-changing path is unprobed.

If a reveal_type result contradicts your understanding of the op (e.g., spatial dims unchanged after a strided conv, or a shaped op returning bare), write a small isolating test, run the checker, and confirm the behavior before proceeding. Either your understanding is wrong (update your mental model) or the checker has a simplification you should document. If a second pass still has unexplained plain or unchanged Tensor results, re-run the environment sanity check and confirm the einops search path before declaring the code untrackable.

This table is your Step 4 input. Do not write assert_type until Step 4 is complete for every BARE entry. The results tell you:

• Shaped type → the checker tracks this op. Write assert_type in Step 5.
• Bare Tensor → shape lost. Investigate in Step 4 before deciding.