Compile-Time Tensor Shape Checking via Staged Shape-Dependent Types

Abstract: When writing programs involving matrices or tensors in general, it is desirable to rule out the inconsistency of tensor shapes (i.e., the generalization of matrix sizes) before actual computation. For this purpose, some languages provide dependent types such as Mat m n, and others offer refinement types to track predicates for shapes. Despite the theoretical maturity, however, such methods are often unhandy for continuous software development due to the requirement of proofs for judging type equality or subtyping; even automated proving is often unsuitable due to its unforeseeable time consumption. To remedy this, our study provides an alternative formalization by using staging. Based on the observation that conditions for the shape consistency can be extracted before running the actual tensor computations in many typical cases, we ensure such consistency by assertions evaluated as compile-time computations, not by proofs. Under this formalization, we can verify the consistency virtually statically in the sense that inconsistencies as to tensor shapes will be immediately detected as assertion failures during compile-time computation. Our work achieves a mathematical guarantee that successfully generated code is always consistent with respect to tensor shapes. Furthermore, to vastly lessen the burden of adding shape- or stage-related descriptions, we (1) allow shape-related arguments to be implicit and infer them in a best-effort manner, and (2) offer a non-staged surface language that seemingly resembles ordinary dependently-typed languages and translate its programs into the staged core language. By a prototype implementation, we confirm that our language is expressive enough to verify a number of programs, including several examples offered by ocaml-torch.

• The paper introduces a compile-time stage that computes and checks tensor shapes, eliminating manual proofs and reducing false positives.
• It employs a two-stage architecture separating shape computations (stage-0) from tensor manipulations (stage-1) to guarantee run-time safety.
• Experiments with the Horsea prototype demonstrate rapid compile-time checks (<0.1s) and robust handling of advanced operations like broadcasting.

Compile-Time Tensor Shape Checking via Staged Shape-Dependent Types

Motivation and Context

The challenge of statically verifying tensor shape consistency is pervasive in modern numerical and machine learning applications. Standard approaches, such as full-spectrum dependent type systems (e.g., Idris) or refinement types (e.g., GraTen), provide expressive means for describing shape properties, but require either manual equational reasoning or engagement with unpredictable automated solvers to establish type equalities and subtyping relationships. In typical software engineering environments, where requirements are volatile and development cycles are rapid, these mechanisms are impractical: they introduce significant annotation overhead, fragile proof obligations, and unpredictable performance costs during type checking. Furthermore, existing approaches are either brittle in the presence of advanced operations (e.g., broadcasting) or generate excessive false positives.

Core Ideas: Staged Shape-Dependent Types

The paper introduces a compile-time approach to shape checking built on staged computation. The essential observation is that for a broad class of tensor-manipulating applications, shape computations and consistency checks are lightweight, side-effect-free, and can be completely resolved at compile time before any actual tensor-valued computation occurs.

The proposed system splits tensor programs into two stages:

1. Stage 0 (Compile-Time): All shape computations, constraints, and consistency assertions for tensor operations are evaluated. This includes evaluating numeric indices and shape expressions, performing shape arithmetic, and inserting and executing contract checks for shape equality or required broadcastability.
2. Stage 1 (Run-Time): Only tensors with fully statically consistent and specialized shapes, according to the results of stage 0, are manipulated. At this point, shape errors are impossible, and tensor code can be further optimized or compiled.

The mechanism relies on a staged, dependently-typed core language where shape terms are evaluated as first-class functional code at compile time. Type-directed contraction checks—such as type equivalence or refinement assertions—are realized as assert-like contracts that are evaluated during the compile-time phase, as opposed to requiring the generation or verification of explicit equational proofs.

Language Design

The language exposes a staged core, parameterized by an explicit distinction between stage-0 (compile-time) and stage-1 (run-time) expressions and types. Stage-1 tensor types (e.g., $\mathtt{Tensor}~\%~m$) are parameterized over stage-0 shape expressions, with the cross-stage persistence operator (%) denoting value transfer from compile to run time.

Shape consistency is ensured by inserting cast or contract assertions wherever type compatibility across applications must be validated (e.g., when passing a tensor of some computed shape as an argument). The system leverages an ML-style staged lambda calculus (MetaML-like), with brackets and escapes for quasiquotation and splicing.

