Design Space 01: Static vs Dynamic Shapes

Series: AI-Native Programming Language Research — Design Space Memos Document 1 of 6 (planned: shapes, AD-through-effects, sparsity, distribution, verification, hardware abstraction)

1. Problem statement

Every tensor operation has a shape — at minimum, rank and per-dimension size, but in practice also dtype, layout, device, and (increasingly) sharding. The fundamental design question is: when is the shape known, and who is responsible for proving it correct?

The two ends of the spectrum:

• Fully static: shapes are part of the type, the compiler proves shape compatibility, and a shape mismatch is a compile error. Ergonomic cost: annotations and shape arithmetic show up in user code. Compilation cost: a new input shape may force recompilation.
• Fully dynamic: shapes are runtime metadata on values, every operation checks them, mismatches are runtime errors. Ergonomic cost: zero up front, but errors surface deep in execution traces. Optimization cost: the compiler can't fuse, plan memory, or pick layouts without speculating.

Neither pure end has won, and the interesting work in the last five years has been at the symbolic middle: shapes are values that may be unknown but are subject to compile-time constraints and arithmetic, and the compiler can specialize, guard, or fall back depending on what it learns.

This memo maps the spectrum, surveys the points along it, compares them, and identifies open problems.

2. What "shape" actually means

A useful distinction before going further. "Shape" in the loose sense is everything the compiler needs to know about a tensor to generate code:

The static-vs-dynamic question applies to each component independently. Most languages are static in dtype and rank but disagree on dimension sizes and sharding. When people say "static shapes" they almost always mean dimension sizes known at compile time.

3. The spectrum

Five distinguishable points, ordered from most dynamic to most static:

3.1 Fully dynamic (NumPy, eager PyTorch, eager TensorFlow)

Shapes are runtime metadata. Every op does a shape check; broadcasting rules are evaluated at call time. There is no static guarantee that A @ B is well-formed.

Strengths: zero annotation burden, trivially supports data-dependent shapes (variable-length sequences, ragged batches, dynamic masking).

Weaknesses: errors surface late and far from cause; compilation can't fuse without speculative tracing; memory planning is reactive (alloc on demand) rather than precomputed; no IDE assistance for shape errors.

