Steven Troxler

Add plumbing for literal-capturing Flag-bound type vars Summary: Pyrefly needs literal-preserving `Flag` type parameters, but adding a new restriction kind creates broad exhaustive-match fanout before source syntax should construct it. Model `Flag` as a passive `Restriction` with its small builtin domain beside `Restriction`, and mechanically teach display, simplification, solving, queries, and reporting to preserve or project that restriction. The end goal here is to allow us to pass flags to the type-level shape DSL with code that looks like this: ``` def f[B: Flag[bool]](b: B) -> Tensor[takes_a_bool(B)] ``` The stubs use this pattern extensively to avoid the need for a combinatorial number of overloads for all kinds of transforms that take flag values. This commit intentionally does not recognize `Flag` annotations or change inference; the following activation commit provides the single source construction point. Splitting passive representation from activation keeps the mechanical type-system fanout reviewable and behavior-neutral. Reviewed By: rchen152 Differential Revision: D114061238 fbshipit-source-id: c9a3acbff0ef9083a75bcadeb57fcaac1c8c1955

Handle overflow in dimension canonicalization Summary: Dimension canonicalization combines coefficients, distributes products, flattens division, and evaluates literal powers. Those transformations used unchecked signed arithmetic, so boundary expressions could panic instead of producing a usable gradual dimension. Keep the public `canonicalize` API unchanged while making its internal normalization fallible. Checked overflow now propagates through the complete normalization before the public boundary converts it to the gradual dimension, preventing partially canonical results. Safe boundary expressions remain exact, including unrelated maximum literals and cancellation involving `i64::MIN`. Reviewed By: rchen152 Differential Revision: D113565991 fbshipit-source-id: 4500e41ef4171bb54ab3f9b6239d7e9af38de6d1

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

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

Represent symbolic MapIntTuples parameters Summary: Concrete `MapIntTuples` sources reduce eagerly, but symbolic sources must remain until specialization. Preserve those maps in the type representation. At the root of a parameter annotation, retain the mapped member type and derive an ordinary `Sequence` view so calls can be checked before the source shapes are inferred. For example, `make[S: IntTuples](s: S) -> MapIntTuples[lambda T: Box[T], S]` maps forward from `S`. In `consume[S: IntTuples](xs: MapIntTuples[lambda T: Box[T], S]) -> S`, the map becomes a parameter pattern for call-site inference. Deferred lambdas support binder-aware substitution, traversal, and equality. Substitution always visits ordinary call arguments before the MapIntTuples-specific fields. Evaluation applies the same source rules recursively inside unions: tuple members are mapped, gradual-like members use the map fallback, `Never` is preserved, and other members are rejected. The representation and solver integration remain isolated in the dedicated `map_int_tuples` modules. Reviewed By: avikchaudhuri Differential Revision: D117994171 fbshipit-source-id: 0a5c113d44fa15ae297578f066766dedacbc4832

Add an Index type parameter restriction Summary: Add the shape-extension-specific `Index` upper bound for Python index values in ordinary annotations. Values remain normal Pyrefly types; the restriction validates the supported index grammar while preserving literals needed for shape computation. The restriction accepts integer, slice, `None`, ellipsis, and integer-sequence components, including gradual combinations. Errors that depend on applying an otherwise well-formed index—such as repeated ellipses or too many axes—remain the responsibility of the indexing operation, so overload fallback cannot hide them. Keep construction, validation, and lowering inside shape-specific modules. Generic type-system code sees only the shared shape-extension restriction boundary introduced by the preceding refactor. Reviewed By: avikchaudhuri Differential Revision: D118397695 fbshipit-source-id: fbbf010cfd4961397357c5e8d6ef9c6f54265fa8

Activate single-source literal-preserving `Flag` parameters Summary: Shape DSL functions need literal `int`, `bool`, and `str` configuration values without changing ordinary generic inference. Activate canonical `shape_extensions.Flag[...]` bounds as a shape-owned PEP 695 restriction with exactly one direct, single-value function parameter as the authoritative binding source. The binding layer records canonical import provenance in an explicit ordinary-versus-shape bound enum, while the isolated shape-Flag solver module owns domain, scope, source, and default validation. Calls precompute Flag variables once from the final type-parameter instantiation and carry a nonempty set in `CallContext`, avoiding both allocation for ordinary calls and variable-table lookups in argument/default hot paths. Direct sources retain literals; wrapped occurrences remain compatibility bounds; runtime defaults outrank type-parameter defaults and return hints. Ordinary user classes named `Flag` remain ordinary typing, including shadows inside the `shape_extensions` module. Indirect aliases, reexports, and unsupported legacy `TypeVar` bounds are diagnosed after normal untyping. Query, constructor, subset, CinderX, hover, and Pysa projections use the exact builtin domain, with Flag-specific diagnostics for invalid values and conflicting constraints. Reviewed By: rchen152 Differential Revision: D113199434 fbshipit-source-id: feae0c87f312d66b49609a6cca72a0a385916def

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

Bootstrap JAX shape stubs Summary: We want to work with the JAX core team on shape typing, and that conversation needs something concrete to react to: stubs they can read, and tests that show what the stubs actually guarantee. Nothing existed for JAX yet. This adds `pyrefly-jax-stubs`, a deliberately small model of JAX built on the harness from the previous commit. It follows the NumPy package rather than the Torch one, because the NumPy layout is the one we want to keep: one suite per test file, discovered by glob, with every file both type checked and executed against the real library so that a stub which is self-consistent but wrong still fails. Registering the package is a one-line addition to `run_tests.py`, so it is covered by the same single CI job as torch and numpy. The stubs cover array creation, broadcasting arithmetic, `matmul`, `reshape`, `transpose`, the axis/keepdims reductions, and the elementwise `jax.nn` activations. `transpose` and `.T` are exact at any rank, including an explicit axis permutation. Shapes are modeled; dtypes are not, matching the Torch stubs. `Array` lives in `jax/_array.pyi` rather than `jax/__init__.pyi` so that `jax.numpy` can refer to it without importing its own parent package, which is also how real JAX arranges it. The type-level DSL is still being built out, so some rules cannot be expressed yet. Where that happens the stub returns a gradual shape and carries a TODO, rather than guessing or rejecting valid code: - `reshape` is exact for an explicit shape. A `-1` placeholder is gradual, because inferring it needs the product of the dimensions, which the DSL does not expose yet; for the same reason a reshape that changes the element count is caught only at runtime. The in-flight Torch migration adds a `dsl.prod` intrinsic that makes both checkable. - `matmul` models 2-D operands, and batched calls are gradual rather than errors. Batched matmul is everywhere in JAX, so rejecting it would be worse than not checking it. The `@` operator stays stricter, because giving it a gradual fallback would also swallow the mismatched-inner-dimension error, which is the most valuable check here. One DSL workaround is called out in the package README because it is not guessable: an `int | tuple[int, ...]` parameter cannot be iterated after narrowing with `is_int_value` alone -- the function silently evaluates to `Unknown` with no diagnostic. Leading with an `is None` check makes the narrowing work, so `reshape_shape` accepts a `None` it then rejects, the same workaround `conv_shape` uses in the Torch stubs. Reviewed By: rchen152 Differential Revision: D117373822 fbshipit-source-id: b958a7eae793a60b204508f440f3dc09867d35fe