Computer Science
Computer Science-2021.gz:143186
Computer Science
ad887cc0812d228016a2fd32bfc82c95119e99d6
CHAD: Combinatory Homomorphic Automatic Differentiation

We introduce Combinatory Homomorphic Automatic Differentiation (CHAD), a principled, pure, provably correct define-then-run method for performing forward and reverse mode automatic differentiation (AD) on programming languages with expressive features. It implements AD as a compositional, type-respecting source-code transformation that generates purely functional code. This code transformation is principled in the sense that it is the unique homomorphic (structure preserving) extension to expressive languages of Elliott's well-known and unambiguous definitions of AD for a first-order functional language. Correctness of the method follows by a (compositional) logical relations argument that shows that the semantics of the syntactic derivative is the usual calculus derivative of the semantics of the original program. In their most elegant formulation, the transformations generate code with linear types. However, the code transformations can be implemented in a standard functional language lacking linear types: while the correctness proof requires tracking of linearity, the actual transformations do not. In fact, even in a standard functional language, we can get all of the type-safety that linear types give us: we can implement all linear types used to type the transformations as abstract types, by using a basic module system. In this paper, we detail the method when applied to a simple higher-order language for manipulating statically sized arrays. However, we explain how the methodology applies, more generally, to functional languages with other expressive features. Finally, we discuss how the scope of CHAD extends beyond applications in AD to other dynamic program analyses that accumulate data in a commutative monoid.


INTRODUCTION
Automatic differentiation (AD) is a technique for transforming code that implements a function into code that computes 's derivative, essentially by using the chain rule for derivatives. Due to its efficiency and numerical stability, AD is the technique of choice whenever we need to compute derivatives of functions that are implemented as programs, particularly in high dimensional settings. Optimization and Monte-Carlo integration algorithms, such as gradient descent and Hamiltonian Monte-Carlo methods, rely crucially on the calculation of derivatives. These algorithms are used in virtually every machine learning and computational statistics application, and the calculation of derivatives is usually the computational bottleneck. These applications explain the recent surge of interest in AD, which has resulted in the proliferation of popular AD systems such as TensorFlow [1], PyTorch [36], and Stan Math [9].
AD, roughly speaking, comes in two modes: forward mode and reverse mode. When differentiating a function R → R , forward mode tends to be more efficient if ≫ , while reverse mode generally is more performant if ≫ . As most applications reduce to optimization or Monte-Carlo integration of an objective function R → R with very large (at the time of this paper, in the order of 10 4 − 10 7 ), reverse mode AD is in many ways the more interesting algorithm [5].
However, reverse AD is also more complicated to understand and implement than forward AD. Forward AD can be implemented as a structure-preserving program transformation, even on languages with complex features [38]. As such, it admits an elegant proof of correctness [21]. By contrast, reverse AD is only well-understood as a compiletime source-code transformation that does not rely on using a run-time interpreter (also called define-then-run style AD) on limited programming languages, such as first-order functional languages. Typically, its implementations on more expressive languages that have features such as higher-order functions make use of interpreted define-by-run approaches. These approaches first build a computation graph during runtime, effectively evaluating the program until a straight-line first-order program is left, and then they perform automatic differentiation on this new program [9,36].
Such approaches firstly have the severe downside that they can suffer from interpretation overhead. Secondly, the differentiated code cannot benefit as well from existing optimizing compiler architectures. As such, these AD libraries need to be implemented using carefully, manually optimized code, that, for example, does not contain any common subexpressions. This implementation process is precarious and labour intensive. Furthermore, some whole-program optimizations that a compiler would detect go entirely unused in such systems.
Similarly, correctness proofs of reverse AD have taken a define-by-run approach or have relied on non-standard operational semantics, using forms of symbolic execution [2,8,31]. Most work that treats reverse AD as a source-code transformation does so by making use of complex transformations which introduce mutable state and/or non-local control flow [37,44]. As a result, we are not sure whether and why such techniques are correct. Furthermore, AD applications (e.g. in machine learning) tend to be run on parallel hardware, which can be easier to target with purely functional code. Another approach has been to compile high-level languages to a low-level imperative representation first, and then to perform AD at that level [22], using mutation and jumps. This approach has the downside that we might lose important opportunities for compiler optimizations, such as map-fusion and embarrassingly parallel maps, which we can exploit if we perform define-then-run AD on a high-level functional representation.
A notable exception to these define-by-run and non-functional approaches to AD is Elliott's work [16], which presents an elegant, purely functional, define-then-run version of reverse AD. Unfortunately, their techniques are limited to first-order programs over tuples of real numbers. The workshop paper [43] by Vytiniotis, Belov, Wei, Plotkin, Manuscript submitted to ACM and Abadi proposes two possible extensions of Elliott's functional AD to accommodate higher-order functions. However, it does not address whether or why these extensions would be correct or establish a more general methodology for applying AD to languages with expressive features. This paper introduces Combinatory Homomorphic Automatic Differentiation (CHAD) and its proof of correctness.
CHAD is based on the observation that Elliott's work in [16] has a unique structure preserving extension that lets us perform AD on various expressive programming language features. We see purely functional higher-order (parallel) array processing languages such as Accelerate [10] and Futhark [19] as particularly relevant platforms for the machine learning applications that AD tends to be used for. With that in mind, we detail CHAD when applied to higher-order functional programs over (primitive) arrays of reals. This paper includes the following: • We introduce CHAD, a categorical perspective on AD, that lets us see AD as a uniquely determined homomorphic (structure-preserving) functor from the syntax of its source programming language ( §3) to the syntax of its target language ( §4).
• We explain, from this categorical setting, precisely in what sense reverse AD is the "mirror image" of forward AD ( §6).
• We detail how this technique lets us define purely functional define-then-run reverse mode AD on a higher-order language ( §7).
• We present an elegant proof of semantic correctness of the resulting AD transformations, based on a semantic logical relations argument, demonstrating that the transformations calculate the derivatives of the program in the usual mathematical sense ( §5 and §8).
• We show that the AD definitions and correctness proof are extensible to higher-order primitives such as a map-operation over our primitive arrays ( §10).
• We show how our techniques are readily implementable in standard functional languages to give purely functional, principled, semantically correct, compositional, define-then-run reverse mode AD ( §9).
• Finally, we place CHAD in a broader context and explain how it applies, more generally, to dynamic program analyses that accumulate information in a commutative monoid ( §11).
We start by giving a high-level overview of the main insights and theorems in this paper in §2.
For a review of the basics of AD, we refer the reader to [5,32]. We discuss recent related work studying AD from a programming languages perspective in §12.

KEY IDEAS
We start by providing a high-level overview of the paper, highlighting the main insights and theorems underlying our contributions.

Aims of Automatic Differentiation
The basic challenge that automatic differentiation aims to solve is the following. We are given a program : real ⊢ : real that takes an -dimensional array of (floating point) real numbers as input and produces an -dimensional array of reals as output. That is, computes some mathematical function : R → R . We want to transform the code of to: 1 1 The program transformations are really called Here, we write R for the space of (co)tangent vectors to R ; we regard R as a commutative monoid under elementwise addition. We write ⊸ for a linear function type to emphasize that derivatives are linear in the sense of being monoid homomorphisms.
Furthermore, we have some more desiderata for these code transformations: (1) we want these code transformations to be defined compositionally, so we can easily extend the source programming language we apply the transformations to with new primitives; (2) we want these transformations to apply to a wide range of programming techniques, so we are not limited in our programming style even if we want our code to be differentiated; (3) we want the transformations to generate purely functional code so we can easily prove its correctness and deploy it on parallel hardware; (4) we want the code size of  2 to be proportional to that of and, generally, as low as possible; this means that we can use forward AD to efficiently compute a column of the Jacobian matrix of partial derivatives, while reverse AD efficiently computes a row of the Jacobian.
In this paper, we demonstrate how the CHAD technique of automatic differentiation satisfies desiderata (1)-(4) -we leave (5) to future work. It achieves this by taking seriously the mathematical structure of programming languages as freely generated categories and by observing that differentiation is compositional according to the chain rule.

The Chain Rule -Pairing and Sharing of Primals and Derivatives
To achieve desideratum (1) of compositionality, it is tempting to examine the chain rule, the key property of compositionality of derivatives. Given : R → R , we write for the function that pairs up the primal function value ( ) with the derivative ( ) of at that acts on tangent vectors . The chain rule then gives the following formula for the derivative of the composition ; of and : T( ; )( ) = (T 1 (T 1 ( )), T 2 ( ); T 2 (T 1 ( ))), where we write T 1 def = T ; 1 and T 2 def = T ; 2 for the first and second components of T respectively. We make two observations: (1) the derivative of the composition ; does not only depend on the derivatives of and but also on the primal value of ; (2) the primal value of is used twice: once in the primal value of ; and once in its derivative; we want to share these repeated subcomputations, to address desiderata (4) and (5). Similarly, we can pair up 's transposed (adjoint) derivative , which propagates not tangent vectors but cotangent vectors: It then satisfies the following chain rule, which follows from the usual chain rule above together with the fact that ( ; ) = ; for linear maps and (transposition is contravariant -note the resulting reversed order of T * 2 and T * 2 for reverse AD): T * ( ; )( ) = (T * 1 (T * 1 ( )), T * 2 (T * 1 ( )); T * 2 ( )). Again, pairing and sharing the primal and (transposed) derivative computations is beneficial.
CHAD directly implements the operations T and T * as source code transformations − → D and ← − D on a functional language to implement forward 2 and reverse mode AD, respectively. These code transformations are defined compositionally through structural induction on the syntax, by exploiting the chain rules above combined with the categorical structure of programming languages.

