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

When it comes to writing programs involving tensors, it is desirable to rule out the inconsistency of tensor shapes (i.e., the generalization of vector lengths or matrix sizes) before running actual computations. For example, consider the function $f$ that takes three matrices $A$, $B$, and $C$ and computes $A(B++C++B)$, where $++$ is a binary operator that vertically concatenates two matrices. For verifying consistency, the following constraints must be met: (1) for the use of $++$, $B$ and $C$ must have the same number of columns; and (2) for the matrix multiplication, the number of $A$’s columns must be equal to the number of $(B++C++B)$’s rows. Checking such constraints during program execution is problematic because it may cause a runtime error due to some trifling shape mismatch bugs only after heavy computations, consuming much time.

A popular approach to addressing this kind of problem is to use dependent type systems. Some languages, such as Idris [6, 7], provide dependent types like $\mathrm{𝚅𝚎𝚌}n$ or $\mathrm{𝙼𝚊𝚝}mn$, which are the type for vectors of length $n$ and the one for matrices of size $m\times n$, respectively, and assign the matrix multiplication the following type: $\forall p,q,r:\mathrm{𝙽𝚊𝚝}.\mathrm{𝙼𝚊𝚝}pq\to \mathrm{𝙼𝚊𝚝}qr\to \mathrm{𝙼𝚊𝚝}pr$. Other methods, such as GraTen [32], offer refinement types like $\{\nu :\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\mid \nu .\mathrm{𝚜𝚑𝚊𝚙𝚎}=[m,n]\}$ or $\{\nu :\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\mid \mathrm{𝚕𝚎𝚗}(\nu .\mathrm{𝚜𝚑𝚊𝚙𝚎})=3\}$. As an example of this kind of approach, consider the function $f$ above implemented in an Idris-like hypothetical language:

Here, $\{j:\mathrm{𝙽𝚊𝚝}\}\mathrm{\cdots }\{n:\mathrm{𝙽𝚊𝚝}\}$ are binders for implicit parameters, and $\mathrm{𝑚𝑎𝑡𝑀𝑢𝑙𝑡}$ and $\mathrm{𝑣𝑒𝑟𝑡𝐶𝑎𝑡}$ stand for the matrix multiplication and $++$, which are assigned types $\forall p,q,r:\mathrm{𝙽𝚊𝚝}.\mathrm{𝙼𝚊𝚝}pq\to \mathrm{𝙼𝚊𝚝}qr\to \mathrm{𝙼𝚊𝚝}pr$ and $\forall p,q,r:\mathrm{𝙽𝚊𝚝}.\mathrm{𝙼𝚊𝚝}pr\to \mathrm{𝙼𝚊𝚝}qr\to \mathrm{𝙼𝚊𝚝}(p+q)r$, respectively. $\mathrm{𝐻𝑒𝑖𝑔ℎ𝑡𝑠𝑀𝑎𝑡𝑐ℎ}km$ is a manually proved lemma used for matching $D$’s actual type with the one required of $D$; since $D$ is of type $\mathrm{𝙼𝚊𝚝}(m+k+m)n$ and the function $(\mathrm{𝑚𝑎𝑡𝑀𝑢𝑙𝑡}A)$ requires its argument to be of $\mathrm{𝙼𝚊𝚝}(k+2\ast m)r^{\prime }$ for some $r^{\prime }$, the type-checker must verify that $m+k+m=k+2\ast m$ holds. Working as a proof, the lemma rewrites the former type to the latter through the $\mathrm{𝐫𝐞𝐰𝐫𝐢𝐭𝐞}$-construct. In general, to type-check an application $M_1{M}_2$, where $M_1$ and $M_2$ are known to have types $\mathrm{𝙼𝚊𝚝}{M}_{11}^{\prime }{M}_{12}^{\prime }\to T$ and $\mathrm{𝙼𝚊𝚝}{M}_{21}^{\prime }{M}_{22}^{\prime }$, respectively, we must prove the type equality $\mathrm{𝙼𝚊𝚝}{M}_{11}^{\prime }{M}_{12}^{\prime }=\mathrm{𝙼𝚊𝚝}{M}_{21}^{\prime }{M}_{22}^{\prime }$, i.e., that, for $i\in \{1,2\}$, argument expressions $M_{1i}^{\prime }$ and $M_{2i}^{\prime }$ always describe the same value in the given context. Languages of such approaches provide some form of mechanism for proving this equality.