Implicit Arguments and Surface Syntax

To reduce the annotation burden inherent in explicit dependent programming, shape-related parameters can be left implicit, subject to a type inference algorithm that reconstructs omitted information using local substitution and type context, inspired by the "let arguments go first" principle. A non-staged surface language (Horsea) resembling standard ML/Idris syntax is provided, and is translated to the staged core via binding-time analysis.

Refinement Types and Error Localization

Stage-0 refinement types (e.g., $\{\,n : \mathtt{Int}~|~n \geq 0\,\}$) allow precise specification of shape-related invariants and serve as contracts to localize the sources of shape inconsistencies during compile-time checks. Failure diagnostics can thus point directly to the call-site or argument culprit, even for computed or parameterized shapes.

Expressiveness

The staged approach accommodates complex tensor operations, including implicit conversion mechanisms such as NumPy/PyTorch-style broadcasting, by representing these as first-class compile-time functions that compute and verify broadcastability or generate shape-specialized code fragments.

Formalization and Metatheory

The system is formalized via an explicitly staged dependent type theory. Compile-time types are interpreted as sets of stage-0 expressions; run-time types are parameterized over compile-time values. Key features of the formalization:

• Elaboration rules insert the minimal necessary set of compile-time assertion contracts at function applications and data flows where type unification or subtype checking is required.
• Reduction semantics guarantee that all shape-related errors are manifest as failures or assertion violations during the compile phase, and that only shape-safe code is ever produced for execution.
• Preservation and Progress hold at both stages, with a cotermination property ensuring that once code is generated, run-time execution is free of shape failures.
• Unlifting is introduced as the mapping from generated code fragments (after compile-time evaluation) back into the simply-typed world for execution and further optimization.

Practical Implementation and Evaluation

A prototype implementation (Horsea) is provided, with a Haskell backend and a translation pipeline from surface syntax down through staged representations. The prototype verifies the shape correctness of a variety of nontrivial ocaml-torch example programs, notably automating the checking of advanced shape manipulations including reshaping and broadcasting, and validates the practicality of the surface syntax and inference scheme.

Performance characteristics: All compile-time shape checks in evaluated programs complete rapidly (typically $<$0.1s), as no heavy proof search is required—only ordinary functional evaluation of the relevant subterms.

Implications and Theoretical Impact

The staged shape-dependent type system achieves several properties relevant to both programming language theory and practical software engineering:

• By eliminating the need for proof artifacts (either manual or automated solver-based), the system seamlessly integrates into iterative and experimental software workflows, supporting rapid refactoring and evolution for applications where type indices (shapes) are ultimately statically known.
• False-positives endemic to overapproximate static analyses are eliminated, since all contracts are checked for concrete shapes corresponding to program instantiations.
• Support for programmable implicit conversions (broadcasting) is achieved without unduly complicating the type system.
• The decoupling of shape computations (stage-0) from tensor computations (stage-1) provides a clear phase distinction, opening additional avenues for compile-time optimization, code specialization, and further static or dynamic verification.

Limitations and Future Directions

The system does not seek to subsume full dependent type verification for universally quantified libraries, nor is it capable of statically proving shape invariants parameterized over pure type-level variables (e.g., generic tensor libraries). However, it is effective for verifying all end-user code instantiated for concrete computation—precisely the class of programs causing the most insidious run-time failures in practice.

Potential future directions include:

• Extending the approach to richer forms of type inference and cross-stage information flow.
• Integration of property-based testing at the stage-0 level for randomized shape space sampling.
• Exploring hybridization with verification-oriented systems for library authors, potentially via property contracts or symbolic solvers for a subset of term-level indices.

Conclusion

By leveraging a type-theoretic staged programming discipline, the paper provides a rigorous compile-time guarantee of tensor shape correctness in realistic programs, without recurring to the annotation and proof burdens of traditional full-spectrum dependent types or refinement-type solvers. The system readily accommodates practical tensor manipulations, scales naturally with advanced program evolution, and provides precise, actionable error localization for the end-user. The approach represents a significant recombination of staged metaprogramming and type-level contracts, yielding a pragmatic, mathematically justified solution to shape safety in tensor programming.

Reference: "Compile-Time Tensor Shape Checking via Staged Shape-Dependent Types" (2604.23807)