Avik Chaudhuri

Port 9 TorchBench models with tensor shape annotations Summary: 9 new model ports with full shape tracking: tacotron2, llama, stargan, sam, tts_angular, drq, dlrm, deeprecommender, background_matting. Key improvements over initial ports: - sam: PromptEncoder[D, ES, MIC] fully typed, window partition/unpartition typed, ImageAttention generic over H/W, unified D across components - llama: RoPE chain typed (arange/outer DSL), KV cache typed with Optional start_pos, build_causal_mask typed - tacotron2: variadic Prenet[*Bs], full decoder forward with autoregressive loop - stargan: Generator and Discriminator generic over spatial dim S - drq: SquashedNormal[*EventShape] with typed distributions - dlrm: tensor fancy indexing verified with assert_type 11 type: ignore remaining (5× A1 algebraic, rest genuine limitations). Stub additions: torch.quantization, optim updates, F.interpolate symint. Reviewed By: stroxler Differential Revision: D97893786 fbshipit-source-id: 863984bee7092fdff92b5df46de35cb771c6ac9e

Port 7 TorchBenchmark models with tensor shape annotations Summary: Ports Demucs, DenseNet, LearningToPaint, SoftActorCritic, SpeechTransformer, Super-SloMo, and U-Net from TorchBenchmark with full tensor shape annotations. Each model demonstrates end-to-end shape inference through complex architectures including pooling, upsampling, LSTM, GLU, PixelShuffle, and None indexing. Also adds fixture stubs for torch.distributions (used by SoftActorCritic) and additional torch/__init__.pyi and nn/functional.pyi function stubs needed by the ported models. Reviewed By: yangdanny97 Differential Revision: D97134636 fbshipit-source-id: 672256498c414706845c51ab5c500d519a2b7328

TypeVar variant for arithmetic Summary: Add torch_shapes.TypeVar and TypeVarTuple classes that support arithmetic operators (N + 1, N * 2, etc.) at Python runtime, unlike typing.TypeVar which raises TypeError. These use __class__ = typing.TypeVar so that isinstance checks pass and Generic[N] works correctly. On the pyrefly side, extend the TypeVar/TypeVarTuple special export recognition to accept torch_shapes as a valid defining module, so torch_shapes.TypeVar is treated identically to typing.TypeVar by the type checker. Also add int_type_var model variants (nanogpt, gptfast) that use torch_shapes.TypeVar instead of PEP 695 syntax, with both type-checked and runnable versions, plus comprehensive runtime and type-checking tests. Reviewed By: stroxler Differential Revision: D95282862 fbshipit-source-id: ec07dbb9efc94b653d6afc4ec9bf6b3d821580f0

port-model: add 7 typed port examples as eval baseline Summary: Adds seven Tensor[]-shaped ports under fbcode/pyrefly/test/tensor_shapes/models/: mobilenetv2, mixtral_moe, wavernn, openpose, apg, finalmlp, masknet. Each was selected as the highest-scoring round from the v7 multi-round eval batch (R1 blind → R2 self-reflection → optional R3) using the verify_port.sh metric (shaped-vs-bare ratio, receipt coverage, smoke tests). They serve two purposes: 1. Reference ports demonstrating the skill's intended output across architectures (vision, MoE transformer, RNN, pose estimation, recommendation MLPs). 2. A baseline corpus future skill iterations can be evaluated against without re-burning compute on the v7 multi-round flow. Also lifts three stub additions surfaced by the mobilenetv2 port into fixtures/torch/nn/__init__.pyi: Conv2d.bias, GroupNorm.weight, GroupNorm.bias. Without these, the init-loop pattern (`for m in self.modules(): if isinstance(m, nn.Conv2d): m.bias`) fails to typecheck. Reviewed By: stroxler Differential Revision: D102635653 fbshipit-source-id: 1508cc0bffd34726a22a3706c84e792890ad91e0