Despite their theoretic maturity and success in many safety-critical fields, however, such verification methods do not seem to be so eagerly applied to relatively typical, continuous software development, especially cases in industry. Although various reasons can be considered for this, we suppose that the following situations would be major factors:

1. 1.

   The externality of requirements: Requirements imposed on software continuously arise due to social situations and users’ preferences, which are not predictable beforehand.

2. 2.

   The sequential nature of software development: Checking is performed repeatedly during development. Although we can restrict properties to check to some lightweight ones for frequently performed verification, it would still be better if we do not have to take much time for each run, in order not to harm productivity.

3. 3.

   Time is literally money: It is costly to let software engineers engaged in development. Moreover, each engineer is available for approximately only 40 hours per week.

Under such circumstances, major issues that hamper adoption of existing methods for tensor shape checking would be some of the following:

1. A.

   Cumbersomeness of future changes: Methods that require manual proving easily make future changes of programs unwieldy; even slight modification of programs may demand nearly complete amendment of proofs. This does not go along with the condition 1 above, i.e., the unpredictable nature of requirements imposed on software.

2. B.

   Unpredictable time consumption by automated proving: To reduce the burden of manual proving, some methods provide automated proving by using back-end solvers, possibly with some part of the syntax restricted to a subset suitable for automation. However, such approaches are often too time-consuming or at least take unpredictable time; slight change of properties to check may drastically increase the elapsed time.

3. C.

   Frequent false-positive errors: As pointed out in some articles like Ascari et al. [3] or a CACM article about the use of static analyses in Facebook (currently known as Meta) [19], methods for verifying detailed properties are likely to cause false positives too easily due to their nature of the overapproximation of program behaviors. Namely, even when the validity is clear for humans, analysis tools often warn, e.g., the existence of a type-level gap and require some form of annotations or proofs. Frequent false-positive errors also add additional cumbersomeness to future changes¹¹1“Theoreticians,” including the authors, tend to take false-positive errors for granted (since we cannot achieve sound and complete verification and thereby overapproximations of some kind are necessarily introduced). However, for an affinity with development workflows, it would even be worth considering to strike a balance between soundness properties and the frequency of false positives. .

4. D.

   Lack of concise support for flexible tensor-handling operations: There are some implicit conversions of tensors frequently utilized in DNN-related programs, such as broadcasting [18, 13]. These conversions are flexible enough to make reasoning tensor shapes non-trivial, at least beyond decidable theories. For example, two tensors of shapes [5, 3, 1, 10] and [3, 4, 10] are addable by broadcasting (specifically, by duplicating the former by $4$ along with the third dimension and the latter by $5$, respectively). Typical type systems have difficulty in supporting such conversions in a concise manner; they will require proofs for the feasibility of the conversion, and whether the proving is manual or automated, that will also lead to some of the issues A–C above.

In essence, for the adoption to continuous development, it would also be crucial to wipe off the concern about the burden arising from the mismatch between verification methods and development workflows, such as the one due to too many false positives or unpredictable time consumption, not only to establish a method to verify the correctness of programs.