3.2 Tracing-based (JAX jit, TensorFlow AutoGraph, PyTorch's torch.jit.trace)

The user writes dynamic-looking code. The system traces an execution with abstract values, recovers a graph specialized to the input shape, and compiles it. Subsequent calls with the same shape hit the cache; new shapes trigger retracing.

Strengths: keeps the dynamic feel; the compiler sees a fully shaped graph and can apply XLA-class optimizations.

Weaknesses: the recompilation cliff. Production LLM serving with variable sequence lengths or batch sizes hits this constantly. Tracing also can't see Python-level control flow that depends on runtime values; users write jax.lax.cond and jax.lax.scan to work around it. Side effects during tracing have notoriously subtle semantics.

3.3 Symbolic shapes (PyTorch 2 / TorchDynamo + Inductor, TVM Relax, MLIR shape dialect, ONNX dynamic axes)

Shapes are first-class symbolic values at compile time. A dimension might be s0 (unknown but consistent across uses), s0 + 1 (arithmetic), or s0 * s1 (products). The compiler reasons about constraints, generates guards, and emits one kernel that handles a family of shapes.

This is where most current research and engineering effort sits, and it's the most pragmatic answer for LLM workloads where sequence lengths and batch sizes are genuinely dynamic.

Strengths: one compilation handles a family of shapes; integrates with autotuning; supports data-dependent shapes via guards and graph breaks.

Weaknesses: symbolic shape inference is hard — equality of s0 * 2 and s0 + s0 requires arithmetic reasoning. Guards are easy to over-specialize, leading to recompilations. Implementations are heuristic-heavy; the design space (which guards to insert, when to specialize, when to fall back) is unsettled.

3.4 Statically annotated, non-dependent (Hasktorch type-level shapes, tch-rs typed tensors, NamedTensor proposals)

Shapes go in types, but as type-level naturals or constants — not as values that can be computed over. You write Tensor [Batch, 768] Float, and the type checker matches them syntactically.

Strengths: catches shape errors at compile time without dependent types; works in mainstream type systems (Haskell, Rust, Scala) with type-level naturals.

Weaknesses: shape polymorphism is painful. Writing a function that works for any sequence length means parameterizing over type-level naturals, and the resulting signatures get heavy. Arithmetic on type-level naturals (e.g., proving M+N = N+M) often requires hand-written equality witnesses.

3.5 Fully dependent (Dex, Futhark, ATS, Idris, F*)

Shapes are ordinary values that happen to appear in types. dot : (n : Nat) -> Vec n F -> Vec n F -> F is a real type. The compiler proves shape compatibility through ordinary type checking.

Strengths: the language has nothing special to say about tensors — they're just instances of a more general dependent type machinery. Shape arithmetic is just arithmetic. Highest level of static guarantee.

Weaknesses: dependent type systems are hard. Type inference is undecidable in general; user-facing error messages are notoriously bad. Existing dependent languages (Idris, Agda, Lean) have research-prototype-level tooling. Dex made this work for arrays specifically with index sets, but the surface syntax has a learning curve.

4. Comparison

The diagonal here matters: tracing and symbolic shapes look superficially similar to dynamic shapes from the user's perspective, but the compiler infrastructure and consequences are completely different.

5. Tradeoff axes in depth

5.1 The recompilation cliff

This is the single biggest practical problem with pure static shape systems. JAX shipped jax.jit in 2018; production users hit recompilation on every new batch size or sequence length. The mitigations have evolved:

1. Bucketing: pad inputs to a small set of shapes. Wastes compute, complicates batching.
2. Polymorphic export (jax2tf, JAX shape polymorphism): export a graph with symbolic dimensions that lowers to an XLA program with dynamic shapes.
3. Symbolic shapes from the start (PyTorch 2): never compile a fully static graph; always parameterize over symbolic dims.

The pattern: what started as a static-shape system is, under load, growing a symbolic-shape layer. This is empirical evidence that pure static shapes are too rigid for production ML, especially LLM serving.

5.2 Shape polymorphism vs ergonomics

The deepest design tension. The function matmul should work for any compatible shapes, which means its type must be polymorphic in shape. Naive dependent typing makes this:

matmul : (m n k : Nat) -> Tensor [m, k] -> Tensor [k, n] -> Tensor [m, n]

That's tolerable. But chain three matmuls, do a softmax, an attention computation, and the type signature dwarfs the body. Real research questions:

• Can shape inference work for tensors as well as Hindley-Milner does for ordinary types?
• Can common shape patterns (broadcasting, reductions, reshapes) be encoded without per-call annotation?
• Should shape variables be implicit (inferred) by default, with explicit annotation as opt-in?

Dex's index-set approach handles this for pointful programs (where you write loops over indices) but not for pointfree programs (where you compose tensor functions). Mojo punts on this entirely so far.

5.3 Data-dependent shapes

A tensor whose shape depends on a runtime value. Examples:

• nonzero(x) — output shape depends on x's contents.
• x[mask] — boolean indexing.
• Variable-length sequences in attention.
• Top-k routing in MoE.
• KV-cache that grows during decoding.

Pure static-shape systems can't express these without falling back to fully dynamic. Symbolic shapes handle them via guards: "compile assuming s0 ≤ 1024, generate a fallback for the violation case." Dependent types handle them via existentials: Σ n. Tensor [n].

The hard cases are control flow depending on shape (e.g., MoE routing where each expert sees a different number of tokens). These are still bolted-on in essentially every system.

5.4 Sharding and shape

Distributed tensors blew up the shape question in the last three years. A [batch, seq, hidden] tensor sharded across a [data, model] mesh has different effective shapes per-device, and the compiler needs to reason about both the logical shape and the physical sharding.

GSPMD, JAX pjit/shard_map, and PyTorch DTensor all handle this with varying levels of compiler integration. The unsolved question: should sharding be in the type? GSPMD treats it as compile-time annotation that the compiler propagates. A more native answer would put Sharded(axis=0, mesh=M) in the tensor's type and let unification handle resharding. Nobody has done this cleanly.

5.5 Memory planning

Static shapes enable upfront memory planning: the compiler computes peak memory, picks an allocation strategy, and emits explicit alloc/free. Dynamic shapes force reactive allocation (caching allocator in PyTorch).

Symbolic shapes can plan if bounds are known: "allocate enough for the maximum sequence length." The cost is overallocation when actual shapes are smaller. Production systems (vLLM's PagedAttention, FlashAttention's tiling) work around this with custom memory managers — another thing that wants to be a language feature.

5.6 LLM-era workloads as a stress test

Modern transformer inference is the worst case for pure static shapes:

• Variable batch sizes (continuous batching).
• Variable sequence lengths.
• Growing KV-caches.
• Speculative decoding with variable acceptance.
• MoE with per-expert variable load.

A language designed today should treat these as first-class rather than as workarounds. None currently do.

6. Open research problems

1. Gradual shape typing. Mix static and dynamic shapes within a single program with a clean interop story. Inspired by gradual typing (Siek, Taha). Dimensions you annotate get checked statically; unannotated dimensions are dynamic. Can the compiler insert guards at the boundary? Can shape inference fill in the gaps?

2. Symbolic shape inference quality. PyTorch 2's symbolic shape engine is heuristic-heavy and not principled. Could you frame it as constraint solving over Presburger arithmetic? Polyhedral analysis (used in HPC compilers) handles this for affine cases — does it generalize?

3. Sharding-shape unification. Treat sharding as part of the shape, with resharding inferred by the compiler when shapes meet at an operation boundary. Looks like a dimension-aware unifier with cost model.

4. Differentiating through dynamic shapes. Reverse-mode AD over a graph with data-dependent shapes is genuinely awkward. What's the right semantics for nonzero, masking, MoE routing under autodiff? Some of this is folklore in the AD literature (Pearlmutter & Siskind, "Reverse-mode AD in a functional framework"); some is genuinely open.

5. Effect-aware shape semantics. When shape depends on a random sample (rejection sampling, top-k) or a side effect, what does "static" even mean? This connects to the AD-through-effects memo (next document).

6. Ergonomic dependent shape types. The Dex/Futhark approach is too heavy for mainstream adoption. The Hasktorch approach is too rigid. Is there a sweet spot — bidirectional type checking, refinement-style inference, effect-track shape constraints — that's expressive and ergonomic?

7. Shape-aware autotuning. When the shape is symbolic, kernel selection becomes a function of shape, not a fixed choice. This is partially solved by Triton-style autotuning but not in a principled way.

7. Recommendation for an AI-native language design

If I were starting from scratch and had to pick a position:

• Symbolic shapes as the default, not static or dynamic. This is where the empirical evidence points.
• Optional dependent annotations for the kernel-author tier, where shape proofs matter for correctness (FlashAttention-class code).
• Gradual typing at the boundary between annotated and unannotated regions, with guard insertion.
• Sharding in the type system, integrated with shape, not separate.
• First-class data-dependent shapes with existentials and refinement, not bolted on.
• MLIR shape dialect or equivalent as the IR-level representation, so the surface language and the compiler share a vocabulary.

The thesis: pure static is too rigid for production, pure dynamic is too unconstrained for the compiler, and the symbolic middle is where the real design work lives. Any language that picks an extreme will end up reinventing the middle anyway, so design for it from day one.

8. References

Foundational papers

• Ragan-Kelley, J. et al. (2013). Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. PLDI 2013. — Schedule/algorithm separation; influential for thinking of shape as a compile-time concern separate from semantics.

• Henriksen, T. et al. (2017). Futhark: purely functional GPU-programming with nested parallelism and in-place array updates. PLDI 2017. — Size-typed array language; one of the cleanest takes on static shapes.

• Vasilache, N. et al. (2018). Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions. — Polyhedral model for tensor programs.

• Chen, T. et al. (2018). TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI 2018.

• Lattner, C. et al. (2021). MLIR: Scaling Compiler Infrastructure for Domain Specific Computation. CGO 2021. — The IR substrate where most modern shape work lives.

Shape-typed array languages

• Paszke, A. et al. (2021). Getting to the Point: Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programming (Dex). ICFP 2021. — Index sets as a more usable foundation for shape-typed array programming.

• Rink, N. et al. Tensors Fitting Perfectly. — Type-level shape encoding in mainstream type systems.

• Hasktorch project — type-level naturals for tensor shapes in Haskell.

Symbolic shapes and modern compilers

• Lazos, A., Ansel, J. et al. (2024). PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. ASPLOS 2024. — TorchDynamo, symbolic shape inference, and the path from Python to compiled graphs.

• Lai, R. et al. (2023). Relax: Composable Abstractions for End-to-End Dynamic Machine Learning. — TVM Unity's first-class symbolic shape IR.

• Bradbury, J. et al. JAX: composable transformations of Python+NumPy programs. — Tracing-based approach with shape_polymorphism extensions for export.

Sharding and distributed tensors

• Xu, Y. et al. (2021). GSPMD: General and Scalable Parallelization for ML Computation Graphs. — Sharding annotations propagated by the compiler.

• JAX pjit / shard_map documentation. PyTorch DTensor RFCs.

Dependent types and verification

• Pearlmutter, B. & Siskind, J. (2008). Reverse-mode AD in a functional framework: Lambda the Ultimate Backpropagator. TOPLAS. — Foundational for thinking about AD over higher-order programs.

• Siek, J. & Taha, W. (2006). Gradual Typing for Functional Languages. Scheme Workshop. — The gradual typing reference; relevant for the gradual-shape-typing direction.

Related work in array programming

• APL family (J, K, Q, BQN) — dynamic but shape-aware; informative on the ergonomic side.

• Single-Assignment C (SaC) — shape-polymorphic functional array language predating Futhark.

Production / engineering writeups (less formal but important)

• vLLM PagedAttention paper (Kwon et al., 2023) — memory planning under dynamic shapes.
• FlashAttention papers (Dao et al., 2022/2023) — kernel-level shape and memory tradeoffs.
• ONNX dynamic axes specification — production interchange format's compromise position.

9. Next memo

Design Space 02: Differentiation through Effects. Reverse-mode AD over pure functions is well-understood; AD over programs with mutation, randomness, exceptions, or I/O is genuinely open. This connects directly to several of the open problems above (data-dependent shapes, MoE routing, probabilistic programming) and is where effect systems intersect with autodiff in non-trivial ways.