CHAD on a First-Order Functional Language
Here, we outline how CHAD looks when applied to programs written in a first-order functional language. We consider this material known, as it is essentially the algorithm of [16]. However, we present it in terms of a -calculus rather than categorical combinators, by applying the well-known mechanical translations between both formalisms [13]. We hope that this presentation makes the algorithm easier to apply in practice.
We consider a source programming language (see §3) where we write , , for types that are either statically sized arrays of real numbers real or tuples * of types , . These types will be called first-order types in this section. 3 We consider programs of type in a typing context Γ = 1 : 1 , . . . , : , where are identifiers. We write such typing of programs in a context as Γ ⊢ : . As long as our language has certain primitive operations (which we represent schematically) Γ ⊢ 1 : real 1 · · · Γ ⊢ : real Γ ⊢ op( 1 , . . . , ) : real such as constants (as nullary operations), (elementwise) addition and multiplication of arrays, inner products and certain non-linear functions like sigmoid functions, we can write complex programs by sequencing together such operations. Fig. 1 (a) and (b) give some examples of some programs we can write, where we write real for real 1 and indicate shared subcomputations with let-bindings.
CHAD transforms the types and programs of this source language into types and programs of a suitably chosen target language (see §4) that is a superset of the source language. CHAD associates to each source language type types of: 2 For forward AD, we can also choose to implement instead together with its chain rule as code transformations. This leads to a different style of forward AD based on a dual numbers representation. [21] gives an analysis of this style of forward AD, similar to the treatment of reverse AD and (non-dual number) forward AD in this paper. Although dual numbers forward AD is more efficient in its memory use and preferable in practical implementations, it does not have an obvious reverse mode variant. See § §11.1 for more discussion. 3 In the rest of the paper, we also consider the unit type 1 a first-order type. These types are also called ground types.
• forward mode primal values The types That is, such linear types in our language need to have a commutative monoid structure (0, +). For example, real is the commutative monoid over real with 0 is the zero vector and (+) as elementwise addition of vectors. Derivatives and transposed derivatives are then linear functions, that is, homomorphisms of this (0, +)-monoid structure. As we will see, we use the monoid structure to initialize and accumulate (co)tangents in the definition of CHAD.
We extend these operations − → D and ← − D to act not only on types but also on typing contexts Γ = 1 : 1 , . . . , : to produce primal contexts and (co)tangent types: To each program Γ ⊢ : , CHAD then associates programs calculating the forward mode and reverse mode derivatives − → D Γ ( ) and ← − D Γ ( ), whose definitions make use of the list Γ of identifiers that occur in Γ: Observing that each program computes a differentiable (infinitely differentiable) function between Euclidean spaces, as long as all primitive operations op are differentiable, the key property that we prove for these code transformations is that they actually calculate derivatives: Theorem A (Correctness of CHAD, Thm. 8.3). For any well-typed program (where and are first-order types, i.e. real and tuples of such types) 1 : 1 , . . . , : ⊢ : Once we fix a semantics for the source and target languages, we can show that this theorem holds if we define − → D and ← − D on programs using the chain rule. The proof works by plain induction on the syntax. For example, we can correctly define reverse mode CHAD on a first-order language as follows (see §7): Here, we write v. for a linear function abstraction (merely a notational convention -it can simply be thought of as a plain function abstraction) and • for a linear function application of : ⊸ to the argument : (which again can be thought of as a plain function application). Furthermore, given a program of tuple type 1 * · · · * , we write proj for its -th projection of type . Similarly, given a program of linear type , we write coproj ( ) for the -th coprojection 0, . . . , 0, , 0, . . . , 0 of type 1 * · · · * and we write idx( ; 1 , . . . , ) = for the index of an identifier in a list of identifiers. Finally, op here is a linear operation that implements the transposed derivative of the primitive operation op. We note that we crucially need the commutative monoid structure on linear types to correctly define the reverse mode derivatives of programs that involve tuples (or -ary operations for ≠ 1). Intuitively, matrix transposition (of derivatives) flips the copying-deleting comonoid structure provided by tuples into the addition-zero monoid structure.

Insight 2.
In functional define-then-run reverse AD, we need to have a commutative monoid structure on types of cotangents to mirror the comonoid structure coming from tuples: copying fan-out in the original program gets translated into fan-in in the transposed derivative, for accumulating incoming cotangents. This leads to linear types of cotangents.
Furthermore, observe that CHAD pairs up primal and (co)tangent values and shares common subcomputations, as desired. We see that what CHAD achieves is a compositional efficient reverse mode AD algorithm that computes the (transposed) derivatives of a composite program in terms of the (transposed) derivatives op of the basic building blocks op. Finally, it does so in a way that satisfies desiderata (1)-(5).
For example, Fig. 1 (c) and (d) display the code that forward and reverse mode CHAD, respectively, generate for the source programs of (a) and (b). This is the code that is actually generated by the CHAD code transformations in our Haskell implementation followed by some very basic simplifications that do not affect time complexity and whose only purpose here is to aid legibility. For more information about how exactly this code relates to the output one gets when applying the forward and reverse AD macros in this paper to the source programs, see Appendix B.

Intermezzo: the Categorical Structure of CHAD
While this definition of CHAD on a first-order language straightforwardly follows from the mathematics of derivatives, it is not immediately clear how to should be extended to source languages with more expressive features such as higherorder functions. Indeed, we do not typically consider derivatives of higher-order functions in calculus. In fact, it is not even clear what a tangent or cotangent to a function type should be, or, for that matter, what a primal associated with a value of function type is. To solve this mystery, we employ some category theory.
Observe that the first-order source language we consider can be viewed as a category Syn with products (see §3): its objects are types , , and morphisms ∈ Syn( , ) are programs : ⊢ : modulo standard -program equivalence (identities are given by variables and composition is done through let-bindings). This category is freely generated by the objects real and morphisms op in the sense that any consistent assignment of objects (real ) and morphisms (op) in a category with products C extends to a unique product preserving functor : Syn → C.
l e t y = 2 * x z = x * y w = c o s z v = y , z , w in v (a) Original program : real ⊢ : real * real * real computing a function : l e t y = x1 * x4 + 2 * x2 z = y * x3 w = z + x4 v = s i n w in v (b) Original program 1 : real, 2 : real, 3 : real, 4 : real ⊢ : real computing a function : . We see that the copying of 4 is translated into the addition of cotangents in 4 ′ . We observe that we can define categories Σ C L and Σ C L (their so-called Grothendieck constructions, or Σ-types, see §6) with objects that are pairs ( 1 , 2 ) with 1 an object of C and 2 an object of L and homsets We prove that these categories have finite products, provided that some conditions are satisfied: namely that C has finite products and L has indexed finite biproducts (or equivalently: has indexed finite products and is enriched over commutative monoids). Indeed, then ∈ ( 1 , 2 ) = ( ∈ 1 , ∈ 2 ). In other words, it is sufficient if our model Manuscript submitted to ACM of linear logic is biadditive. In particular, the categorical model of linear logic that we can build from the syntax of our target language for CHAD, LSyn : CSyn → Cat, satisfies our conditions (in fact, it is the initial model that does so), so Σ CSyn LSyn and Σ CSyn LSyn have finite products. By the universal property of the source language Syn, we obtain a canonical definition of CHAD. The definitions following from this universal property reproduce the definitions of CHAD that we have given so far.
Intuitively, the linear types represent commutative monoids, implementing the idea that (transposed) derivatives are linear functions in the sense that ( )(0) = 0 and ( )( + ′ ) = ( )( ) + ( )( ′ ). We have seen that this commutative monoid structure is important when writing down the definitions of AD as a source-code transformation.
Seeing that a higher-order language can be viewed as a freely generated Cartesian closed category Syn, it is tempting to find a suitable target language such that Σ CSyn LSyn and Σ CSyn LSyn are Cartesian closed. Then, we can define CHAD on this higher-order language via Thm. B.