Basically, $f^{\prime }$ takes $j$, $k$, $m$, and $n$ as stage-$0$ parameters, and produces code for $(A,B,C)\mapsto A(B++C++B)$ specialized for those parameters. Here, $\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}$ takes three parameters $p$, $q$, and $r$ and returns $⟨{\mathrm{𝑣𝑒𝑟𝑡𝐶𝑎𝑡}}_{p,q,r}⟩$, where ${\mathrm{𝑣𝑒𝑟𝑡𝐶𝑎𝑡}}_{p,q,r}$ is the specialized operation that vertically concatenates $p\times r$ and $q\times r$ matrices. The stage-$0$ built-in function $\mathrm{𝑔𝑒𝑛𝑀𝑎𝑡𝑀𝑢𝑙𝑡}$ also takes three parameters $p$, $q$, and $r$ and returns $⟨{\mathrm{𝑚𝑎𝑡𝑀𝑢𝑙𝑡}}_{p,q,r}⟩$, where ${\mathrm{𝑚𝑎𝑡𝑀𝑢𝑙𝑡}}_{p,q,r}$ is the matrix multiplication operation specialized for $p\times q$ and $q\times r$ (if $p$, $q$, and $r$ are all non-negative; an assertion failure will be raised otherwise).

To reflect the operational behavior, code-generating functions are assigned types peculiar to staged computation. For example, $\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}$ is assigned type $(p:\mathrm{𝙸𝚗𝚝})\to (q:\mathrm{𝙸𝚗𝚝})\to (r:\mathrm{𝙸𝚗𝚝})\to ⟨\mathrm{𝙼𝚊𝚝}\%p\%r\to \mathrm{𝙼𝚊𝚝}\%q\%r\to \mathrm{𝙼𝚊𝚝}\%(p+q)\%r⟩$ by combining dependent function types $(x:T_1^{\left(0\right)})\to T_2^{\left(0\right)}$ and code types $⟨T^{\left(1\right)}⟩$, meaning that $p$, $q$, and $r$ are available at stage $0$ and the resulting code for matrix computation is at stage $1$. By the same token, $\mathrm{𝑔𝑒𝑛𝑀𝑎𝑡𝑀𝑢𝑙𝑡}$ has type $(p:\mathrm{𝙸𝚗𝚝})\to (q:\mathrm{𝙸𝚗𝚝})\to (r:\mathrm{𝙸𝚗𝚝})\to ⟨\mathrm{𝙼𝚊𝚝}\%p\%q\to \mathrm{𝙼𝚊𝚝}\%q\%r\to \mathrm{𝙼𝚊𝚝}\%p\%r⟩$.

Although the type of matrices is indexed by size information as before, the type system is not equipped with non-trivial type equality to identify. Instead, the language offers stage-$0$ casts of the form $⦇T_1⊲T_2⦈$ and inserts them through type-checking in order for the casts to be evaluated at compile time to assert the equality of two types. Basically, the above program will be elaborated to the following in a type-guided manner, for example:

In this way, we can verify the consistency during code-generating compile-time computation. We expect that, in many cases, compile-time computation does not take much time and that inconsistencies will be immediately detected as assertion failures. If the stage-$0$ evaluation succeeds, i.e., produces a code fragment without causing any failure, the resulting code fragment is guaranteed to be consistent with respect to specific tensor shapes. In this sense, we can do tensor shape checking without relying on either manual or automated proving, and thus mitigate the aforementioned issues A, B, and C (We defer to Section 2.2 how the remaining issue D can be resolved in our method). One may see this approach as a compile-time version of manifest contracts [21, 30, 29, 58], hybrid type checking [22, 44], or some similar methods like Lemay et al. [45]. Our approach might also be able to be seen as some kind of staging-based foundation of the template metaprogramming in C++ [36].

Nonetheless, now we have to write many “annotation-like” arguments in exchange for the absence of proofs (like the three arguments in $\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}(m+k)mn$). Manually adding staging constructs may also be cumbersome. In Section 2, we illustrate that most of these descriptions can actually be cleared away by introducing surface languages.

We suppose that our method alleviates the issues A–D from the following perspectives:

1. A.

   Since it only requires type annotations for binders and some implicit arguments that cannot be inferred, it would not hamper future changes to a large extent.

2. B.

   Since compile-time computations can be described in a usual functional language, one can expect that users can easily estimate the elapsed time for shape checking. Furthermore, for typical cases, the elapsed time will be quite instant (say, less than 0.1 second).

3. C.

   Our method does not cause false positive errors about tensor shapes in the sense that equations of tensor shapes are tested by concrete values after specializing them.

