Steven Troxler

Partition legal positions for `TypeVar` vs `SymVar` kinds Summary: **Context** This commit wraps up a set of commits implementing a new `SymVar` type variable "kind" (where by kind I mean like TypeVar vs TypeVarTuple vs ParamSpec). The idea is that TypeVar vs SymVar can appear legally in different places and have different syntactic and semantic rules - their domains and uses are in fact disjoint, just as TypeVar and TypeVarTuple are disjoint (TypeVarTuple can only appear in unpackings). This commit is the final cutover which partitions where the TypeVar vs SymVar kinds are legal: all use of SymVar is banned in normal type position, whereas all use of TypeVar is banned in symbolic-integer positions (inside of `Dim[...]` and inside of the shape in `Tensor[[...]]` or `ndarray[[...]]`, other than when unpacked via `*Elements[...]`) **This diff** Pyrefly now partitions legal quantified-variable positions between ordinary type contexts and symbolic shape contexts. Bare variables used in `Dim[...]`, shaped-array shape lists, `SizeTuple[...]`, or scalar shape expressions must have `SymVar` kind, so normal `TypeVar` parameters are rejected there. Conversely, `SymVar` parameters are rejected in ordinary type positions, keeping symbolic dimensions from flowing into normal type arguments. The remaining exception is the legacy `Generic[...]` spelling needed for existing class generic plumbing: it can still mention `SymVar` so classes can bind symbolic dimensions, but uses of those variables are checked according to the partitioned context they appear in. The existing `SizeTuple` carrier path for `*Elements[Shape]` is also preserved for variadic shape unpacking. This tightens the invariant started by the shape-arithmetic cutover: scalar entries in `Type::Dim`, `Type::Size`, and shaped-array shapes come from `SymVar`-kind variables rather than ordinary `TypeVar`s, while ordinary type slots remain reserved for normal type variables. The tensor-shape stubs, examples, and corpus are updated to mark symbolic scalar shape parameters as `SymVar`, including module/stub signatures whose return shapes depend on scalar dimensions. The new shape DSL tests cover rejected ordinary `TypeVar` shape uses, rejected `SymVar` ordinary type uses, the preserved `Generic[...]` compatibility path, and the preserved `SizeTuple` carrier behavior. Reviewed By: kinto0 Differential Revision: D111140929 fbshipit-source-id: f3a0e908b3dbb3f0e241a9bbb24e6f95b053cb30

Enforce `SymIntVar` solver answers as `SymInt` Summary: `SymIntVar` is now a distinct symbolic-integer kind, so solver state should not be able to specialize it to an ordinary type. This change normalizes explicit `SymIntVar` class arguments, bounds, and solver answers through a shared symbolic-integer recovery path, preserving ordinary `TypeVar` behavior while making `SymIntVar` fall back to gradual `SymInt[int]` after diagnostics instead of leaking non-symbolic answers. It also makes mixed `TypeVar`/`SymIntVar` var unification preserve the `SymIntVar` root so the kind invariant survives inference. The recovery behavior mirrors the rest of Pyrefly kind checking: invalid user syntax still gets an error, but downstream types are kept structurally valid. Focused tests cover bad specialization, bad call recovery, ordinary `TypeVar` solving to `SymInt`, `SymIntVar` inference chains, bounded `TypeVar` interaction, and shaped-array overload behavior. Reviewed By: kinto0 Differential Revision: D111623181 fbshipit-source-id: 6716aa762e85b6c57c4e8a4fc184160834e7a4e0

Restrict shape integer parsing to shape contexts Summary: This code is aimed at making it safe to use `tensor-types = true` in an existing codebase, if shape stubs are not being used, by making the tensor-shape-specific behaviors in Pyrefly all require explicit opt-in via stubs `tensor-shapes = true` should be an fbcode no-op for ordinary Python type forms unless the user opts into shape-aware stubs. This removes the fallback that converted integer literals in every generic type argument or `TypeVar` default into shape sizes, so ordinary generics keep rejecting `Box[3]` and `T = TypeVar("T", default=3)` while registered shaped arrays and `Dim[...]` continue to parse shape integers through their local shape parsers. Reviewed By: avikchaudhuri Differential Revision: D106722729 fbshipit-source-id: d712c59ffded1dacee69bd7f9fb2486df4d91fb5