Insight 3.
To understand how to perform CHAD on a source language with language feature (e.g., higher-order functions), we need to understand the categorical semantics of language feature (e.g., categorical exponentials) in categories of the form Σ C L and Σ C L . Giving sufficient conditions on a model of linear logic L : C → Cat for such a semantics to exist yields a suitable target language for CHAD as the initial such model LSyn : CSyn → Cat , with the definition of the algorithm falling from the universal property of the source language.
2.5 Cartesian Closure of Σ C L and Σ C L and CHAD of Higher-Order Functions Having had this insight, we identify conditions on a locally indexed category L : C → Cat that are enough to guarantee that Σ C L and Σ C L are Cartesian closed (see §6).
Then, Σ C L and Σ C L are Cartesian closed with, respective, exponentials: In particular, if we extend our target language with (linear) powers, (linear) copowers and (Cartesian) function types, then LSyn : CSyn → Cat satisfies the conditions of Thm. C, so we can extend Thm. B to our higher-order source language. In particular, we find the following definitions of CHAD for primals and (co)tangents to function types: Interestingly, we see that for higher-order programs, the primal transformations are no longer the identity. Indeed, the  On programs, we obtain the following extensions of our definitions for reverse AD: With regards to ctx of with respect to the context variables from the cotangent of and obtain the transposed derivative ′ arg of with respect to its function arguments from 's primal. We combine these two halves of the transposed derivative with 's transposed derivative (which we get from its cotangent) to get the correct transposed derivative for the function application .

Proving CHAD Correct
With these definitions in place, we turn to the correctness of the source-code transformations. To phrase correctness, we first need to construct a suitable semantics with an uncontroversial notion of semantic differentiation (see §5).
We choose to work with a semantics in terms of the category Set of sets and functions 4 , noting that any function : R → R has a unique derivative as long as is differentiable. We will only be interested in this semantic notion of derivative of first-order functions for the sake of correctness of AD, and we will not concern ourselves with semantic derivatives of higher-order functions. We interpret the required linear types in the category CMon of commutative monoids and homomorphisms.
By the universal properties of the syntax, we obtain canonical, structure-preserving (homomorphic) functors − : Syn → Set, − : CSyn → Set and − : LSyn → CMon once we fix interpretations R of real and well-typed (differentiable) interpretations op for each operation op. These functors define a concrete denotational semantics for our source and target languages.
Having constructed the semantics, we can turn to the correctness proof (of §8). Because calculus does not provide an unambiguous notion of derivative at function spaces, we cannot prove that the AD transformations correctly implement mathematical derivatives by plain induction on the syntax. Instead, we use a logical relations argument over the semantics.
Insight 5. Once we show that the (transposed) derivatives of primitive operations op are correctly implemented, correctness of (transposed) derivatives of all other programs follows from a standard logical relations construction over the semantics that relates a curve to its primal and (co)tangent curve. By the chain rule for (transposed) derivatives, all CHAD transformed programs respect the logical relations. By basic calculus results, CHAD therefore has to compute the (transposed) derivative.
In §8 we present an elegant high-level formulation of this correctness argument, using categorical logical relations techniques (subsconing). To make this argument accessible to a wider audience of readers, we present here a low-level description of the logical relations argument. The reader may note that these arguments look significantly different from the usual definitions of logical relations. That difference is caused by the non-standard Cartesian closed structure of Σ C L and Σ C L and the proof is entirely standard when viewed from the higher level of abstraction that subsconing gives us.
We first sketch the correctness argument for forward mode CHAD. By induction on the structure of types, writing ( , ′ ) for the product pairing of and ′ , we construct a logical relation on types as We extend the logical relation to typing contexts Γ = 1 : 1 , . . . , : as Γ def = 1 * ··· * . Then, we establish the following fundamental lemma, which says that all well-typed source language programs respect the logical relation.
The proof goes via induction on the typing derivation of . The main remaining step in the argument is to note that any tangent vector at 1 * · · · * R , for first-order , can be represented by a curve R → 1 * · · · * .
Similarly, for reverse mode CHAD, we define, by induction on the structure of types, a logical relation on types (and, as before, we also define Γ = 1 * ··· * for typing contexts Γ = 1 : 1 , . . . , : ): Then, we establish the following fundamental lemma.
For any source language program Γ ⊢ : and any : The proof goes via induction on the typing derivation of . Correctness follows from the fundamental lemma by observing that any cotangent vector at 1 * · · · * R , for first-order , can be represented by a curve R → 1 * · · · * . We obtain our main theorem, Thm. A, but now for our CHAD algorithms applied to a higher-order source language.