4. D.

   As explained in Section 2.2, our method accommodates operations that perform implicit shape conversions such as broadcasting.

We here note that our formalization is not something that completely replaces theorem-proving approaches; our language ensures the shape consistency only after concrete tensor shapes are given for compile-time code generation. By contrast, for verifying the consistency of a tensor-handling library, it is essentially necessary to prove that the library works for any possible combination of tensor shapes, but our method cannot ensure properties that cannot be judged by evaluation, e.g., essentially universally quantified propositions. In this sense, the target population of our language design is end users who are implementing specific heavy tensor computations rather than authors of tensor-handling libraries. Nonetheless, it is possible to utilize our method for testing libraries in a way like property-based testing [12], i.e., by running code generation with many randomly selected combinations of tensor shapes. Also, although in a somewhat awkward manner and with the aid of the so-called $\mathrm{𝐫𝐮𝐧}$-primitive [66, 69, 31, 42], our method accommodates programs that handle tensors whose sizes are known only at runtime; Section 6.2 discusses how to achieve this.

The rest of the paper is organized as follows: First, Section 2 gives an overview of our method by using running examples, and Section 3 describes the formalization of the staged core language ${\lambda }^{⟨⟩⦇⦈}$ and proves its metatheoretic safety properties. Section 4 extends ${\lambda }^{⟨⟩⦇⦈}$ with implicit parameters/arguments and gives how to infer omitted arguments. Section 5 provides Horsea as a proof-of-concept, non-staged surface language, and explains the basics of its translation to the staged language by using BTA. After that, Section 6 discusses further extension of our language with some features necessary for real-world use. Finally, Section 7 describes a prototype type-checker implemented based on our method and its example use cases involving ocaml-torch, Section 8 discusses the related work, and Section 9 concludes the paper. Figure 1 depicts the overall procedure of our method explained in Sections 3–5.

For implicit parameters/arguments, we first introduce $\{x:T_1^{\left(0\right)}\}\to T_2^{\left(0\right)}$, a variant form of stage-$0$ dependent function types. Functions of this type work exactly the same as those of type $(x:T_1^{\left(0\right)})\to T_2^{\left(0\right)}$ from the operational perspective, but they allow users to omit arguments and require the type-checker to infer an expression that should substitute $x$ for each context. For example, $\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}$ is now assigned type $\{p:\mathrm{𝙽𝚊𝚝}\}\to \{q:\mathrm{𝙽𝚊𝚝}\}\to \{r:\mathrm{𝙽𝚊𝚝}\}\to ⟨\mathrm{𝙼𝚊𝚝}\%p\%r\to \mathrm{𝙼𝚊𝚝}\%q\%r\to \mathrm{𝙼𝚊𝚝}\%(p+q)\%r⟩$, meaning $p$, $q$, and $r$ can be implicit. Users can also define a function with implicit arguments by using variant $\lambda$-abstractions of the form $(\lambda \{x:T^{\left(0\right)}\}.M^{\left(0\right)})$. In addition, when one wants to specify arguments explicitly for implicit parameters, applications of the form $M_1^{\left(0\right)}\{M_2^{\left(0\right)}\}$ can be used.