Rewrite model ports with generic blocks and recursive patterns Summary: Rewrite model ports to use class-level type parameters, ModuleList with recursive chain patterns, and overload for base/recursive cases: - bert: super() cleanup - dcgan: DCGAN config class, GenUpStage/DiscDownStage generic blocks, recursive _chain with overload and exponential return types - demucs: generic EncoderBlock/DecoderBlock, overload on center_trim, recursive encoder-decoder with skip connections - densenet: _dense_chain as free function with overload - gptfast: remove duplicate license, Optional cleanup - resnet: recursive _chain with overload, ShapePreservingActivation - soft_actor_critic: remove trailing whitespace - super_slomo: major rewrite with ModuleList + recursive patterns, separate encode/decode ModuleLists - unet: ModuleList recursive pattern with skip connections Also adds test_overload_narrowing.py and test_recursive_skip.py. Reviewed By: stroxler Differential Revision: D97425585 fbshipit-source-id: 2f5c0c7db4d68e42a027411308e3c543f78a9980

jaxtyping integration Summary: Add jaxtyping annotation support to tensor shape type inference. To ease adoption it makes sense to parse and print syntax like Float[Tensor, "batch channels"], which already exists out in the wild, as an alternative to (or mixed with) "native" syntax Tensor[N, M], which is experimental. Fortunately this extension is fairly lightweight. We reuse the existing internal generics machinery as well as all existing mechanisms for declaring / computing shapes of torch ops. Details: - New TensorSyntax enum (Native / Jaxtyping) to track annotation origin. Transparent to type equality — syntax doesn't affect type identity, only display rendering. - Added module-scoped map from name to Quantified: same dim name always maps to same Quantified TypeVar, so dims are shared across functions in a module. - Supported syntax: named dims, integer dims, variadic *name, anonymous _, ellipsis ..., broadcast # (stripped, treated as regular dim), combined *#, arithmetic (n+1, n-1), parenthesized arithmetic, scalar "", leading/trailing whitespace. We don't care about dtypes yet. Special case: jaxtyping's shape string, when specifying just a name, is treated as a string literal, not a forward reference to be resolved. - Extended function tparams computation: visits callable types to find jaxtyping-created Quantified dims, adds them to the function's Forall. Added mixed-syntax detection: errors if both native Tensor[N, M] and jaxtyping Float[Tensor, "N M"] appear in the same function. This takes care of most common use patterns of jaxtyping. Main gap: not sure how dims are shared across fields and methods of a class - so that's currently a strict advantage of native syntax, and in particular, without that we cannot port nanogpt.py and gptfast.py tests over to jaxtyping. Reviewed By: stroxler Differential Revision: D94942050 fbshipit-source-id: 717b5bf39f73a297073f6cb1785d826e4834f8b9

runtime impact of tensor type annotations Summary: This diff demonstrates some known shortcomings of native syntax for tensor type annotations with shapes: (1) torch.Tensor is not subscriptable (2) typing.TypeVar does not allow arithmetic. These kinds of errors happen early (on definitions) and have nothing to do with type checking tensor shapes per se. We also demonstrate that `from __future__ import annotations` is a workaround for both shortcomings. In particular, both errors do not manifest at run time when such types appear purely as, well, annotations. Moreover, type checking continues to work, as demonstrated by adding `from __future__ import annotations` to nanogpt and gptfast - both of these extensively use tensor shapes and guard against regressions, and the fact that these annotations are not "evaluated" at run time does not seem to affect the type checking. However, as expected and also demonstrated, evaluation of `assert_type` at run time triggers the same errors. What to do with this information will be the subject of future work, but for now, documenting these known shortcomings via tests makes sense. Reviewed By: stroxler Differential Revision: D95090840 fbshipit-source-id: ad276bea4118ffda35dce96f814bc82a551fd678

runnable model variants for runtime testing Summary: Create runnable copies of the future_annotations model variants with assert_type calls stripped, plus runtime tests that instantiate and run forward passes. This verifies that the typed models can actually execute at Python runtime. Reviewed By: stroxler Differential Revision: D95416793 fbshipit-source-id: 2169e231d9c111d79bc40dbc229296d41fb6eaaf