A Practical Implementation in Haskell
Next, we address the practicality of our method (in §9). The code transformations we employ are not too daunting to implement and they are well-behaved in the sense that the derivative code they generate grows linearly in the size of the original source code. However, the implementation of the required linear types presents a challenge. Indeed, types like !(−) ⊗ (−) and (−) ⊸ (−) are absent from languages such as Haskell and OCaml. Luckily, in this instance, we can implement them using abstract data types by using a (basic) module system: Insight 6. Under the hood, ! ⊗ can consist of a list of values of type * . Its API ensures that the list order and the difference between ++ [( , ), ( , ′ )] and ++ [( , + ′ )] (or and ++ [( , 0)]) cannot be observed: as such, it is a quotient type. Meanwhile, ⊸ can be implemented as a standard function type → with a limited API that enforces that we can only ever construct linear functions: as such, it is a subtype.
This idea leads to our reference implementation of CHAD in Haskell (available at https://github.com/VMatthijs/CHAD), which generates perfectly standard simply typed functional code that is given a bit of extra type-safety with the linear types, implemented as abstract types. To illustrate what our method does in practice, we consider two programs of our higher-order source language, that we may want to differentiate in Fig. 2 (a) and (b). The forward and reverse mode derivatives that our CHAD implementation generates for these programs are listed in Fig. 2 (c) and (d), again modulo minor simplifications that aid legibility but have no significant run-time implications. 5 In §9, we also phrase the correctness proof of the AD transformations in elementary terms, such that it holds in the applied setting where we use abstract types to implement linear types. We show that our correctness results are l e t f = z . x * z + 1 z s = r e p l i c a t e x ys = map f z s in ys (a) Original program : real ⊢ : real computing a function l e t f = x 2 i . x1 * x 2 i ys = map f x2 w = sum ys in w

Fig. 2.
Forward and reverse AD illustrated on higher-order functional array processing programs. The parts of the programs that involve AD on higher-order functions are marked in blue. Observe that in (c) (respectively, (d)), the primal value associated with the function from program (a) (respectively, (b)) computes both the original function as well as its derivative (respectively, transposed derivative) with respect to its argument (respectively, 2 ). In (c), the tangent ′ to is produced by propagating forwards the tangent ′ to the context variable that captures, by using 's derivative with respect to . This lets us correctly propagate forwards the contributions to ′ from both 's dependence on its argument and on its context variable . Dually, in (d), the cotangent to , which we construct from the cotangent ′ , is consumed by propagating it backwards to the cotangent 1 ′ to the context variable 1 that captures, by using 's transposed derivative with respect to 1. Meanwhile, the adjoint 2 ′ is constructed using the part of the primal of that captures 's transposed derivative with respect to 2 .
meaningful, as they make use of a denotational semantics that is adequate with respect to the standard operational semantics. Furthermore, to stress the applicability of our method, we show in §10 that it extends to higher-order (primitive) operations, such as map.
Finally (in §11), we zoom out and reflect on how this method generalizes. The crux of CHAD is in the following steps: • view the source language as a freely generated category Syn with some appropriate structure S (such as Cartesian closure, coproducts, (co)inductive types, iteration), generated from objects real and morphisms op; • find a suitable target language LSyn (with linear types arising from the effect of commutative monoids) for the translation such that Σ CSyn LSyn and Σ CSyn LSyn are categories with the structure S; in our experience, this is possible for most common choices of S corresponding to programming language constructs; • then, by the universal property of Syn, we obtain unique structure preserving (homomorphic) functors We conclude by observing that, in this sense, CHAD is not specific to automatic differentiation at all. We can choose different generators than real and op for Syn and different mappings of these generators under

-CALCULUS AS A SOURCE LANGUAGE FOR AUTOMATIC DIFFERENTIATION
As a source language for our AD translations, we can begin with a standard, simply typed -calculus which has ground types real of statically sized 6 arrays of real numbers, for all ∈ N, and sets Op 1 ,..., of primitive operations op for all , , 1 , . . . , ∈ N. These operations will be interpreted as differentiable 7 Examples to keep in mind for op include • constants ∈ Op for each ∈ R , for which we slightly abuse notation and write () as ; • elementwise addition and product (+), ( * ) ∈ Op , and matrix-vector product (★) ∈ Op · , ; • operations for summing all the elements in an array: sum ∈ Op 1 ; • some non-linear functions like the sigmoid function ∈ Op 1 1 .
We intentionally present operations in a schematic way, as primitive operations tend to form a collection that is added to in a by-need fashion, as an AD library develops. The precise operations needed will depend on the applications. In statistics and machine learning applications, Op tends to include a mix of multi-dimensional linear algebra operations and mostly one-dimensional non-linear functions. A typical library for use in machine learning would work with multi-dimensional arrays (sometimes called "tensors"). We focus here on one-dimensional arrays as the issues of how precisely to represent the arrays are orthogonal to the concerns of our development.
The types , , and terms , , of our AD source language are as follows: The typing rules are in Fig. 3. We use the usual conventions for free and bound variables and write the captureavoiding substitution of with in as [ / ]. We employ the usual syntactic sugar , .
[ fst / , snd / ], and we write real for real 1 . As Fig. 4 displays, we consider the standard -equational theory for our language, where equations hold on pairs of terms of the same type in the same context. We could consider further equations for our operations, but we do not as we will not need them.
This standard -calculus is widely known to be equivalent to the free Cartesian closed category Syn generated by the objects real and the morphisms op (see [27]).
Syn has the following well-known universal property.

LINEAR -CALCULUS AS AN IDEALISED AD TARGET LANGUAGE
As a target language for our AD source code transformations, we consider a language that extends the language of §3 with limited linear types. We could opt to work with a full linear logic as in [3] or [6]. Instead, however, we will only include the bare minimum of linear type formers that we actually need to phrase the AD transformations. The resulting language is closely related to, but more minimal than, the Enriched Effect Calculus of [14]. We limit our language in this way because we want to stress that the resulting code transformations can easily be implemented in existing functional languages such as Haskell or OCaml. As we discuss in §9, the idea will be to make use of a module system to implement the required linear types as abstract data types.
In our idealised target language, we consider linear types (aka computation types) , , , in addition to the Cartesian types (aka value types) , , that we have considered so far. We think of Cartesian types as denoting sets and linear types as denoting sets equipped with an algebraic structure. The Cartesian types will be used to represent sets of primals. The relevant algebraic structure on linear types, in this instance, turns out to be that of a commutative monoid, as this algebraic structure is needed to phrase automatic differentiation algorithms. Indeed, we will use the linear types to denote sets of (co)tangent vectors. These (co)tangents form a commutative monoid under addition.
Concretely, we extend the types and terms of our language as follows: | v linear identifier We work with linear operations lop ∈ LOp 1 ,..., 1 ,..., ; ′ 1 ,..., ′ , which are intended to represent differentiable functions that are linear (in the sense of respecting 0 and +) in the last arguments but not in the first . We write in LOp as the derivatives of each (Cartesian) operation op ∈ Op.
Manuscript submitted to ACM let 6. Equational rules for the idealised, linear AD language, which we use on top of the rules of Fig. 4. In addition to standard -rules for !(−) ⊗ (−)-and ⊸-types, we add rules making (0, +) into a commutative monoid on the terms of each linear type as well as rules which say that terms of linear types are homomorphisms in their linear variable.
In addition to the judgement Γ ⊢ : , which we encountered in §3, we now consider an additional judgement Γ; v : ⊢ : . While we think of the former as denoting a function between sets, we think of the latter as a function from the set which Γ denotes to the set of monoid homomorphisms from the denotation of to that of .
Example 5.1. The real numbers R form a commutative monoid with 0 and + equal to the number 0 and addition of numbers.
Example 5.2. Given commutative monoids ( ) ∈ , we can form the product monoid ∈ with underlying set Given a set and a commutative monoid , we can form the power monoid ⇒ def = ∈ as the -fold self-product monoid.
Ex. 5.2 gives the categorical product in CMon. We can, for example, construct a commutative monoid structure on any . . , − 1} ⇒ R by combining the one on R with the power monoid structure.
Example 5.3. Given commutative monoids ( ) ∈ , we can form the coproduct monoid ∈ with underlying set Given a set and a commutative monoid , we can form the copower monoid ! ⊗ def = ∈ as the -fold self-coproduct monoid. We Ex. 5.3 gives the categorical coproduct in CMon.
Example 5.4. Given commutative monoids and , we can form the commutative monoid ⊸ of homomor- Ex. 5.4 gives the categorical internal hom in CMon. Commutative monoid homomorphisms ! ⊗ → are in 1-1- Finally, a category C is called CMon-enriched if we have a commutative monoid structure on each homset C( , ′ ) and function composition gives monoid homomorphisms C( , ′ ) → C( ′ , ′′ ) ⊸ C( , ′′ ). In a category C with finite products, these products are well-known to be biproducts (i.e. simultaneously products and coproducts) if and only if C is CMon-enriched (for more details, see for example [17]

Abstract Denotational Semantics
By the universal property of Syn (Prop. 3.1), the language of §3 has a canonical interpretation in any Cartesian closed category (C, 1, ×, ⇒), once we fix C-objects real to interpret real and C-morphisms op ∈ C( Dom(op) , real ) to interpret op ∈ Op 1 ,..., . That is, any Cartesian closed category with such a choice of objects and morphisms is a categorical model of the source language of §3. We interpret types and contexts Γ as C-objects and Γ : We interpret terms Γ ⊢ : as morphisms in C( Γ , ): We discuss how to extend − to apply to the full target language of §4 by defining an appropriate notion of categorical model for the target language of §4.
• L interprets primitive types and operations: we have a choice real ∈ ob L to interpret real and, for each lop ∈ LOp 1 ,..., In particular, any biadditive model of intuitionistic linear/non-linear logic [6,17,35] is such a categorical model, as long as we choose interpretations for primitive types and operations.
Next, we turn to the interpretation of our target language in such models, which gives an operational intuition of the different components of a categorical model. We can interpret linear types as objects of L: We can interpret ⊸ as the C-object ⊸ def = ⊸ . Finally, we can interpret terms Γ ⊢ : as morphisms in C( Γ , ) and terms Γ; v : ⊢ : as in L( Γ )( , ): Observe that we interpret 0 and + using the biproduct structure of L. P 5.6. The interpretation − of the language of §4 in categorical models is both sound and complete with respect to the +-equational theory: The proof is a minor variation of syntax-semantics correspondences developed in detail in chapters 3 and 5 of [41], where we use the well-known result that finite products in a category are biproducts iff the category is enriched over commutative monoids [17]. Soundness follows by case analysis on the +-rules. Completeness follows by the construction of the syntactic model LSyn : CSyn → Cat: • CSyn extends its full subcategory Syn with Cartesian ⊸-types; • Objects of LSyn( ) are linear types of our target language. • All type formers are interpreted as one expects based on their notation, using introduction and elimination rules for the required structural isomorphisms.

Sets and Commutative
Monoids. Throughout this paper, we have a particularly simple instance of the abstract semantics of our languages in mind, as we intend to interpret real as the usual Euclidean space R (considered as a set) and to interpret each program 1 : real 1 , . . . , : real ⊢ : real as a function R 1 × . . . × R → R . Similarly, we intend to interpret real as the commutative monoid R and each program 1 : real 1 , . . . , : real ; : real ⊢ Manuscript submitted to ACM : real as a function R 1 × . . . × R → R ⊸ R . That is, we will work with a concrete denotational semantics in terms of sets and commutative monoids. Some readers will immediately recognize that the free-forgetful adjunction Set ⇆ CMon gives a model of full intuitionistic linear logic [35]. In fact, seeing that CMon is CMon-enriched, the model is biadditive [17].
However, we do not need such a rich type system. For us, the following suffices. Define CMon( ), for ∈ ob Set, to have the objects of CMon and homsets CMon( )( , ) def = Set( , ⊸ ). Identities are defined as ↦ → ( ↦ → ) and composition ; CMon( ) is given by ↦ → ( ( ); CMon ( )). Given ∈ Set( , ′ ), we define change-of-base CMon( ′ ) → CMon( ) as CMon( )( ) def = ; Set . CMon(−) defines a locally indexed category. By taking C = Set and L(−) = CMon(−), we obtain a concrete instance of our abstract semantics. Indeed, we have natural isomorphisms Note that the derivatives are not merely linear in the sense of preserving 0 and +. They are also multiplicative in the sense that ( )( )( · ) = · ( )( )( ). We could have captured this property by working with vector spaces rather than commutative monoids. However, we will not need this property to phrase or establish correctness of AD.
Therefore, we restrict our attention to the more straightforward structure of commutative monoids.
Defining real def = R and interpreting each lop ∈ LOp as the (differentiable) function lop : it is intended to represent, we obtain a canonical interpretation of our target language in CMon.

Operational Semantics
In this section, we describe an operational semantics for our source and target languages. We consider call-by-value evaluation, but similar results can be obtained for call-by-name evaluation 10 . We present this semantics in big-step style.
Finally, we show that our denotational semantics are adequate with respect to this operational semantics, showing that the denotational semantics are sound tools for reasoning about our programs. 10 In fact, we conjecture our target language to be pure in the sense that reductions are confluent.
(Γ; v : real 1 * · · * real : real 1 * · · * real ) ⇓ (Γ; v : real 1 * · · * real ⊢ v : 7. The big-step call-by-value operational semantics ⇓ for the source and target languages. In the first rule, we intend to indicate that ⇓ unless is a linear function between tuples of real arrays in the sense that Γ; v : real 1 * · · * real ⊢ : real 1 * · · * real . We consider the following programs values, where we write for op() and (v) for lop(; v): , , ::= values . We then define the big-step reduction relation ⇓ , which says that a program evaluates to the value , in Fig. 7.
To be able to define this semantics, we assume that our languages contain, at least, nullary operations op = for all constants ∈ R and nullary linear operations for all linear maps (matrices) ∈ R ⊸ R . For all operations op and linear operations lop, we assume that an intended semantics op and lop is specified as (functions on) vectors of reals. As a side-note, we observe that this operational semantics has the following basic properties: Manuscript submitted to ACM Subject reduction and termination are proved through a standard logical relations argument similar to the ones in [39]. Determinism is observed by simply noting that all rules in the definition of ⇓ have conclusions ⇓ with disjoint .
In fact, seeing that every well-typed program has a unique value such that ⇓ , we write ⇓ for this .
We assume that only first-order types are observable (i.e.: have decidable equality on their values): , ::= first-order Cartesian types  we have that This is proved by induction on the structure of terms. This is proved by induction on the definition of ⇓: note that every operational rule is also an equation in the semantics.
Then, adequacy follows. That is, the denotational semantics is a sound means for proving observational equivalences of the operational semantics.

PAIRING PRIMALS WITH (CO)TANGENTS, CATEGORICALLY
In this section, we show that any categorical model L : C → Cat of our target language gives rise to two Cartesian closed categories Σ C L and Σ C L . We believe that these observations of Cartesian closure are novel. Surprisingly, they are highly relevant for obtaining a principled understanding of AD on a higher-order language: the former for forward AD, and the latter for reverse AD. Applying these constructions to the syntactic category LSyn : CSyn → Cat of our target language, we produce a canonical definition of the AD macros, as the canonical interpretation of the -calculus in the Cartesian closed categories Σ CSyn LSyn and Σ CSyn LSyn . In addition, when we apply this construction to the denotational semantics CMon : Set → Cat and invoke a categorical logical relations technique, known as subsconing, we find an elegant correctness proof of the source code transformations. The abstract construction delineated in this section is in many ways the theoretical crux of this paper.

Grothendieck Constructions on Strictly Indexed Categories
Recall that for any strictly indexed category, i.e. a (strict) functor L : C → Cat, we can consider its total category (or Grothendieck construction) Σ C L, which is a fibred category over C (see [23, ). Furthermore, given a strictly indexed category L : C → Cat, we can consider its fibrewise dual category L : C → Cat, which is defined as the composition C L − − → Cat − − → Cat.
Thus, we can apply the same construction to L to obtain a category Σ C L .
6.2 Structure of Σ C L and Σ C L for Locally Indexed Categories § §6.1 applies, in particular, to the locally indexed categories of §5. In this case, we will analyze the categorical structure of Σ C L and Σ C L . For reference, we first give a concrete description.

Manuscript submitted to ACM
Then, observing that the composition in these Σ-types of categories is precisely the chain rule, we see that the paired-up derivative Tand transposed derivative T * of § §2.2 define functors T : CartSp → Σ Set CMon T * : CartSp → Σ Set CMon .
As we will see in §7, we can implement (higher-order extensions of) these functors as code transformations As we will see, we can derive these code transformations by examining the categorical structure present in Σ C L and Σ C L for categorical models L : C → Cat of the target language in the sense of §5. We believe the existence of this categorical structure is a novel observation. We will make heavy use of it to define our AD algorithms and to prove them correct. • terminal object 1 = (1, 1) and binary products

P . We have (natural) bijections
We observe that we need L to have biproducts (equivalently: to be CMon enriched) in order to show Cartesian closure. Furthermore, we need linear ⇒-types and Cartesian ⊸-types to construct exponentials. Codually, we also obtain the Cartesian closure of Σ C L . However, for concreteness, we write out the proof by hand. • terminal object 1 = (1, 1) and binary products

P . We have (natural) bijections
Observe that we need the biproduct structure of L to construct finite products in Σ C L . Furthermore, we need Cartesian ⊸-types and !(−) ⊗ (−)-types, but not biproducts, to construct exponentials.
Interestingly, we observe that the exponentials in Σ C L and Σ C L are not fibred over C (unlike their products, for example). Indeed ( 1 , 2 ) ⇒ ( 1 , 2 ) has first component not equal to 1 ⇒ 1 . In the context of automatic differentiation, this has the consequence that primals associated with values of function type are not equal to itself. Instead, as we will see, they include both a copy of and a copy of its (transposed) derivative. These primals at higher-order types can be contrasted with the situation at first-order types, where values are equal to their associated primal, as a result of the finite products being fibred.

NOVEL AD ALGORITHMS AS SOURCE-CODE TRANSFORMATIONS
As Σ CSyn LSyn and Σ CSyn LSyn are both Cartesian closed categories by Theorems 6.1 and 6.2, the universal property of the source language (Pro. 3.1) gives us the following definition of forward and reverse mode CHAD as a canonical homomorphic functors. In this section, we discuss Manuscript submitted to ACM • the interpretation of the above functors as a type-respecting code transformation; • how we can define the basic definitions
Furthermore, all variables used in the source code transformations below are assumed to be freshly chosen. Here, we write Γ for the list of identifiers 1 , . . . , that occur in the typing context Γ = 1 : 1 , . . . , : . As we will see later, we need to know these context identifiers to define the code transformation. Equivalently, we can pair up the primal and (co)tangent computations as • a combined forward mode primal and tangent computation • a combined reverse mode primal and cotangent computation We prefer to work with these combined primal and (co)tangent code transformations as this allows us to share common subexpressions between the primal and (co)tangent computations using let-bindings. In writing down the definitions of CHAD on programs, we make sure to choose sensible representatives of these +-equivalence classes: ones that share common subexpressions through let-bindings. While these let-bindings naturally do not affect correctness of the transformation, they let us avoid code explosion at compile-time and unnecessary recomputation at run-time.
Finally, we note that, due to their definition from a universal property, our code transformations automatically respect equational reasoning in the sense that Γ ⊢ = : implies that  In §4, we have assumed that there are suitable terms (for example, linear operations) 1 : real 1 , · · · , : real ; v : real 1 * · · · * real ⊢ op( 1 , . . . , ; v) : real to represent the forward and reverse mode derivatives of the primitive operations op ∈ Op 1 ,..., . Using these, we For example, for elementwise multiplication ( * ) ∈ Op , , which we interpret as the usual elementwise product ( * ) def = ( * ) : R × R → R , we need, by the product rule for differentiation, that

The Implied Forward Mode CHAD Definitions
We define the types of (forward mode) primals − → D ( ) 1 and tangents − → D ( ) 2 associated with a type as follows: Observe that the type of primals associated with a function type is not equal to the original type. This is a consequence of the non-fibred nature of the exponentials in Σ-types Σ CSyn LSyn of categories ( §6).
For programs , we define their efficient CHAD transformation − → D Γ ( ) as follows: We explain and justify these transformations in the next subsection after discussing the transformations for reverse CHAD.

The Implied Reverse Mode CHAD Definitions
We define the types of (reverse mode) primals ← − D ( ) 1 and cotangents ← − D ( ) 2 associated with a type as follows: Again, we associate a non-trivial type of primals to function types as exponentials are not fibred in Σ CSyn LSyn ( §6).
For programs , we define their efficient CHAD transformation ← − D Γ ( ) as follows: We now explain and justify the forward and reverse CHAD transformations. The transformations for variables, tuples and projections implement the well-known multivariate calculus facts about (transposed) derivatives of differentiable functions into and out of products of spaces. The transformations for let-bindings add to that the chain rules for T and T * of § §2.2. The transformations for -abstractions split the derivative of a closure . into the derivative ′ with respect to the function argument and the derivative snd ( ) with respect to the captured context variables; they store ′ together with the primal computation of . in the primal associated with the closure and they store snd ( ) in the (co)tangent associated with the closure. Conversely, the transformations for evaluations extracts those two components of the (transposed) derivative ′ ctx (w.r.t. context variables) and ′ arg (w.r.t. function argument) from the (co)tangent and primal, respectively, and recombine them to correctly propagate (co)tangent contributions from both sources.

Sharing of Common Subexpressions
Through careful use of let-bindings, we have taken care to ensure that the CHAD code transformations we specified have the following good property: for every program former [ 1 , . . . , ] that takes subprograms 1 ,. . . , (for example, function application takes two subprograms and ), we have that  Of course, our use of let-bindings has the additional benefit at run-time that repeated subcomputations are performed only once and their stored results are shared, rather than recomputing every time their results are needed. We have taken care to avoid any unnecessary computation in this way, which we hope will benefit the performance of CHAD in practice. However, we leave a proper complexity and practical performance analysis to future work.

PROVING REVERSE AND FORWARD AD SEMANTICALLY CORRECT
In this section, we show that the CHAD code transformations described in §7 correctly compute mathematical deriva- This logical relations proof can be phrased in elementary terms, but the resulting argument is technical and would be hard to discover. Instead, we prefer to phrase it in terms of a categorical subsconing construction, a more abstract and elegant perspective on logical relations. We discovered the proof by taking this categorical perspective, and, while we have verified the elementary argument (see § §2.6), we would not otherwise have come up with it.

Subsconing.
Logical relations arguments provide a powerful proof technique for demonstrating properties of typed programs. The arguments proceed by induction on the structure of types. Here, we briefly review the basics of categorical logical relations arguments, or subsconing constructions. We restrict to the level of generality that we need here, but we would like to point out that the theory applies much more generally.
Consider a Cartesian closed category (C, 1, ×, ⇒). Suppose that we are given a functor : C → Set to the category Set of sets and functions that preserves finite products in the sense that (1) 1 and ( × ′ ) ( ) × ( ′ ).
Then, we can form the subscone of , or category of logical relations over , which is Cartesian closed, with a faithful Cartesian closed functor 1 to C that forgets about the predicates [24]: • objects are pairs ( , ) of an object of C and a predicate ⊆ ; • morphisms ( , ) → ( ′ , ′ ) are C morphisms : → ′ that respect the predicates in the sense that ( )( ) ⊆ ′ ; • identities and composition are as in C; • (1, 1) is the terminal object, and binary products and exponentials are given by In typical applications, C can be the syntactic category of a language (like Syn), the codomain of a denotational semantics − (like Set), or a product of the above, if we want to consider -ary logical relations. Typically, tends to be a hom-functor (which always preserves products), like C(1, −) or C( 0 , −), for some important object 0 . When applied to the syntactic category Syn and = Syn(1, −), the formulae for products and exponentials in the subscone clearly reproduce the usual recipes in traditional, syntactic logical relations arguments. As such, subsconing generalises standard logical relations methods.

Subsconing for Correctness of AD
We apply the subsconing construction above to where we note that Set, Σ Set CMon, and Σ Set CMon are Cartesian closed (given the arguments of §5 and §6) and that the product of Cartesian closed categories is again Cartesian closed. Let us write where we write for the semantic derivative of and (−) for the matrix transpose (see §5).  Most of the work is now in place to show correctness of AD. We finish the proof below. To ease notation, we work with terms in a context with a single type. Doing so is not a restriction as our language has products, and the theorem holds for arbitrary terms between first-order types. P . Seeing that our language has tuples, we may assume without loss of generality that Γ = : . We use the logical relations for general to show differentiability of all programs. Next, the case of = 1 suffices to show that CHAD computes correct derivatives.
First, we observe that sends differentiable functions R → to differentiable functions R → , as respects the logical relations. Observing that R for some , as is a first-order type, we can choose = . Then, contains ( , ( , ℎ)) for a differentiable isomorphism . It therefore follows that is differentiable.
Second, we focus on the correctness of forward AD, where we use the definition of composition in Set × Σ Set CMon. Therefore, and, by the chain rule, Evaluating the former at 0 gives ( ) = Manuscript submitted to ACM

PRACTICAL RELEVANCE AND IMPLEMENTATION IN FUNCTIONAL LANGUAGES
Most popular functional languages, such as Haskell and OCaml, do not natively support linear types. As such, the transformations described in this paper may seem hard to implement. However, as we will argue in this section, we can easily implement the limited linear types used in phrasing the transformations as abstract data types by using merely a basic module system, such as that of Haskell. The key idea is that linear function types ⊸ can be represented as plain functions → and copowers ! ⊗ can be represented as lists or arrays of pairs of type * .
To substantiate that claim, we provide a reference implementation of CHAD operating on strongly typed deeply embedded DSLs Haskell at https://github.com/VMatthijs/CHAD. This section explains how that implementation relates to the theoretical development in the rest of this paper. It is rather short because our implementation almost exactly follows the theoretical development in §3, 4, 5 and 7.

Implementing Linear Functions and Copowers as Abstract Types in Functional Languages
Based on the denotational semantics, ⊸ -types should hold (representations of) functions from to that are homomorphisms of the monoid structures on and . We will see that these types can be implemented using an abstract data type that holds certain basic linear functions (extensible as the library evolves) and is closed under the identity, composition, argument swapping, and currying. Again, based on the semantics, ! ⊗ should contain (representations of) finite maps (associative arrays) =1 ! ⊗ of pairs ( , ), where is of type , and is of type , and where we identify + ! ⊗ + ! ⊗ ′ and + ! ⊗ ( + ′ ).
To implement this idea, we consider abstract types LFun( , ) of linear functions and Copower( , ) of copowers.
Their programs are generated by the following grammar We note that these abstract types give us precisely the functionality and type safety of the linear function and copower types of our target language of §4. Indeed, we can define a semantics and type preserving translation (−) from that target language to our source language extended with these LFun( , ) and Copower( , ,) types, for which 9.2 Implementing the API of LFun( , ) and Copower( , ,) Types We observe that we can implement this API of LFun( , ) and Copower( , ) types, as follows, in a language that extends the source language with types List( ) of lists (or arrays) of elements of type . Indeed, we implement LFun( , ) under the hood, for example, as → and Copower( , ) as List( * ). The idea is that LFun( , ), which arose as a right adjoint in our linear language, is essentially a subtype of → . On the other hand, Copower( , ), which arose as a left adjoint, is a quotient type of List( * ). We achieve the desired subtyping and quotient typing by exposing only the API of Fig. 8 and hiding the implementation. We can then implement this interface as follows. 11 Here, we write lop for the function that lop is intended to implement, [ ] for the empty list, :: for the list consisting of with prepended on the front, and fold over in from = for (right) folding an operation over a list , starting from . The monoid structure 0 , + can be defined by induction on the structure of types, by using, for example, type classes in Haskell or OCaml's module system or (in any functional language with algebraic data types) by using reified types. Furthermore, the implementer of the AD library can determine which linear operations lop to include within the implementation of LFun. We expect these linear operations to include various forms of dense and sparse matrix-vector multiplication as well as code for computing Jacobian-vector and Jacobian-adjoint products for the operations op that avoids having to compute the full Jacobian. Another option is to simply include two linear operations op and op for computing the derivative and transposed derivative of each operation op.

Maintaining Type Safety throughout the Compilation Pipe-Line in our Reference Implementation
In a principled approach to building a define-then-run AD library, one would shield this implementation using the abstract data types Copower( , ) and LFun( , ) as we describe, both for reasons of type safety and because it conveys the intuition behind the algorithm and its correctness. By using such abstract data types in our Haskell implementation combined with GADTs and type families, we achieve a fully type-safe (well-scoped, well-typed De Bruijn) implementation of the source and target languages of §3 and §4 with their semantics of §5 and statically type-checked code transformations of §7.
However, nothing stops library implementers from exposing the full implementation rather than working with abstract types. In fact, this seems to be the approach [43] have taken. A downside of that "exposed" approach is that the transformations then no longer respect equational reasoning principles. In our reference implementation, we include a compiler from the (linearly typed) target language to a less type-safe "concrete" target language (implementing § §9.2 as a compilation step): essentially the source language extended with list (or array) types. 12 This demonstrates definitely that CHAD implements a notion of compile-time AD code transformation that takes in and spits out standard functional code without any custom semantics.

Compiling Away Copowers
As a final observation on this implementation, we would like to note that while the proposed implementation of copowers as lists is generally applicable, more efficient implementation strategies can often be achieved in practice. In fact, in unpublished follow-up work to this paper led by Tom Smeding, we show that when we implement CHAD to operate on Accelerate [34], we can optimize away uses of copower types.

ADDING HIGHER-ORDER ARRAY PRIMITIVES
The aim of this paper is to answer the foundational question of how to perform (reverse) AD at higher types. The problem of how to perform AD of evaluation and currying is highly challenging. For this reason, we have devoted this paper to explaining a solution to that problem in detail, working with a toy language with ground types of black-box, sized arrays real with some first-order operations op. However, many of the interesting applications only arise once we can use higher-order array primitives such as map and fold on real .
Our definitions and correctness proofs extend to this setting with standard array processing primitives including map, fold, filter, zipWith, permute (aka scatter), ba permute (aka gather), generate (aka build), and array indexing. We plan to discuss these primitives as well as CHAD applied to dynamically sized arrays in detail in an applied followup paper, which will focus on an implementation of CHAD to operate on Accelerate [34].
To illustrate the idea behind such an extension, we briefly discuss the case for map here and leave the rest to future work. Suppose that we add operations 13 In our practical API of § §9.1, the required target language primitives correspond to Γ ⊢ sum : 13 We prefer to work with this elementary formulation of maps rather than the usual higher-order formulation of Γ ⊢ : real → real Γ ⊢ : real Γ ⊢ map( , ) : real because it makes sense in the wider context of languages without function types as well and because it simplifies its CHAD correctness proof. Note that both are equivalent in presence of function types: map( , ) = map( . , ) and map( . , ) = map( . , ).
are the usual functional programming idiom for mapping a unary function over an array, zipping two arrays with a binary operation and for taking the sum of the elements in an array. Note that we assume that we have types for length-arrays of elements of type , here, generalizing the arrays real of elements of type real. We present a correctness proof for this implementation of the derivatives in Appx. A.
Applications frequently require AD of higher-order primitives such as differential and algebraic equation solvers, e.g. for use in pharmacological modelling in Stan [40]. Currently, derivatives of such primitives are derived using the calculus of variations (and implemented with define-by-run AD) [7,18]. Our proof method provides a more lightweight and formal method for calculating, and establishing the correctness of, derivatives for such higher-order primitives.
Indeed, most formalizations of the calculus of variations use infinite-dimensional vector spaces and are technically involved [26].

Memory Use of CHAD's Forward AD
CHAD has formulated reverse and forward AD to be precisely each other's categorical dual. The former first computes the primals in a forward pass and next the cotangents in a reverse pass. Dually, the latter first computes the primals in a forward pass and next the tangents in another forward pass. Seeing that the two forward passes in forward AD have identical control-flow, it can be advantageous to interleave them and simultaneously compute the primals and tangents. Such interleaving greatly reduces the memory consumption of the algorithm, as it means that all primals no longer need to be stored for most of the algorithm. We present such an interleaved formulation of forward AD in [21].
While it is much more memory efficient, it has the conceptual downside of no longer being the mirror image of reverse AD. Furthermore, these interleaved formulations of forward AD work by operating on dual numbers. That is, they use an array-of-structs representation, by contrast with the struct-of-arrays representation used to pair primals with tangents in CHAD. Therefore, an SoA-to-AoS optimization is typically needed to make interleaved implementations of forward AD efficient [38].
Finally, we note that such interleaving techniques do not apply to reverse AD, as we need to have completed the forward primal pass before we can start the reverse cotangent pass, due to the dependency structure of the algorithm.

Applying CHAD to Richer Source Languages
The core observations that let us use CHAD for AD on a higher-order language were the following: (1) there is a class of categories with structure S (in this case, Cartesian closure) such that the source language Syn that we want to perform AD on can be seen as the freely generated S-category on the operations op; (2) we identified structure T that suffices for a CMon-enriched strictly indexed category L : C → Cat to ensure that Σ C L and Σ C L are S-categories; CHAD applies equally to source languages with other choices of S, provided that we can follow steps 1-4.
In particular, in [30] it is shown how CHAD applies equally to languages with sum types and (co)inductive types (and tuple and function types). In that setting, the category LSyn is a genuine strictly indexed category over CSyn, to account for the fact that the (co)tangent space to a space of varying dimension depends on the chosen base point.
That is, in its most principled formulation, the target language has (linear) dependent types. However, we can also work with a simply typed target language in this setting, at the cost of some extra type safety. In fact, our Haskell implementation already supports such a treatment of coproducts.
As discussed in §10, S can also be chosen to include various operations for manipulating arrays, such as map, fold, filter, zipWith, permute (aka scatter), ba permute (aka gather), generate (aka build), and array indexing. We plan to describe this application of CHAD to array processing languages and we are implementing CHAD to operate on the Accelerate parallel array processing language.
In work in progress, we are applying CHAD to partial features such as real conditionals, iteration, recursion and recursive types. Our Haskell implementation of CHAD already supports real conditionals, iteration and recursion.
The challenge in this setting is to understand the subtle interactions between the CPO-structure needed to model recursion and the commutative monoid structure that CHAD uses to accumulate (co)tangents.

CHAD for Other Dynamic Program Analyses
As noted by [43], source-code transformation AD has a lot of similarities with other dynamic program analyses such as dynamic symbolic analysis and provenance analysis.
In fact, as the abstract perspective on CHAD given in § §11.2 makes clear, CHAD is in no way tied to automatic differentiation. In many ways, it is much more general, and can best be seen as a framework for applying dynamic program analyses that accumulate data (either by going through the program forwards or backwards) in a commutative monoid to functional languages with expressive features. In fact, by varying the definitions of To give one more example application of such an analysis, beyond AD, dynamic symbolic analysis and provenance analysis, one can note that for a source language Syn, generated from a base type R that is a commutative semi-ring, we have a notion of algebraic (or formal) derivative of any polynomial 1 : R, . . . , : R ⊢ op( 1 , . . . , ) : R [28].
CHAD can be used to extend and compute this notion of derivative for arbitrary functional programs generated from the polynomials op as basic operations. The particular case of such formal derivatives for the Boolean semi-ring R = B are used in [45] to feed into a gradient descent algorithm to learn Boolean circuits. CHAD makes this method applicable to more general (higher-order) programs over (arrays of) Booleans.

RELATED WORK
This work is closely related to [21] and [20], which introduced a similar semantic correctness proof for a dual-numbers version of forward mode AD and higher-order forward AD, using a subsconing construction. A major difference is that this paper also phrases and proves correctness of reverse mode AD on a -calculus and relates reverse mode to forward mode AD. Using a syntactic logical relations proof instead, [4] also proves correctness of forward mode AD.
Again, it does not address reverse AD.
Manuscript submitted to ACM [12] proposes a similar construction to that of §6, and it relates it to the differential -calculus. This paper develops sophisticated axiomatics for semantic reverse differentiation. However, it neither relates the semantics to a sourcecode transformation, nor discusses differentiation of higher-order functions. Our construction of differentiation with a (biadditive) linear target language might remind the reader of differential linear logic [15]. In differential linear logic, (forward) differentiation is a first-class operation in a (biadditive) linear language. By contrast, in our treatment, differentiation is a meta-operation.
Importantly, [16] describes and implements what are essentially our source-code transformations, though they were restricted to first-order functions and scalars. After completing this work, we realized that [43] describes an extension of the reverse mode transformation to higher-order functions in a similar manner as we propose in this paper, but without the linear or abstract types. Though that paper did not derive the algorithm or show its correctness, it does discuss important practical considerations for its implementation and offers a dependently typed variant of the algorithm based on typed closure conversion, inspired by [37].
Next, there are various lines of work relating to correctness of reverse mode AD that we consider less similar to our work. For example, [31] define and prove correct a formulation of reverse mode AD on a higher-order language that depends on a non-standard operational semantics, essentially a form of symbolic execution. [2] does something similar for reverse mode AD on a first-order language extended with conditionals and iteration. [8] defines a beautifully simple AD algorithm on a simply typed -calculus with linear negation (essentially, a more finely typed version of the continuation-based AD of [21]) and proves it correct using operational techniques. [33] extends this work to apply to recursion. Furthermore, they show with an impressive operational argument that this simple algorithm, surprisingly, corresponds to true reverse mode AD with the correct complexity under an operational semantics with a "linear factoring rule". While this is a natural operational semantics for a linear -calculus, it is fundamentally different from normal call-by-value or call-by-name evaluation (under which the generated code has the wrong computational complexity). For this reason, we require a custom interpreter or compiler to use this reverse AD method in practice.
Very recently, [25] specified another purely functional reverse AD algorithm, which appears similar to, and, which we conjecture to be equivalent to an implementation of the techniques of [33]. These formulations of reverse mode AD all depend on non-standard run-times and hence fall into the category of "define-by-run" formulations of reverse mode AD, for our purposes. Meanwhile, we are concerned with "define-then-run" formulations: source-code transformations producing differentiated code at compile-time that can then be optimized during compilation with existing compiler tool-chains (such as the Accelerate [11], Futhark [19] and TensorFlow [1] frameworks for generating highly performant GPU code). While we can compile such define-by-run transformations together with their interpreter to achieve a source-code transformation (hence a sort of define-then-run transformation), the resulting code recursively traverses an AST, so does not obviously seem suitable for generating, via existing tool-chains, optimized machine code for the usual parallel hardwares that we use as targets for AD, such as GPUs and TPUs.
Finally, there is a long history of work on reverse mode AD, though almost none of it applies the technique to higherorder functions. A notable exception is [37], which gives an impressive source-code transformation implementation of reverse AD in Scheme. While very efficient, this implementation crucially uses mutation. Moreover, the transformation is complex and correctness is not considered. More recently, [44] describes a much simpler implementation of a reverse AD code transformation, again very performant. However, the transformation is quite different from the one considered in this paper as it relies on a combination of delimited continuations and mutable state. Correctness is not considered, perhaps because of the semantic complexities introduced by impurity.
Our work adds to the existing literature by presenting a novel, generally applicable method for compositional sourcecode transformation (forward and) reverse AD on expressive functional languages without a need for a non-standard runtime, by giving a method for compositional correctness proofs of such AD algorithms, and by observing that the CHAD method and its correctness proof are not limited to AD but apply generally to dynamic program analyses that accumulate data in a commutative monoid.

A CHAD CORRECTNESS FOR HIGHER-ORDER OPERATIONS SUCH AS MAP
We extend the proofs of Lemmas 2.1 and 2.2 to apply to the map-constructs of §10.

A.2 Extending the Induction Proof of the Fundamental Lemma for Forward CHAD
First, we focus on extending the induction proof of the fundamental lemma for forward CHAD to apply to maps.
Assume the induction hypothesis that and respect the logical relation. We show that map( . , ) does as well.
Seeing that ( , ( , ℎ)) ∈ Γ by assumption, it is enough, by definition of Γ, :real , to show that By definition of real , it is enough to show that Seeing that respects the logical relation by our induction hypothesis, it is enough to show that which is true by assumption.

A.3 Extending the Induction Proof of the Fundamental Lemma for Reverse CHAD
Next, we extend the fundamental lemma for reverse CHAD to apply to maps. Assume the induction hypothesis that and respect the logical relation. We show that map( . , ) does as well. (Where we assume all terms all well-typed.) Suppose that ( , ( , ℎ)) ∈ Γ . We want to show that D Γ (map( . , )) ( ( )))( )))) ∈ real .