The reconstruction of implicit arguments can be done in a type-guided manner. For example, consider the subexpression $\sim \mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}BC$, and let $N_1^{\left(0\right)}$, $N_2^{\left(0\right)}$, and $N_3^{\left(0\right)}$ denote expressions that should be reconstructed for the three parameters. According to the type assigned to $\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}$, the reconstructed subexpression $\sim (\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}\{N_1^{\left(0\right)}\}\{N_2^{\left(0\right)}\}\{N_3^{\left(0\right)}\})$ will have type $\mathrm{𝙼𝚊𝚝}\%N_1^{\left(0\right)}\%N_3^{\left(0\right)}\to \mathrm{𝙼𝚊𝚝}\%N_2^{\left(0\right)}\%N_3^{\left(0\right)}\to \mathrm{𝙼𝚊𝚝}\%(N_1^{\left(0\right)}+N_2^{\left(0\right)})\%N_3^{\left(0\right)}$. Since this subexpression is applied to $B$, which has type $\mathrm{𝙼𝚊𝚝}\%m\%n$, the domain type $\mathrm{𝙼𝚊𝚝}\%N_1^{\left(0\right)}\%N_3^{\left(0\right)}$ must be identical to $\mathrm{𝙼𝚊𝚝}\%m\%n$ after evaluation. In this case, we can infer that it suffices to just substitute $N_1^{\left(0\right)}$ and $N_3^{\left(0\right)}$ with $m$ and $n$, respectively. Similarly, by comparing $\mathrm{𝙼𝚊𝚝}\%N_2^{\left(0\right)}\%N_3^{\left(0\right)}$ with the type of $C$ (i.e., $\mathrm{𝙼𝚊𝚝}\%k\%n$), we can judge that $k$ can be used as $N_2^{\left(0\right)}$. We have found that this kind of reconstruction can be formalized by using a technique adapted from Xie and Oliveira’s “let arguments go first” [73]; other well-trodden approaches such as Hindley–Milner-like unification might also work, but our formalization seems more concise (and still effective enough) in that it suffices to track variables for substitution only locally. Although the reconstruction is essentially incomplete (i.e., even when there exist appropriate expressions, it cannot always infer them), our reconstuction algorithm is fairly effective for typical use cases, as we will report later. Even when it cannot infer an argument, the algorithm can still report the position where the user should explicitly give arguments.

To this end, we use refinement types [24, 22, 55, 44] of the form $\{\nu :B\mid N^{\left(0\right)}\}$ for stage-0 types. The annotations can be modified from $\mathrm{𝙸𝚗𝚝}$ to $\mathrm{𝙽𝚊𝚝}$, which abbreviates $\{\nu :\mathrm{𝙸𝚗𝚝}\mid \nu \ge \mathrm{𝟶}\}$:

Then, consider an application $f^{\prime }{M}^{\left(0\right)}$. Now that the type for $j$ tracks the precondition required of $j$, the above application elaborates to $f^{\prime }({⦇⊲\mathrm{𝙽𝚊𝚝}⦈}^{\mathrm{\ell }}{N}^{\left(0\right)})$ (assuming $M^{\left(0\right)}$ elaborates to $N^{\left(0\right)}$), where ${⦇⊲T_{\mathrm{rfn}}^{\left(0\right)}⦈}^{\mathrm{\ell }}$ is a cast function to assert that the argument satisfies the predicate of the refinement type $T_{\mathrm{rfn}}^{\left(0\right)}$, and $\mathrm{\ell }$ is a label that points to the original application. This will emit a failure if $N^{\left(0\right)}$ evaluates to a negative integer and report the position of the call site $\mathrm{\ell }$. Compared to the original situation where failures are raised by $\mathrm{𝑔𝑒𝑛𝑀𝑎𝑡𝑀𝑢𝑙𝑡}$ used in $f^{\prime }$, programmers can describe more detailed preconditions expected of arguments, and the type-checker can point to the location where some precondition was violated. We allow refinement types only for stage $0$, and thereby the evaluation of such casts happens only at compile time, i.e., runtime evaluation is still free from assertion failures.

We note that refinement types are also beneficial for error localization as to implicit conversion; with refinement types, we can assign $\mathrm{𝑔𝑒𝑛𝐴𝑑𝑑}$ the following more natural type:

where $\mathrm{𝙽𝚊𝚝𝙻𝚒𝚜𝚝}$ is a basetype for shapes that intuitively works as $\mathrm{𝙻𝚒𝚜𝚝}\mathrm{𝙽𝚊𝚝}$, and $\mathrm{𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡𝑎𝑏𝑙𝑒}$ judges whether a pair of two shapes is in the domain of $\mathrm{𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡}$. As a sideline, $\mathrm{𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡}:(x:\mathrm{𝙽𝚊𝚝𝙻𝚒𝚜𝚝})\to \{\nu :\mathrm{𝙽𝚊𝚝𝙻𝚒𝚜𝚝}\mid \mathrm{𝑏𝑟𝑜𝑎𝑑𝑐𝑎𝑠𝑡𝑎𝑏𝑙𝑒}x\nu \}\to \mathrm{𝙽𝚊𝚝𝙻𝚒𝚜𝚝}$ is now a total function.