Tensor indexing and algebraic simplification improvements Summary: - Slice step support: IndexOp::Slice gains step field, apply_step computes ceil_div for strided slices (e.g. tensor[::2]) - None/NewAxis indexing: tensor[None] inserts a size-1 dimension - Symbolic negative bounds: adjust_negative handles -1*X from unary negation - Negative dim annotations: parse_dimension_expr handles UnaryOp::USub - Algebraic simplification: distributive law (c*(a+b) → c*a+c*b) and floor-div literal extraction ((a+k*d+b)//d → k+(a+b)//d) in dimension canonicalization - Dim <: float subtyping (numeric tower) - Type::Tensor in unary operator dispatch - Optional[Dim[T]] unwrapping in inject_module_attrs Reviewed By: yangdanny97 Differential Revision: D97134529 fbshipit-source-id: 9af8c2b239741ce616f37f82789e006f1d8dee88

Add shape-preserving module stubs and tensor op coverage Summary: Add comprehensive stubs for nn.Module subclasses: 17 activation modules, 9 normalization modules, 5 dropout modules, Identity, basic Conv1d/2d/3d and ConvTranspose1d/2d/3d (channel tracking only), 12 pooling modules, 10 loss modules, and basic Flatten/Unflatten. Also add tensor creation methods (new_zeros/ones/empty/full), randint, addmm, cross to torch stubs, and F.log_softmax/softmin/linear/embedding/rms_norm/dropout1d/2d/3d to functional stubs. Reviewed By: yangdanny97 Differential Revision: D96524714 fbshipit-source-id: 5e5edd3cbd444c9a3960a852a0a0aca4ef2e12ac

Port BERT from TorchBenchmark with shape annotations Summary: Port BERT model from TorchBenchmark with full shape annotations using Module-as-Callable for SublayerConnection and PEP 696 TypeVar defaults for TokenEmbedding and BERT hidden dimensions. Includes LayerNorm, MultiHeadedAttention, PositionwiseFeedForward, TransformerBlock, BERTEmbedding, and BERTLM with assert_type smoke tests. Reviewed By: yangdanny97 Differential Revision: D96524719 fbshipit-source-id: da70e42fff69a05bd90433208788dbe90f73fcc4

Port ResNet and DCGAN from TorchBenchmark with shape annotations Summary: Port ResNet (Phlippe variant) and DCGAN models from TorchBenchmark with full shape annotations. ResNet uses Sequential chaining, Conv formulas, Flatten, and ShapePreservingActivation union type. DCGAN uses Sequential with Conv2d/ConvTranspose2d directly. Includes model_port_changes.md documenting all changes from original sources. Reviewed By: yangdanny97 Differential Revision: D96524711 fbshipit-source-id: caaf599040db552bdb5f8bf84ec18e66db868834

Variadic DSL support and nn.Flatten shape inference Summary: Add variadic shape handling to the meta-shape DSL: Val::Unpacked variant for shapes with prefix/middle/suffix, variadic-aware Slice/Concat/TensorNew operations, and soft errors (ShapeError::Unsupported) for operations that require concrete lengths (Len/Enumerate/Zip/Index/ListComp). Register nn_flatten_forward_ir in the DSL registry and add inject_module_attrs to resolve module field values (start_dim, end_dim) for DSL evaluation. Upgrade nn.Flatten to class Flatten[S=1, E=-1] with typed fields. Reviewed By: stroxler Differential Revision: D96524716 fbshipit-source-id: 4061ca6ec1bee08d717e239d78115ffcbcc08a54

website docs for tensor shapes Summary: Collecting and organizing a bunch of material on tensor shapes into some public-facing docs. 1. Overview (previously in a google doc) 2. Setup 3. Tutorials (series, from basic to advanced: previously in a style guide) 4. AI porting (Claude skill) 5. API reference 6. Contributing to PyTorch stubs (fixtures / DSL definitions for library ops and modules) Reviewed By: stroxler Differential Revision: D102399686 fbshipit-source-id: cae1552151f2463482610d688ff401ff6d100d86

Type::NNModule: nn.Module instances with captured constructor args Summary: Introduces Type::NNModule, a new type variant that wraps ClassType + a field map of captured __init__ arguments. This enables DSL forward functions to access constructor parameters (e.g., kernel_size, stride) directly from the type, without requiring generic type params on every nn.Module stub. - NNModuleType struct with Hash/Ord/Visit/VisitMut/TypeEq impls - maybe_wrap_nn_module in construct_class: captures init args via registry - NNModule call dispatch delegates to underlying class __call__ - inject_module_attrs injects NNModule fields into DSL bound_args - Literal promotion skip for NNModule types (preserves shape info) - NNModule subtyping, attribute delegation, display, variance The register_init_forward API is added but no modules are registered yet — that comes in the next diff. Reviewed By: stroxler Differential Revision: D97134552 fbshipit-source-id: 52f00e5095f0bf6a0810c74db82c39b0b857ad6c

Add SizeExpr::Pow for exponential shape patterns Summary: Add exponentiation support to tensor dimension arithmetic: - SizeExpr::Pow(base, exp) variant in dimension.rs - canonicalize_pow with same-base grouping and literal absorption (e.g., 8 * 2**I → 2**(I+3)) - Display, contains_any, compare_type support - Parse ** operator in type annotations (expr.rs) - Variance inference for Pow (variance_inference.rs) - Extend is_dim_operand to recognize Literal[int] as dimension operand and convert SizeExpr::Literal back to Literal[int] when both operands are literal ints (operators.rs) - Test cases in test_canonicalization.py This enables exponential shape patterns like Tensor[B, C // 2**I, H * 2**I, W * 2**I] used in recursive chain architectures (DCGAN, ResNet, DenseNet). Reviewed By: stroxler Differential Revision: D97425570 fbshipit-source-id: afaf0c8e4cd4ade53d3fc7d6fdc575b56e34e158

Fix DSL crash on variadic tensor shapes and improve broadcast handling Summary: The view/reshape DSL crashes when a tensor has variadic *Bs dims (e.g. from Linear.forward) because prod() and sum() call as_list() on an Unpacked value. Fix by adding Unpacked guards that return Unsupported, matching the pattern already used by len(), enumerate(), and zip(). Broadcasting two tensors with different TypeVarTuples (e.g. *Ts vs *Us) produced a hard error even though in practice the batch dims are often identical. Fix by degrading to shapeless batch dims (any_tuple middle) instead of erroring, consistent with the existing unbounded-middle path. Also allow Tensor[*tuple[Any, ...], N] as a valid annotation for "unknown batch dims + known suffix", which is the natural way to annotate the degraded broadcast result. Reviewed By: kinto0 Differential Revision: D101657005 fbshipit-source-id: a640b38c33e9b2e818b34768a785637f8c739c14

Fix tuple Unpacked-vs-Unpacked subset check Summary: The bidirectional check in Tuple::Unpacked subset comparison was incorrect for TypeVarTuples. It wrapped both middles in tuples and checked both directions, causing tuple[*Cs] ⊆ tuple[*Bs] to fail (comparing *Cs ⊆ tuple[tuple[*Cs]] after binding). Fix: single-direction check based on which side has extras: - No extras: l_middle ⊆ u_middle directly - Only got extras: fold got side, bind want middle - Only want extras: fold want side, bind got middle - Both extras: error (cross-structural) This fixes variadic TypeVarTuple usage in tensor shape methods, e.g. forward[*Bs](self, x: Tensor[*Bs, D]) with multiple variadic-preserving function calls (F.relu, F.dropout). Reviewed By: stroxler Differential Revision: D97893749 fbshipit-source-id: 7a36527053deb70db58a22465fe83af7b5749d9e

Preserve tensor shape through arithmetic with Any/Unknown operands Summary: binop_types short-circuits to Any when either operand is Any, which is correct for generic types (the other side's __radd__ could return anything). But for shaped Tensors this is too conservative: Tensor's arithmetic dunders accept any numeric type and return Self, so the shape is always preserved regardless of the scalar operand's type. Without this fix, Tensor[B, 1] / (2**n - 1.0) loses its shape because 2**n returns Any (typeshed's int.__pow__ catch-all for non-literal exponents), which propagates through subtraction to Unknown, and then the division short-circuits to Unknown. With this fix, the Tensor's __truediv__(Unknown) dispatches normally and returns Self = Tensor[B, 1]. Reviewed By: kinto0 Differential Revision: D101700350 fbshipit-source-id: 229b3e994af447ba50da57b8e9a184bc492e08a0

Fix (a*b)//b → a simplification when b is an expression Summary: When the distributive law expands B*(2*A-1) into -B + 2*A*B, division by (2*A-1) couldn't cancel because the numerator is a sum, not a product. Add try_factor_sum_and_cancel in canonicalize_division: extract common non-literal factors from sum terms, rebuild as product, then cancel with denominator. This is the reverse of distribution — "un-distributing" the sum inside division context while keeping expanded form as the canonical representation. Reviewed By: samwgoldman Differential Revision: D97893688 fbshipit-source-id: b6518b4972ce142b2599cb29046c3a29b819610e