B TERM SIMPLIFICATIONS IN THE IMPLEMENTATION
Our implementation 14 of the AD macros that is described in §9 includes a number of simplification rules on the concrete target language whose only purpose is to make the produced code more readable and clear (without changing its asymptotic runtime cost). The motivation for these rules is to be able to generate legible code when applying the AD macros to example programs. In this appendix, we list these simplification rules explicitly and show the implementation's output on the four example programs in Figs. 1 and 2 under these simplification rules. We do this to illustrate that: (1) the simplifications given here are evidently meaning-preserving, given the + rules in Figs. 4 and 6, and are standard rules that any optimizing compiler would apply; (2) the thus simplified output of the AD macros from §7 is indeed equivalent to the differentiated programs in Figs. 1 and 2.
The simplification rules in question are given below in Table 1 let substitution (if is cheap or used at most once in e) pairProj 1 fst , pair pairProj 2 snd , pair pairEta fst , snd pair letProj 1 fst (let = in ) let = in fst let substitution letProj 2 snd (let = in ) let = in snd let substitution plusZero 1 plus zero equational rule plusZero 2 plus zero equational rule plusPair plus , , plus , plus equational rule plusLet 1 plus (let = in ) let = in plus let substitution plusLet 2 plus (let = in ) let = in plus let substitution algebra 0 * , 0 * 0 (etc.) basic algebra letLamPairSplit let = . , in let 1 = . in let 2 = . The last column in the table shows the justification for the simplification rule: "let substitution", " pair", " pair", "lambda substitution" and " lambda" refer to the corresponding rules in Fig. 4. The equational rules are either from Fig. 6 in the case of + 0 = and its symmetric variant, from the type-based translation rules in § §9.2 in the case of plus on pairs, or otherwise general laws that hold for zero (i.e. 0) and/or the array combinators in question.
Note that all of the rules preserve time complexity of the program through careful sharing of values with let-bindings.
These let-bindings could only be argued to increase work if the value is only used once in the body of the let -but in that case, the 'letInline' rule will eliminate the let-binding anyway.