Migrate torch shapes to SizeTuple carriers Summary: The torch shape corpus now uses the same carrier representation as the NumPy corpus: `Tensor` takes a single `SizeTuple`-bounded shape argument, and shape literals use the list-syntax sugar that Pyrefly normalizes only in shaped-array carrier slots. This removes the old TypeVarTuple surface from the trunk stubs and tests without adding dedicated `Tensor[2, 3]` sugar. This change is necessary if we want our shape type system to be able to represent a user class that is generic over multiple tensors with different shapes, which is necessary to handle many things we'll encounter in the library ecosystem (for example torchrec's `JaggedTensor`. It's impossible to spell that with type var tuples, because a class can only have a *single* type var tuple parameter (TypeVarTuple is the type level equivalent of `*args`, you can't have multiple variadic parameters/arguments by definition) Reviewed By: yangdanny97 Differential Revision: D110505959 fbshipit-source-id: d600e5acedce6175d21b12bf4acb6a4cdbaa9b31

Recover invalid `SymIntTuple` middles Summary: `SymIntTuple` can contain an unpacked middle that was valid when constructed but becomes invalid later through substitution or solver expansion. That mirrors `SymInt` symbolic leaves: the representation may be non-canonical briefly, but normalization must not let invalid structure leak into shape operations or display. Make `SymIntTuple::unpacked` the recovery boundary for those post-substitution cases. Concrete tuple middles preserve rank by replacing invalid elements with gradual `SymInt[int]`; unknown-rank gradual middles preserve known affixes; invalid non-tuple middles degrade to shapeless. At the same time, keep front-door tuple-carrier validation strict, so invalid user-written tuple carriers still produce annotation errors instead of being silently repaired. This also makes scalar dimension validation kind-aware so ordinary `TypeVar`, `ParamSpec`, and `TypeVarTuple` cannot be treated as scalar symbolic dimensions, while preserving whole-shape `TypeVar` carriers and variadic middles where they are valid. Reviewed By: kinto0 Differential Revision: D111584047 fbshipit-source-id: 1d1959958321da4ee62198a27e8c0335ef03cd5e

Plug user-reachable panics in the DSL evaluator Summary: Audit of most eval_dsl_expr / eval_dsl_body / eval_call / eval_binop for panic sites that can be triggered by user-authored shape_dsl_function definitions (as opposed to IR bugs or internal assertions), and fix each one. The idea here is to get reasonable behavior, not panics, when a DSL function is misused in a way that we somehow miss during the type checking phase. Ideally, this would never occur - we validate the DSL ahead of time, so only a bug in Pyrefly would actually make the code reachable (hence no unit tests), but sooner or later we probably will encounter a failure to validate the DSL (i.e. a soundness bug) at which point this change should make it easier to understand what went wrong. This is best-effort and a couple of panics remain where this is much more convenient, which is probably fine - again this is just to make debugging easier if we encounter problems in the future Reviewed By: rchen152 Differential Revision: D106003742 fbshipit-source-id: 1220f1322ee3663a76a0112791b5a834bb6df1bf