Given a program in Horsea, we perform BTA to find out which parts can be at compile time and insert brackets and escapes accordingly. For this conversion, built-in functions in Horsea (e.g., $\mathrm{𝑣𝑒𝑟𝑡𝐶𝑎𝑡}$) are associated with those in the staged language (e.g., $\mathrm{𝑔𝑒𝑛𝑉𝑒𝑟𝑡𝐶𝑎𝑡}$).

The most essential part of the syntax is that stage-$1$ tensor types $\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\%M^{\left(0\right)}$ have a stage-$0$ expression $M^{\left(0\right)}$ of type $\mathrm{𝙽𝚊𝚝𝙻𝚒𝚜𝚝}$ to represent a tensor shape⁸⁸8While $\mathrm{𝙽𝚊𝚝}$ is just a shorthand for $\{\nu :\mathrm{𝙸𝚗𝚝}\mid \nu \ge \mathrm{𝟶}\}$, $\mathrm{𝙽𝚊𝚝𝙻𝚒𝚜𝚝}$ is a base type. This will simplify metatheory. . This gap between tensor types and their argument expressions as to stages ensures, for example, that all the stage-$1$ binders of the form $(\lambda x:\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\%[M_1^{\left(0\right)},M_2^{\left(0\right)}].\mathrm{\dots })$ will be $(\lambda x:\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\%[n_1,n_2].\mathrm{\dots })$ after compile-time computation and thereby that every tensor in generated code has a specialized shape. Stage-$1$ types $\mathrm{𝚅𝚎𝚌}\%M_1^{\left(0\right)}$ and $\mathrm{𝙼𝚊𝚝}\%M_1^{\left(0\right)}\%M_2^{\left(0\right)}$ for vectors and matrices can be provided as syntax sugar of $\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\%[M_1^{\left(0\right)}]$ and $\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}\%[M_1^{\left(0\right)},M_2^{\left(0\right)}]$, respectively. The symbol $\%$ in the notation indicates this gap by an analogy to the notion of cross-stage persistence [31, 40, 66, 74]. By contrast, stage-$0$ tensor types $\mathrm{𝚃𝚎𝚗𝚜𝚘𝚛}s$ have fixed shapes; the metavariable $s$ ranges over the set of finite sequences of natural numbers, and these types can be thought of as a family of countably infinite base types. Stage-$0$ tensor types occur mostly in stage-$0$ terms obtained by unlifting generated code and are seldom written by users.

In response to the setting of stage-$1$ tensor types, stage-$0$ function types can be dependent ones. This allows, for instance, $\mathrm{𝑔𝑒𝑛𝑀𝑎𝑡𝑀𝑢𝑙𝑡}$ to be assigned $(p:\mathrm{𝙽𝚊𝚝})\to (q:\mathrm{𝙽𝚊𝚝})\to (r:\mathrm{𝙽𝚊𝚝})\to ⟨\mathrm{𝙼𝚊𝚝}\%p\%q\to \mathrm{𝙼𝚊𝚝}\%q\%r\to \mathrm{𝙼𝚊𝚝}\%p\%r⟩$. However, this does not apply to stage $1$; by restricting stage-$1$ function types to non-dependent ones, compile-time assertions are considerably simplified, and many realistic programs can still be supported, as exemplified in Section 7. We also use the notation $T_1^{\left(0\right)}\to T_2^{\left(0\right)}$ for $(x:T_1^{\left(0\right)})\to T_2^{\left(0\right)}$, where $x\notin fv(T_2^{\left(0\right)})$.

The target syntax is basically an “assertion-included” variant of the source syntax. It consists of stage-$b$ assertive terms $N^{(b)}$ and assertive types $T^{(b)}$ defined by the following:

Stage-$0$ assertions have two forms. The first one is ${⦇⟨T_1^{\left(1\right)}⟩⊲⟨T_2^{\left(1\right)}⟩⦈}^L$, which judges that the stage-$1$ types $T_1^{\left(1\right)}$ and $T_2^{\left(1\right)}$ syntactically coincide after evaluation. Interestingly, adding assertions of this form covers all the necessary cases for checking type equality itself. The attachment $L$ stands for the source of errors when the assertion fails. The other one ${⦇⊲\{\nu :B\mid N_1^{\left(0\right)}\}⦈}^L$ is used for checking the validity of downcasts. In operational terms, it basically works as an identity function, but before returning the argument (say $c$) as is, it checks whether the argument satisfies the predicate $N_1^{\left(0\right)}$, i.e., whether $[c/\nu ]N_1^{\left(0\right)}$ evaluates to $\mathrm{𝚝𝚛𝚞𝚎}$. The variant ${⦇\{\nu :B\mid N_1^{\left(0\right)}\},N_2^{\left(0\right)},c⦈}^L$ called an active check is the intermediate form of this assertion process, where $N_2^{\left(0\right)}$ is a “refinement proposition” (i.e., the application of the refinement predicate $N_1^{\left(0\right)}$ to the tested value $c$) under evaluation. It also keeps $c$ separately so that it will evaluate to $c$ when the assertion passes. Only the first one of these two assertion forms is new; the latter device is ported from the context of manifest contracts [30, 58].

Typing judgments are defined as $\Gamma {⊢}^b{M}^{(b)}:T^{(b)}\rightsquigarrow N^{(b)}$ (for $b\in \{0,1\}$), which can be read as “under the type environment $\Gamma$, the source expression $M^{(b)}$ has type $T^{(b)}$ and elaborates to $N^{(b)}$ by assertion insertion.” The structure of type environments are defined by: $\Gamma \mathit{::=}\bullet \mid \Gamma ,x:{(T^{\left(0\right)})}^0\mid \Gamma ,x:{(T^{\left(1\right)})}^1$, i.e, $\Gamma$ tracks the stage at which each variable was bound. Figure 2 shows the typing rules for these judgments. The rules S0-Brkt and S1-Esc are peculiar to staged computation and are natural extensions from the literature. S0-Abs and S1-Abs elaborate the type annotations by using the judgments $\Gamma {⊢}^b{S}^{(b)}\rightsquigarrow T^{(b)}$. Also, for constants, we have S0-Cst0, S0-CstP, and S1-CstP. Here, we use two environments ${ℬ}_0$ and ${ℬ}_{\%}$; the former maps stage-$0$-specific built-in functions $p$ to stage-$0$ types $T^{\left(0\right)}$, and the latter works similarly for stage-agnostic constants $c$. The metavariable ${\tau }^{\left(1\right)}$ and its unlifting $\downarrow ({\tau }^{\left(1\right)})=T^{\left(0\right)}$ will be introduced shortly. Entries are like the following:

The most distinctive typing rules are those for applications, i.e., S0-App and S1-App. Here, $T_1^{\left(1\right)}{||}^1{T}_2^{\left(1\right)}$ judges type compatibility, i.e., type equivalence ignoring the difference of argument expressions, and $\Gamma {⊢}_L{T}_1^{\left(0\right)}⊲T_2^{\left(0\right)}\rightsquigarrow N_0^{\left(0\right)}$ generates assertive cast terms $N_0^{\left(0\right)}$. Figure 3 displays the rules for these judgments. Among these rules, I-Code is the core of the cast term generation; it checks the compatibility of given two types and simply produces an “atomic” cast term ${⦇⟨T_1^{\left(1\right)}⟩⊲⟨T_2^{\left(1\right)}⟩⦈}^L$. I-Arr, which works for functions, is another characteristic rule. It produces two cast terms for domain types and codomain types, respectively, and combines them. Thanks to this rule, we can handle higher-order programs without any hindrance. This mechanism is partially inspired by GraTen [32] and ${\mathrm{F}}_{\mathrm{H}}^{\sigma }$ [58].