B.1 First-order example programs
The output of our implementation for the forward derivative of Fig. 1 (a) and the reverse derivative of Fig. 1 (b) can be found below in Fig. 9.
l e t two = 2 . 0 y = two * x z = x * y in y , z , c o s z , x ' . l e t y ' = two * snd x ' z ' = p l u s ( x * y ' ) ( y * snd x ' ) in y ' , z ' , ( 0 . 0 − s i n z ) * z ' (a) The implementation's forward AD output on Fig. 1 (a). Compare this with Fig. 1 (c). l e t two = 2 . 0 y = x1 * x4 + two * x2 w = y * x3 + x4 in s i n w , v ' . l e t w ' = c o s w * v ' y ' = x3 * w' in z e r o , x4 * y ' , two * y ' , y * w' , p l u s w ' ( x1 * y ' ) (b) The implementation's reverse AD output on Fig. 1 (b). Compare this with Fig. 1 (d). The simplification rules listed above have already been applied (otherwise the output would, indeed, be much less readable). The only change we made to the literal text output of the implementation is formatting and renaming of variables.
For both programs, we note that in the implementation, environments are encoded using snoc-lists: that is, the environment Γ = 1 : real, 2 : real, 3 : real is represented as (( , 1 : real), 2 : real), 3 : real. Hence, for this Γ, − → D (Γ) 2 would be represented as ((1 * real) * real) * real. This has an effect on the code in Fig. 9, where in subfigure (a), the variable x ' has type 1 * real instead of real, which it had in Fig. 1 (c). Furthermore, the output in Fig. 9 (b), which is the cotangent (i.e. adjoint) of the environment, has type (((1 * real) * real) * real) * real, meaning that the 'zero' term in the result has type 1 and is thus equal to .
We hope that it is evident to the reader that these outputs are equivalent to the programs given in Fig. 1 (c) and (d).
B.2 Second-order example program Fig. 2 (a) The implementation's forward derivative of Fig. 2 (a) is shown below in Fig. 10. This version contains a let-bound function 'g' which does not occur in the code of Fig. 2 (c). However, inlining this function in the two places where it l e t f ' = z . d . p l u s ( x * snd d ) ( z * snd ( f s t d ) ) z s = v r e p l i c a t e x g = z . l e t d f u n = f ' z in ( x * z + 1 . 0 , denv . p l u s ( snd ( f s t ( f s t denv ) ) z ) ( d f u n ( z e r o , snd denv ) ) ) in ( vmap ( z . f s t ( g z ) ) z s , x ' . l e t zs ' = v r e p l i c a t e ( snd x ' ) in p l u s ( v z i pW i t h ( z . z ' . snd ( g z ) ( z e r o , z ' ) ) z s zs ' ) ( vmap ( z . snd ( g z ) ( ( ( x ' , z . f ' z ( x ' , z e r o ) ) , zs ' ) , z e r o ) ) z s ) ) is used does not increase work, because due to pair projection and equational rules concerning 'zero' and 'plus', only one half of the 'plus' expression in 'g' remains at both of the invocation sites of 'g'. A version with 'g' manually inlined and simplified using the stated rules is shown in Fig. 11. (Our automatic simplifier is currently not smart enough to prove that inlining 'g' does not increase work, and hence keeps it let-bound.) First, we notice that the type of the variable x ' here is 1 * real instead of real because of the snoc-list representation of environments, as was discussed in the previous subsection. Because of this, the expression snd x' in Fig. 11 is equivalent to the expression x' in Fig. 2 (c). Knowing this, let us work out how the implementation got this output code, and how it is equivalent to the code in Fig. 2 (c).
For the purposes of this explanation, the most important component of the source of Fig. 2 (a) is its first line: 'let f = z. x * z + 1 in ... '. Recall from the forward AD macro from § §7.4: Hence, the binding of is transformed to a binding of − → D Γ ( ), which is then used somehow in the body. Since is the term ' z. x * z + 1' here, and since the macro rule for lambda abstraction is as follows: This is the term that appears on the right-hand side of a let-binding in the forward AD transformed version of the code from Fig. 2 (a). Inlining of and some further simplification yields: wherein we recognize f and f ' from Fig. 2 (c).
The implementation proceeded differently in the simplification of − → D , ( . * + 1), namely by splitting the lambda bound to using 'letLamPairSplit'. The second of the resulting two lambda functions is .
After seeing this, we leave it to the reader to show that the code in Fig. 11 and Fig. 2 (c) is equivalent under forward and reverse let-substitution.