Tighten SymInt operator precision Summary: Symbolic integer operations should stay precise only when the result is known to remain a valid symbolic integer. In particular, exponentiation can leave integer shape algebra when the exponent is negative or only known dynamically, and unary integer operations should follow ordinary `int` behavior rather than manufacturing symbolic shape expressions. Add conservative sign reasoning for `SymInt`, use Python floor-division semantics during canonicalization, and reject provably negative exponents in symbolic integer annotations. Runtime `SymInt` exponentiation now preserves precision only for sound nonnegative exponent cases, returns `float` for negative exponents, and falls back to gradual results for unknown symbolic exponents. Unary `+`, `-`, and `~` on `SymInt` values route through normal `int` behavior. The floor-division canonicalization remains deliberately conservative: rewrites such as `(a // b) // c` are only applied when the outer divisor is proven safe, avoiding unsound rewrites for negative or potentially fractional power divisors. Tests cover the edge cases that drove the conservatism, plus `SymVar`/`TypeVar` partitioning for pow annotations. **Note on corpus changes** The corpus changes correctly expose a loss of accuracy due to our limited type system. As a rule, `a // b // c != a // (b * c)` - they are equal only if `c` is a positive integer. But because we lack constraints on integers, we cannot rule out that the second divisor in `a // 2 // 2**(Depth-1)` is less than 1. As a result, the simplification to `a // 2 ** Depth` does not hold under the current complexity of our model of integers. Note that `a // 2 ** (Depth -1) // 2` actually *would* simplify since we do know that `2` is a nonnegative integer. This loss of fidelity is actually correct in the sense that it is the only correct answer with an integer arithmetic system this constrained, and closely related to the known issues we have where `N // M * M` cannot be simplified without knowing constraints that the type system doesn't have. There are actually currently not one but two constraints here preventing Pyrefly from simplifying: - First, we have no ability to syntactically express knowing that `Depth > 1` - Second, because we lack *all* forms of reasoning over integers, even if we did know `Depth > 1` we would not currently have any rule to derive that `Depth - 1 >= 0`... we would need a nontrivial model of integer arithmetic (more than just canonicalization) to derive this knowledge Reviewed By: kinto0 Differential Revision: D111381832 fbshipit-source-id: f849d8b11a17cf30cbf5a30ec13afc8c44784c0a

Canonicalize SymInt after bound expansion Summary: Solver variable expansion can leave symbolic integer expressions in a non-canonical form when a bound variable is replaced by a concrete or gradual symbolic integer. That makes downstream shape reasoning inspect the wrong representation even though the semantic value is already known. Canonicalize only `Type::SymInt` nodes after `expand_with_bounds` resolves variables. The post-pass is intentionally narrower than general solver simplification, so expansion does not rewrite unions, tuples, callables, or other non-symbolic-integer type structure as a side effect. In effect, this is just "flattening" nested representations that can occur through the `Symbolic` branch after substitution or Var expansion, so that the canonical form is the "flat" form where all `Symbolic` branches are leaves. The regression coverage exercises literal folding, gradual propagation, quantified-leaf preservation, nested tuple-splice traversal, and a guard that non-`SymInt` types are not broadly simplified. Reviewed By: kinto0 Differential Revision: D111346482 fbshipit-source-id: 0b2505507aca874cc545ca5640c897470d9f3968

Rename Type::Size and SizeExpr to SymInt Summary: This is a pure rust refactor - I am not yet changing the display or the Python `shape_extensions.Size` spelling, just to keep the actual diff size manageable (I'll change that in the next commit); this commit is a pure rename + comment cleanup in rust code only. Reviewed By: yangdanny97 Differential Revision: D111297983 fbshipit-source-id: 839340dcbe21b127d8d33b8d4b656d343575e702

Gate residual reads by outermost expansion var Summary: **This stack** Adds a "callable residual" type that allows us to capture structure when a generic or overloaded callable is passed to a higher-order function. The residuals can later be "exploded" to recover generic and overload structure in outputs (either Callable outputs, or classes whose methods capture the callable structure of an argument - e.g. callback protocol classes). The full design details are [here](https://github.com/stroxler/pyrefly-docs/tree/main/callable-residuals) **This diff** Thread the outermost query var through expand-time recursion and only expose payloads when that query var is in ; non-target reads flatten residuals to their precomputed fallback. All reads of solver `variables` map go through this check - if two vars got unified via `union-find` but one of them isn't actually part of a comparsion that can produce residuals (like comparing a generic callable argument with a callable parameter in a higher-order function), then we will prevent the residual from appearing in the result and instead we will just flatten to a plain fallback type. The relevant read points are: - expand_var - record_recursive - is_subset_eq when we look up answers This allows us to cleanly enforce boundaries where residual types get flattened; without it, and expand call could potentially produce a residual somewhere where we aren't expecting it. This is only necessary becuase the current union-find algorithm conflates variables rather than setting bounds between them, and as a result the residuals can appear in places they don't belong; if we do a major refactor of the solver in the future this can go away. Reviewed By: rchen152 Differential Revision: D102011864 fbshipit-source-id: 68e1e4bbae4c39d1d7cb862bcd826bfb9f646b92

Tighten type signature of the `SizeExpr` data structure Summary: `SizeExpr` is the representation we use for shape arithmetic, but its arithmetic nodes still carried arbitrary `Type` operands. That made the type less descriptive than the invariant we want for `Type::Size`, and kept every consumer responsible for remembering which nested `Type` values were actually dimension leaves. This changes `SizeExpr` into a closed arithmetic tree with an explicit `Symbolic` leaf for scalar dimension symbols. Shape construction and tuple-carrier conversion normalize scalar dimensions into that internal form, while the variadic `SizeTuple` carrier middle remains a carrier type as before. Traversal, jaxtyping display, variance inference, and size subset comparison now recurse through the tree and unwrap symbolic leaves at comparison boundaries. Reviewed By: kinto0 Differential Revision: D111182215 fbshipit-source-id: 030832bd62ba48ffe379d39be03cac9b3226edd6

Store SymIntTuple with shape-specific data Summary: SymIntTuple should enforce shape invariants directly instead of relying on tuple storage conventions. This replaces tuple-backed storage with a private representation whose fixed dimensions are SymInt and whose variadic middle remains Type, while keeping tuple-carrier and ordinary tuple APIs as explicit projections. The constructor and projection paths preserve current recovery behavior for invalid carriers and keep the SymIntVar/TypeVar partition: scalar dimensions use SymInt/SymIntVar, while ordinary TypeVar carriers remain whole-shape middles. Nested tuple carriers are normalized recursively so concrete middles flatten and unbounded middles recover to gradual shape evidence. Reviewed By: kinto0 Differential Revision: D111623182 fbshipit-source-id: 0961c5c24799e77aae2543c514fafbbf21f04324

Add shape tuple carrier conversion groundwork Summary: **This Stack** I realized talking to Joren and looking at numpy stubs that numpy already has a limited kind of shape type, based on a type variable bounded by `tuple[int]`, where a concrete shape can be encoded as `tuple[Literal[...], ...]`. This support is limited since it only has access to PEP484 types so there are no symbolic shapes (just `int` to get rank-only checks) and there's no type-level arithmetic. But we should aim to be able to write numpy stubs that are compatible, which means not assuming the shape is a TypeVarTuple as in `Tensor[*Shape]`, but also allowing a tuple-valued regular type var as in `ndarray[Shape, Dtype]`. Assuming we go forward with this, it would mean pytorch shapes are written like `Tensor[3, 4, 5]` versus numpy being written like `ndarray[(3, 4, 5), float]` in shorthand, or `ndarray[tuple[Literal[3], Literal[4], Literal[5]], float]` when using a backward-compatible mode that works with existing type annotations. The changes shouldn't be too difficult and backwad-compatible (in other words the torch stubs won't have to change relative to the v1.1 release, which is good) **This commit** Add canonical conversion helpers between `ShapedArrayShape` and user-facing tuple carrier types so future NumPy-style shaped arrays can share the same internal shape projection as existing Tensor-style arrays. This keeps unsupported carrier elements explicit, maps literal dimensions through `Literal[n]`, unwraps `Dim[x]` carriers to raw internal dimensions, and normalizes unbounded tuple carriers to unknown-rank shapeless shapes. Reviewed By: yangdanny97 Differential Revision: D108936921 fbshipit-source-id: c682af8e8efecee4f74e0b11e62e34bbf25ffe4c

Require shaped-array registration Summary: Enabling `tensor-types = True` should be a no-op for projects that use ordinary PyTorch stubs. The shaped-array machinery now has an explicit opt-in via `shape_extensions.shaped_array(shape="Shape")`, so remove the remaining compatibility paths that treated any `torch.Tensor` as shaped. Keep PyTorch-specific efficiency lints working for normal `torch.Tensor`, but classify shaped arrays as PyTorch only when their concrete base class is actually `torch.Tensor`. Register the website sandbox tensor stubs explicitly so the shape examples continue using the new opt-in model. Shape-transform operations also need to compare shaped-array bases without double-checking the shape `TypeVarTuple`, because indexing and DSL transforms update `ShapedArrayType.shape` while the underlying class arguments may still carry the previous shape. Erasing just the shape `TypeVarTuple` for base-class compatibility preserves non-shape parameters such as dtype/device while leaving the actual dimension check in the shaped-array shape comparison. The registered shaped-array boundary also exposed a shape DSL binding assumption: DSL evaluation should use the call-site operand shape, not the bare `Tensor` parameter annotation used by fixture stubs. Record unhinted argument types for meta-shape binding, canonicalize registered class types only inside the DSL boundary, and allow unions whose members all have the same shaped-array shape to provide that shape. This keeps ordinary unregistered `torch.Tensor` inert while preserving shape information through control-flow joins in the shape-aware fixtures. Reviewed By: rchen152 Differential Revision: D106685115 fbshipit-source-id: a8cd3c3dc1df093d41d4657969a35b30739ed250

Rename SymInt/SymIntVar/SymIntTuple to Int/IntVar/IntTuple Summary: Rename the public shape-integer types from `SymInt`/`SymIntVar`/`SymIntTuple` to `Int`/`IntVar`/`IntTuple`. Once tensor-shape typing is adopted in a codebase it becomes a load-bearing part of that codebase, so these names should take high precedence, and `Int`/`IntVar`/`IntTuple` read better than the `Sym`-prefixed spellings. The name `Int` is coherent, not just shorter: the runtime values described by the `Int[...]` category of static types are exactly `int`. In particular the gradual/bare `Int` (≡ `Int[int]`) denotes exactly the set `int` in set-theoretic terms, differing from `int` only in static-typing behavior. Collisions on `Int` are possible but acceptable: a user can `from shape_extensions import Int as SymInt`, or `import shape_extensions as se; se.Int[...]` in a module that needs it — the usual name-collision inconvenience. The average user imports nothing named `Int`, so this is a win on average; collisions on `IntVar`/`IntTuple` are unlikely. This is a mechanical rename spanning the internal representation (`Type::Int`, the `Int` enum, `IntTuple`, `QuantifiedKind::IntVar`), the public `shape_extensions` API and `.pyi` stubs, type display, tests, and the tensor-shape corpus. The internal shape-DSL (`meta_shape_dsl`'s `DslType::SymInt` and the `shape_extensions.dsl.symint` surface marker) is intentionally left on the old spelling: it is slated to be reworked into type-level functions, so renaming it now would be throwaway churn. In unit-test fixtures that deliberately juxtapose builtin `int` with shape `Int`, identifiers use case to disambiguate (e.g. `returns_Int_from_int`), which reads clearly next to the annotations. Reviewed By: yangdanny97 Differential Revision: D111825021 fbshipit-source-id: e85d924d520514b1cf72cdc6d2d449ddda84406a

Only allow `SymVar` (not `TypeVar`) to use arithmetic ops Summary: **Context** This commit wraps up a set of commits implementing a new `SymVar` type variable "kind" (where by kind I mean like TypeVar vs TypeVarTuple vs ParamSpec). The idea is that TypeVar vs SymVar can appear legally in different places and have different syntactic and semantic rules - their domains and uses are in fact disjoint, just as TypeVar and TypeVarTuple are disjoint (TypeVarTuple can only appear in unpackings). This commit is the final cutover which partitions where the TypeVar vs SymVar kinds are legal: all use of SymVar is banned in normal type position, whereas all use of TypeVar is banned in symbolic-integer positions (inside of `Dim[...]` and inside of the shape in `Tensor[[...]]` or `ndarray[[...]]`, other than when unpacked via `*Elements[...]`) **This diff** Pyrefly already distinguishes scalar symbolic dimensions from ordinary type variables at binding time, but shape expression parsing still allowed ordinary `TypeVar` symbols to participate in integer arithmetic. That undermines the staged `SymVar` migration because expressions like `Tensor[[N + 1]]` and `Dim[N + 1]` can keep treating any type variable as a dimension variable. Track whether a dimension expression is being parsed in a bare shape position or in arithmetic. Bare positions continue to accept ordinary type variables for this stage, preserving the incremental migration, but operands of unary and binary shape arithmetic must now be `SymVar`-kind. Wrapper forms such as `D[...]`, `D(...)`, and `Dim[...]` preserve the incoming context so they cannot launder an ordinary type variable into arithmetic. The tensor-shapes corpus now marks the symbols that participate in shape arithmetic or numeric class shape arguments as `SymVar`, and convolution stubs mark spatial method parameters as `SymVar` because their return types compute output sizes. Reviewed By: kinto0 Differential Revision: D111135442 fbshipit-source-id: 03e1fc41316e7780e7e4c884fcf5823d5cdb9839

Rename internal tensor model Summary: The internal shape-carrying type is no longer specific to PyTorch: registered classes such as `numpy.ndarray` and eventually `jax.Array` can now use the same machinery. Rename the internal `Tensor*` model to `ShapedArray*` so the code distinguishes Pyrefly’s library-neutral shape abstraction from concrete user-facing classes like `torch.Tensor`. This is intended to be mechanical and behavior-preserving; the DSL spelling and existing `tensor_shapes` feature naming remain unchanged for now, and rendered types still use the concrete base class name. Reviewed By: avikchaudhuri Differential Revision: D106675657 fbshipit-source-id: 4637e3ebe7574ff00e4ab0c3313a0428dcb47165

Avoid per-origin `shape_extensions` misses Summary: Pyrefly now derives tensor shape availability by checking whether `shape_extensions` is resolvable, but Buck/source-db configs made a missing marker look origin-sensitive and could pay one negative lookup per file. Add a conservative `SourceDatabase::may_contain_module` hook and use it only for the `shape_extensions` marker: when the source DB can prove the marker is absent and the remaining fallback lookup is origin-independent, cache the miss under the origin-independent key. The optimization deliberately stays find-only. It does not add `shape_extensions` to normal module imports or content dependencies, and origin-sensitive configs still use the regular per-origin lookup path. Reviewed By: avikchaudhuri Differential Revision: D107602378 fbshipit-source-id: 3743820a8adf1ccf04dc7a53b381367878d88d86

Add solver residual-answer plumbing scaffolding Summary: **This stack** Adds a "callable residual" type that allows us to capture structure when a generic or overloaded callable is passed to a higher-order function. The residuals can later be "exploded" to recover generic and overload structure in outputs (either Callable outputs, or classes whose methods capture the callable structure of an argument - e.g. callback protocol classes). The full design details are [here](https://github.com/stroxler/pyrefly-docs/tree/main/callable-residuals) **This diff** Why: Residual-aware solving needs explicit solver-side slots for residual metadata and visibility-gated answers, but we want to introduce that structure without changing solving behavior yet. What: Add Bounds residual metadata storage, ResidualIdentity scaffolding, and Variable::ResidualAnswer, while keeping all solver flows that read solved vars treating ResidualAnswer like Answer for now. Why it works: The solver can now carry residual-oriented state through existing pathways without behavioral shifts, preparing follow-up commits to add capture, gating, and flattening semantics incrementally. Reviewed By: rchen152 Differential Revision: D102011885 fbshipit-source-id: a1b9f57a46f44d8357d81c9de50c34b215f3f00a

Add `SymVar` quantified kind for legacy generic syntax Summary: **Context** Sam, Avik, and I discussed and concluded that for a number of reasons - especially because a PEP would require runtime support for arithmetic ops, but there are also semantic reasons in the type theory - symbolic integer type variables should be a distinct "kind" from normal TypeVar. By this we mean "kind" in the sense that TypeVarTuple and ParamSpec are distinct kinds: it's a type variable with a different domain, legal in disjoint places from `TypeVar` with different syntactic rules defined such as allowing symbolic arithmetic... in other words, a new QuantifiedKind **This Diff** `shape_extensions.SymVar` is the legacy runtime-compatible spelling for symbolic shape variables, but Pyrefly still represented it as an ordinary `TypeVar`. Add `QuantifiedKind::SymVar` and preserve that kind through legacy `SymVar(...)` declarations while continuing to use the existing `TypeVar` validation and generic behavior for now. This keeps the behavior coherent for the legacy syntax: `SymVar` works anywhere the old shape-extension `TypeVar` worked, including `Dim[...]` and tuple-carrier shaped-array forms, while the kind is now distinct so later diffs can make normal type variables and symbolic shape variables legal in disjoint places. The underlying principle here is that symbolic integers are a type variable "kind", analogous to TypeVarTuple or ParamSpec in that they behave differently and follow different rules on syntax, semantics, and where they can legally appear than TypeVars do. This is something Sam, Avik, and I discussed in some depth. Reviewed By: yangdanny97 Differential Revision: D111135444 fbshipit-source-id: 4dc4806068bcc3de85786bf759a305d57e5090cf