(b)
When the implementation performs reverse AD on the code in Fig. 2 (b) and simplifies the result using the simplification rules in Table 1, the result is the code that is shown below in Fig. 12.   Fig. 2 (b). The implementation writes operations on scalar arrays (real ) with a 'v' prefix; list operations are wri en without a prefix. Contrast this to Fig. 2, where all arrays are implemented as lists and hence there are no pure-array operations which would have been wri en using a 'v' prefix.
Firstly, note the evalOp EScalProd. Since scalar multiplication is implemented as an operation (op), it does not take two separate arguments but rather a pair of the two scalars to multiply. Due to application of the 'pairEta' simplification rule from Table 1, the argument to the multiplication was reduced from fst p, snd p to just 'p', preventing the prettyprinter from showing fst p * snd p; instead, this becomes a bare operation application to the argument 'p'. In the code in Fig. 2 (d), this lambda is the argument to the 'map' in the cotangents block, meaning that 'p' stands for x2i, y' .
In this example, a copower structure is created because of the usage of a function abstraction in the code to be differentiated using reverse AD. Here, this copower was interpreted using lists as described in § §9.2. The 'toList' function converts an array of scalars (i.e. a value of type real ) to a list of scalars. The code in Fig. 2 (d) goes further and interprets all arrays as lists, effectively removing the distinction between arrays originating from arrays in the source program and lists originating from copower values. Together with recognizing the snoc-list representation of environments, it becomes sufficient to inline some let-bindings in Fig. 2 (d) to arrive at code equivalent to the code shown here.