# Denotationally Correct, Purely Functional, Efficient Reverse-mode Automatic
Differentiation

Mathieu Huot<EMAIL_ADDRESS>University of OxfordOxfordUK and  Amir
Shaikhha  University of EdinburghEdinburghUK

###### Abstract.

Reverse-mode differentiation is used for optimization, but it introduces
references, which break the purity of the underlying programs, making them
notoriously harder to optimize. We present a reverse-mode differentiation on a
purely functional language with array operations. It is the first one to
deliver a provably efficient, purely functional, and denotationally correct
reverse-mode differentiation. We show that our transformation is semantically
correct and verifies the cheap gradient principle. Inspired by PROPs and
compilation to categories, we introduce a novel intermediate representation
that we call ‘unary form’. Our reverse-mode transformation is factored as a
compilation scheme through this intermediate representation. We obtain
provably efficient gradients by performing general partial evaluation
optimizations after our reverse-mode transformation, as opposed to manually
derived ones. For simple first-order programs, the obtained output programs
resemble static-single-assignment (SSA) code. We emphasize the modularity of
our approach and show how our language can easily be enriched with more
optimized primitives, as required for some speed-ups in practice.

categorical semantics, automatic differentiation, functional programming

††copyright: none

##  1\. Introduction

Deep learning is moving towards increasingly sophisticated optimization
objectives that employ tensors and operations on tensors. Reverse-mode
Automatic Differentiation (AD) is a technique to automatically compute the
gradient of objective functions of the form $\mathbb{R}^{n}\to\mathbb{R}$.
Such functions appear a lot in practice: for instance, as loss functions in
machine learning.

In order to reach the efficiency of the usual imperative version of reverse-
mode, the transformations, even in functional languages (Pearlmutter and
Siskind, 2008), resort to references. The lack of purity in reverse-mode is
known to make it significantly harder to optimize and parallelize. None of the
current implementations of reverse-mode in a functional setting (Wang et al.,
2019; Pearlmutter and Siskind, 2008; Baydin et al., 2016) are pure, and often
complicated heuristics with no guarantees are used, e.g. (Team, 2017). As a
result, to optimize for efficiency, a hand-crafted reverse-derivative must be
given, for every important non-trivial operation. In other words, abstracting
away from imperative code is still a hurdle that functional implementations
need to overcome.

In this paper, we define a purely functional (without references or control
mechanisms such as state monads), denotationally correct, and provably
efficient reverse-mode AD. To do so, we define the Unary Normal Form (UNF)
representation inspired by PROPs (MacLane, 1965) and compilation to categories
(Elliott, 2017). We can easily define and prove correctness of reverse-mode on
this representation. The whole reverse-mode transformation is obtained by
compiling the language to this Intermediate Representation (IR), applying the
simpler reverse-mode transformation, and compiling again to the original
language. After standard optimizations, the output program looks like SSA
(Cytron et al., 1989) or ANF (Sabry and Felleisen, 1993), which leads to more
efficient implementations.

|  This Paper |  (Wang et al., 2019) |  (Shaikhha et al., 2019) |  (Huot et al., 2020) |  (Brunel et al., 2019) |  (Abadi and Plotkin, 2020) |  (Barthe et al., 2020) |  (Pearlmutter and Siskind, 2008) |  (Elliott, 2018) |  (Sherman et al., 2021) |  (Vytiniotis et al., 2019) |  (Mak and Ong, 2020) |  (Vákár, 2021) |  (Manzyuk, 2012) |  (Cockett et al., 2019) |  (Cruttwell et al., 2019) |  (Krawiec et al., 2022) |  (Paszke et al., 2021b)
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---
Reverse Mode | ● | ● | ○ | ● | ● | ● | ○ | ● | ● | ○ | ● | ● | ● | ○ | ● | ○ | ● | ●
Complexity | ● | ◐ | ○ | ○ | ◐ | ◐ | ○ | ● | ○ | ○ | ◐ | ○ | ○ | ○ | ○ | ○ | ● | ●
Pure Derivatives | ● | ○ | ● | ● | ● | ● | ● | ○ | ● | ● | ● | ● | ● | ● | ● | ● | ● | ●
Correctness | ● | ○ | ○ | ● | ● | ● | ● | ○ | ● | ● | ○ | ● | ● | ● | ● | ● | ● | ○
Tensor Support | ● | ● | ● | ○ | ○ | ○ | ○ | ○ | ○ | ○ | ● | ○ | ● | ○ | ○ | ○ | ○ | ●
HO Functions | ○ | ● | ● | ● | ● | ○ | ● | ● | ○ | ● | ● | ● | ● | ● | ○ | ● | ● | ○
Recursion | ◐ | ● | ◐ | ◐ | ○ | ● | ○ | ● | ○ | ◐ | ○ | ○ | ○ | ○ | ○ | ○ | ○ | ○
Conditional | ● | ● | ● | ● | ○ | ● | ● | ● | ○ | ● | ○ | ○ | ○ | ○ | ○ | ○ | ○ | ○
Table 1. Comparison of different functional differentiable programming
frameworks. $\CIRCLE$ means that the property is verified, and $\Circle$ means
that it is absent in the work. $\LEFTcircle$ for complexity means that the
proof is not fully covered, and for recursion, that it does not support
general recursion but map, reduce and/or fold. HO stands for higher-order.
Correctness is ticked if a proof is formalized in the paper.

###  1.1. Examples

We introduce the general idea of efficient reverse-mode in a functional
setting through the following examples.

###### Example 1.1 (First-order term).

Let us consider the term let w${}_{1}:\mathbf{R}$ = x1 * x2 in let
w${}_{2}:\mathbf{R}$ = w1 * x1 in w2 in the context
$\Gamma:=\\{x_{1}:\mathbf{R},x_{2}:\mathbf{R},x_{3}:\mathbf{R}\\}$.
After an (inefficient) reverse-mode transformation, we obtain:

let w${}_{1}:\mathbf{R}$,w${}_{1}^{\prime}:\mathbf{R}^{4}\to\mathbf{R}^{3}$ =
$\langle$ x1 * x2, fun (y1,$\ldots$, y4) $\rightarrow\;$ (y1+x2*y4, y2+x1*y4,
y3)$\rangle$ in
---
let w${}_{2}:\mathbf{R}$,w${}_{2}^{\prime}:\mathbf{R}^{5}\to\mathbf{R}^{3}$ =
$\langle$ w1*x1, fun (y1,$\ldots$, y5) $\rightarrow\;$
w${}_{1}^{\prime}$(y1+w1*y5, y2, y3, y4+x1*y5)$\rangle$
in w${}_{2}^{\prime}$(0,0,0,0,1)

The part (0,0,0,0,1) corresponds to initializing the tangent variables in the
imperative reverse-mode algorithm. After some general partial evaluation
techniques that will be detailed further in the paper, we obtain:

let w${}_{1}:\mathbf{R}$ = x1 * x2 in
---
let w${}_{2}:\mathbf{R}$ = w1 * x1 in
let
y${}_{1}:\mathbf{R}$,y${}_{2}:\mathbf{R}$,y${}_{3}:\mathbf{R}$,y${}_{4}:\mathbf{R}$,y${}_{5}:\mathbf{R}$
= 0,0,0,0,1 in
let y${}_{1}^{\prime}:\mathbf{R}$ = y1+w1*y5 in
let y${}_{4}^{\prime}:\mathbf{R}$ = y4+x1*y5 in
(y${}_{1}^{\prime}$+x2*y${}_{4}^{\prime}$, y2+x1*y${}_{4}^{\prime}$, y3)

This is very close to the SSA form (Cytron et al., 1989) of what the
imperative reverse-mode differentiation of our initial term would be.

This term can be further optimized via constant propagation and algebraic
simplifications to give

let w${}_{1}:\mathbf{R}$ = x1 * x2 in
---
let w${}_{2}:\mathbf{R}$ = w1 * x1 in
(w1+x2*x1, x1*x1, 0)

###### Example 1.2 (Simple operations on arrays).

On arrays, three simple operations of interest are the dot-product of two
vectors, and the product or sum of the elements of a vector. In a functional
setting, these can be defined as follows:

prod(A:$\mathbf{A}[\mathbf{R}]^{n}$) | := | reduce * 1 A
---|---|---
sum(A:$\mathbf{A}[\mathbf{R}]^{n}$) | := | reduce + 0 A
dot(A:$\mathbf{A}[\mathbf{R}]^{n}$,B:$\mathbf{A}[\mathbf{R}]^{n}$) | := | reduce + 0 (map2 * A B)

where reduce is a known fold-left operator for which the function argument is
associative. It is notably faster to execute than a fold-left, as it is
parallel-friendly.

The gradient of each of these expressions with respect to A is:

$\nabla_{A}$prod(A) | := | map2 * (scanr * 1 (shift1L A)) (shift1R (scanl * 1 A))
---|---|---
$\nabla_{A}$sum(A) | := | map (x $\rightarrow\;$ 1) A
$\nabla_{A}$dot(A,B) | := | B

where

  * •

scanl is the scan-left operator that returns all the intermediate results of
fold-left,

  * •

scanr is the scan-right operator that returns all the intermediate results of
fold-right,

  * •

shift1L [v1,$\ldots$,vn] is the shift-left operator and returns
[v2,$\ldots$,vn], and

  * •

shift1R [v1,$\ldots$,vn] is the shift-right operator and returns
[v1,$\ldots$,vn-1].

These gradients are a few examples among numerous ones which are usually
derived by hand, and are here obtained automatically as special cases of our
work.

${\mathbf{Source}}$${\mathbf{Target}}$${Source\mathbf{UNF}}$${Target\mathbf{UNF}}$$\scriptstyle{\text{efficient
}{\scalebox{0.8}{$\scriptstyle\overleftarrow{\mathcal{D}}$}}(Fig.~{}\ref{fig:direct_diff_macro})}$$\scriptstyle{(Fig.~{}\ref{fig:source_to_unf})}$$\scriptstyle{\text{optim
}(Fig.~{}\ref{fig:optim})}$$\scriptstyle{{\scalebox{0.8}{$\scriptstyle\overleftarrow{\mathcal{D}}$}}(Fig.~{}\ref{fig:diff_macro})}$$\scriptstyle{(Fig.~{}\ref{fig:unf_to_target})}$
Figure 1. Outline of the compilation scheme.

###  1.2. Contributions

We propose a source-code transformation on a simple purely functional language
for purely functional reverse-mode differentiation. Our transformation
consists of a compilation scheme that is outlined in Figure 1. We make the
following contributions:

  * •

We present our work with a simple yet expressive array-based language (with
constructs such as map2 and reduce) in Section 3. We show how to directly
compute an efficient reverse-mode AD for the expressions of this program (top
of Figure 1). Furthermore, we show how to extend our work to a richer language
in Section 7.

  * •

One of the key insights behind efficient reverse-mode AD is to only consider
unary operators. Inspired by this insight and following Intermediate
Representations (IR) such as SSA and ANF, we introduce a novel IR, which we
call UNF (Section 4). We introduce an alternative and easier-to-follow
compilation pipeline for efficient reverse-mode AD (bottom of Figure 1).

  * •

We prove complexity guarantees for the programs transformed under reverse-mode
AD. Furthermore, we show a list of optimizations that can further improve the
constant factors (Section 5).

  * •

We prove the correctness of our transformations (top/bottom parts of Figure 1)
by defining a denotational semantics of our languages using multicategories
and concategories (Section 6).

Next, we recall rudiments of automatic differentiation, forward and reverse-
mode differentiation.

##  2\. Reverse-mode Automatic Differentiation

###  2.1. Rudiments of AD and dual numbers

To find the gradient $\nabla f$ of a function $f:\mathbb{R}\to\mathbb{R}$
compositionally, need both $f$ and $\nabla f$ when calculating $\nabla(f;g)$.
This is the reason why we are more generally interested in transforming a
function $f:\mathbb{R}^{n}\to\mathbb{R}$ into a function
$g:(\mathbb{R}\times\mathbb{R})^{n}\to\mathbb{R}\times\mathbb{R}$ in such a
way that for every $f_{1},\dots,f_{n}:\mathbb{R}\to\mathbb{R}$,

$(f_{1},\nabla f_{1},\dots,f_{n},\nabla
f_{n});g=((f_{1},\dots,f_{n});f,\nabla((f_{1},\dots,f_{n});f))$.

The idea of AD is to systematically transform a differentiable function
$f:\mathbb{R}^{n}\to\mathbb{R}$ into a function
$g:\mathbb{R}^{2n}\to\mathbb{R}^{2}$ which captures $f$ and all its partial
derivatives. An intuition for $g$ is often given in terms of dual numbers. The
transformed function operates on pairs of numbers, $(x,x^{\prime})$, and it is
common to think of such a pair as $x+x^{\prime}\epsilon$ for an
‘infinitesimal’ $\epsilon$. The main two ways in which AD is performed in
practice is by operator overloading or by source code transformation (see e.g.
(Griewank and Walther, 2008) Chapter 6). Our approach focuses on a source code
transformation, which is better fitted for compilation and optimizations.

###  2.2. Reverse-mode Automatic Differentiation

A potential computational problem shows up when one wants to compute the full
gradient of a function $\mathbb{R}^{n}\to\mathbb{R}$, for a large $n$.
Forward-mode only computes one directional derivative, for instance one
partial derivative. This implies $n$ passes must be performed through the
forward derivative to compute the whole gradient. By using the symmetry in the
chain rule, there is a way to compute the whole gradient faster, and this
method is reverse-mode automatic differentiation. Suppose given a function
$f=f_{n}\circ\ldots\circ f_{1}:\mathbb{R}^{n}\to\mathbb{R}$. Mathematically,
forward mode essentially computes
$(\operatorname{J}f)v=\operatorname{J}f_{n}(\operatorname{J}f_{n-1}(\ldots(\operatorname{J}f_{1}v))\ldots)$
for a direction $v\in\mathbb{R}^{n}$. Reverse-mode, on the other hand,
computes
$(\operatorname{J}f)^{T}v=\operatorname{J}^{T}f_{1}(\operatorname{J}^{T}f_{2}(\ldots(\operatorname{J}^{T}f_{n}v))\ldots)$
for a vector $v\in\mathbb{R}$. In particular, taking $v=1$ computes the
gradient of $f$.

Because the computation flow of the function is reversed, the actual
implementation of reverse-mode is quite tricky. Reverse-mode AD is only well-
understood as a source-code transformation on limited programming languages.
Typically, its implementations on more expressive languages that have features
such as higher-order functions and conditionals make use of 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 evaluate this new program (Carpenter et al., 2015;
Paszke et al., 2017). Such approaches have the severe drawback that the
obtained code cannot benefit from existing optimizing compilers. As such, the
implementation process is tedious and labor-intensive as these AD libraries
need to be implemented using carefully, manually optimized code. In addition,
some whole-program optimizations that a compiler would detect are completely
missed.

###  2.3. Inefficiency of purely functional reverse-mode AD

Following (Pearlmutter and Siskind, 2008), there is a simple way to define an
inefficient yet purely functional reverse-mode transformation for first-order
programs. We review a slight modification of their transformation, which is
also better explained through an example.

Let us consider the term
$x_{1}:\mathbb{R},\ldots,x_{n}:\mathbb{R}\vdash\exp(\cos(x_{i}))$. To compute
its gradient, following the chain rule, we need the Jacobian matrices of
$\cos$ at $x_{i}$ and of $\exp$ at $\cos(x_{i})$. Instead of considering these
as operations from $\mathbb{R}\to\mathbb{R}$, we consider them as functions
from the whole context. So $\cos$ and $\exp$ are seen as functions
$\mathbb{R}^{n}\to\mathbb{R}$. However, by simply doing this, we lose
compositionality. So we modify $\cos$ to also return its context. It is now
seen as a function
$\llbracket\cos\rrbracket:\mathbb{R}^{n}\to\mathbb{R}^{n+1}$. Similarly,
$\exp$ is transformed. It also needs to take the return value of
$\llbracket\cos\rrbracket$ as an extra argument, the one it will actually use
and not simply return. We thus obtain
$\llbracket\exp\rrbracket:\mathbb{R}^{n+1}\to\mathbb{R}^{n+2}$. Now the
jacobians matrices $\operatorname{J}\llbracket\cos\rrbracket\in Mat_{n,n+1}$,
$\operatorname{J}\llbracket\exp\rrbracket\in Mat_{n+1,n+2}$ compose nicely.
The same can be done for binary operators and let bindings. This transforms a
first-order program to a function $f:\mathbb{R}^{n}\to\mathbb{R}^{n+m}$ of the
form $f_{m}\circ\ldots\circ f_{1}$. If the original program was of type
$\mathbb{R}$, then the return value of the original program is the last
component of $f$. Following the mathematical presentation of reverse-mode
above, the gradient of the original program is then obtained as

$\nabla f$ | = | $\operatorname{J}^{T}f(0,\ldots,0,1)=\operatorname{J}^{T}f_{1}(J^{T}f_{2}(\ldots(\operatorname{J}^{T}f_{m}(0,\ldots,0,1))\ldots)$
---|---|---

To actually reverse the order of computation needed for this transpose of
Jacobians, we use a simple continuation; $f_{i}$ is turned into
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}{f_{i}}:=<f_{i},\lambda
Y.Y\circ J^{T}f_{i}>$ where $Y:\mathbb{R}^{n+i-1}\to\mathbb{R}^{n}$. We
recover compositionality by noting that $<f_{i+1}(f_{i}),(\lambda Y.Y\circ
J^{T}f_{i+1})(\lambda Y.Y\circ J^{T}f_{i})>$ reduces to $<f_{i+1}\circ
f_{i},\lambda Y.Y\circ J^{T}f_{i}\circ J^{T}f_{i+1}>$, and thus by induction
we can obtain $<f,\lambda Y.YJ^{T}f>$. By applying the identity continuation
$\mathbb{R}^{n}\to\mathbb{R}^{n}$ on the second component and then the result
to $(0,\ldots,0,1)$, we have obtained a purely functional way to compute
$\nabla f$.

This purely functional implementation has the following issues in terms of
efficiency:

Issue 1. If we see the term as a directed graph, reverse mode back propagates
from the end of the graph to the starting nodes via every path. However, it is
hard to keep track of all this information in parallel in a functional way.
Mutation is usually key for these cases; the imperative version of reverse
mode for a binary operator $op(x,y)$ adds
$x^{\prime}\mathrel{+}=\partial_{1}op(x,y);y^{\prime}\mathrel{+}=\partial_{2}op(x,y)$,
where $\partial_{i}op$ are the partial derivatives of $op$.

Issue 2. Each $J^{T}f_{i+1}$ is a potentially huge matrix if $n$ or $m$ is
big.

Issue 3. We have to carry a continuation and $\beta$-reduce a lot of higher-
order functions.

###  2.4. Insights for efficient purely functional reverse-mode AD

Overall, we use the following three insights to solve the inefficiency
associated with the purely functional implementations of reverse-mode AD.

Insight 1. One of the key simple ideas that we used was to transform every
operator into a unary one. This essentially trivializes the computation flow
to a line. Even if the starting program was a straight-line program, having
non-unary operators was a source of inefficiency and justified the use of
mutation in the first place. By returning every variable every time, the
problem of using a variable several times does not need to be dealt with via
mutation. This simple idea of transforming a program into essentially a
straight line is what our new intermediate representation UNF allows.

Insight 2. If we look at $J^{T}f_{i+1}$, we notice that this function is
almost the identity, except at the last row. Even on the last row, if the
original term was a unary or binary operator like $cos,exp,+,*$, the row is
zero except for at most two indices (one for unary operators). This means we
can use a more compact representation $J^{T}\llbracket
op(x_{i},x_{j})\rrbracket:=\lambda(y_{1},\ldots,y_{n+i}).(y_{1},\ldots,y_{n+i})+[i]\partial_{1}op(x_{i},x_{j})+[j]\partial_{2}op(x_{i},x_{j})$,
where $[k]$ means that the element is added at the $k$-th index of the tuple.

Insight 3. We know in advance that the Jacobian functions are going to be
applied one to another, and we can use partial evaluation to $\beta$-reduce
all of these $\lambda$s. Because each function is almost the identity, we
obtain a lot of substitutions of the form $[x/y]$ where both $x$ and $y$ are
variables. This allows us to drastically reduce the size of
$\operatorname{J}^{T}f$. In fact, for simple programs, this is basically
enough to get an efficient purely functional reverse derivative
transformation. We develop this idea further for a richer language.

##  3\. Simple pure reverse-mode differentiation

###  3.1. Source Language

We consider a standard language, and give it a standard call-by-value
operational semantics. It consists of a first-order functional language with
arrays and a few typical second-order array operations. The types T, terms e,
and typing rules are given in Figure 2. We have included a minimal set of
array operations for the sake of illustration, it is not hard to add more. See
Section 7.

Core Grammar | Description
---|---
T | ::= | $\mathbf{R}$ | Real Type
| $\mid$ |  T $\times$ T | Product Type
| $\mid$ | $\mathbf{A}[\mathbf{R}]^{n}$ | Real Array Type of size $n$
e | ::= |  x $\mid$ c | Variable & Real constant
| $\mid$ | let x = e in e | Variable Binding
| $\mid$ |  $\langle$ e, e $\rangle$ $\mid$ $\pi_{1}$(e) $\mid$ $\pi_{2}$(e) | Pair Constructor/Destructor
| $\mid$ |  e op2 e $\mid$ op1 e | Binary/Unary operations
| $\mid$ |  map2 (x,y.e) e e $\mid$ reduce (x,y.e) e e | Array map2 & reduce

|  $\Gamma\vdash$ x: T
---
(x: T$\in\Gamma$)  $\Gamma\vdash$ e1: T1 $\Gamma\vdash$ e2: T2
---
$\Gamma\vdash$ $\langle$e1,e2$\rangle$: T1 $\times$ T2
$\Gamma\vdash$ e: T1 $\times$ T2
---
$\Gamma\vdash$ $\pi_{i}$e: Ti
($i\in\\{1,2\\}$)
|  $\Gamma\vdash$ e1: T1 $\Gamma$, x: T1 $\vdash$ e2: T2
---
let x = e1 in e2: T2
$\Gamma\vdash$ e: $\mathbf{R}$
---
$\Gamma\vdash$ op1 e : $\mathbf{R}$
$\Gamma\vdash$ e1: $\mathbf{R}$ $\Gamma\vdash$ e2: $\mathbf{R}$
---
$\Gamma\vdash$ e1 op2 e2 : $\mathbf{R}$
|  $\Gamma$, x: $\mathbf{R}$, y: $\mathbf{R}$ $\vdash$ e1: $\mathbf{R}$
$\Gamma$ $\vdash$ e2: $\mathbf{A}[\mathbf{R}]^{n}$ $\Gamma$ $\vdash$ e3:
$\mathbf{A}[\mathbf{R}]^{n}$
---
$\Gamma\vdash$ map2 (x,y.e1) e2 e3: $\mathbf{A}[\mathbf{R}]^{n}$
|  $\Gamma\vdash$ c: $\mathbf{R}$
---
x: $\mathbf{R}$, y: $\mathbf{R}$ $\vdash$ e1: $\mathbf{R}$ $\Gamma$ $\vdash$
e2: $\mathbf{R}$ $\Gamma$ $\vdash$ e3: $\mathbf{A}[\mathbf{R}]^{n}$
---
$\Gamma\vdash$ reduce (x,y.e1) e2 e3: $\mathbf{R}$

Figure 2. Grammar and type system of the source language.

For scalar operations, we assume a given set of operations, including $+$ and
$*$. op1 and op2 denote respectively a unary and a binary operation on reals.
These operations represent smooth total functions, but again, this can be
easily generalized (§7). Typical examples include cos, exp, +, *. We use infix
notation for binary operators.

The reduce operator is a fold-left operator for which the function is assumed
to be associative, and the provided initial value should be a unit of the
binary operation. It is a well-known parallel-friendly construct. For the sake
of simplicity in the presentation, the bound function in reduce is restricted
to having no free variables. Furthermore, as our main focus is on AD and not
array processing, we currently restrict to arrays of reals. We show how to
lift these restrictions and how to differentiate some other array operators in
Section 7 and in the supplementary material.

###  3.2. Target Language

The target language of our source-code transformation is an extension to the
source language. It is a higher-order language, as our purely functional
reverse-mode introduces a continuation. The set of scalar operations should
also be closed under partial differentiation. In more detail, for every unary
scalar operation op1, we assumed a given operator $\partial$op1 whose
semantics should be the derivative of op1, e.g. $\partial$sin=cos. Similarly,
for every binary operator op2, we assume given operators $\partial_{1}$op2,
$\partial_{2}$op2, respectively representing the first and second partial
derivative of op2.

Similarly, the target language contains more array primitives, which are used
to define the reverse derivatives of array operations. Scan left scanl is
similar to fold left, but also stores all the intermediate results in an
array, which it returns. In the same vein, scan right scanr performs a fold
left by reading the array from right to left and stores the intermediate
results in an array from right to left.

Finally, we add two new shift operators shift1L and shift1R. They take an
array of size $n$, and respectively forget the first and the last element of
the array. These somewhat ad-hoc operators naturally show up when
differentiating fold-like operators.

The grammar for types and terms along with the type system of the target
language are presented in Figure 3. Our lambda abstractions take $n$
arguments, as we are not concerned with partial applications in this work. In
fact, the lambda abstractions introduced by reverse-mode will be removed
during partial evaluation, and the notation with lambda abstractions having
$n$ bound variables makes reading slightly easier. We note that we don’t
actually need the full power of higher-order because we only use lambda
abstractions over variables of ground types and let expressions binding such
lambda abstractions. We only need the target language to be second-order.

Core Grammar | Description
---|---
T | ::= | $\ldots$ | Same as Source
| $\mid$ |  T$\times\ldots\times$T |  $n$-tuples
| $\mid$ | T$\rightarrow\;$T | Function Type
e | ::= | $\ldots$ | Same as Source
| $\mid$ | fun (x1,$\ldots$,xn) $\rightarrow\;$ e | Lambda Abstraction
| $\mid$ | e(e${}_{1}\ldots$en) | Function Application
| $\mid$ | $\langle$e,$\ldots$,e$\rangle$ | Tuples
| $\mid$ |  scanl (x,y.e) e e $\mid$ scanr (x,y.e) e e | Array scan left and right
| $\mid$ |  shift1L e $\mid$ shift1R e | Array left/right shifting

|  $\Gamma$, x: $\mathbf{R}$, y: $\mathbf{R}$ $\vdash$ e1: $\mathbf{R}$
$\Gamma$ $\vdash$ e2: $\mathbf{R}$ $\Gamma$ $\vdash$ e3:
$\mathbf{A}[\mathbf{R}]^{n}$
---
$\Gamma\vdash$ scanl (x,y.e1) e2 e3: $\mathbf{A}[\mathbf{R}]^{n+1}$
|  $\Gamma$, x: $\mathbf{R}$, y: $\mathbf{R}$ $\vdash$ e1: $\mathbf{R}$
$\Gamma$ $\vdash$ e2: $\mathbf{R}$ $\Gamma$ $\vdash$ e3:
$\mathbf{A}[\mathbf{R}]^{n}$
---
$\Gamma\vdash$ scanr (x,y.e1) e2 e3: $\mathbf{A}[\mathbf{R}]^{n+1}$
|  $\Gamma$, x1: G1, $\ldots$, xn: Gn $\vdash$ e: T
---
$\Gamma\vdash$ fun (x1,$\ldots$,xn) $\rightarrow\;$ e:
G${}_{1}\times\ldots\times$Gn$\rightarrow\;$T
|  $\Gamma$ $\vdash$ e: G${}_{1}\times\ldots\times$Gn $\rightarrow\;$ T
$\Gamma$ $\vdash$ ei: Gi for all $1\leq i\leq n$
---
$\Gamma\vdash$ e(e${}_{1}\ldots$en): T
|  $\Gamma$ $\vdash$ e: $\mathbf{A}[\mathbf{R}]^{n+1}$
---
$\Gamma\vdash$ shift1L e: $\mathbf{A}[\mathbf{R}]^{n}$
$\Gamma$ $\vdash$ e: $\mathbf{A}[\mathbf{R}]^{0}$
---
$\Gamma\vdash$ shift1L e: $\mathbf{A}[\mathbf{R}]^{0}$
|  $\Gamma$ $\vdash$ e: $\mathbf{A}[\mathbf{R}]^{n+1}$
---
$\Gamma\vdash$ shift1R e: $\mathbf{A}[\mathbf{R}]^{n}$
$\Gamma$ $\vdash$ e: $\mathbf{A}[\mathbf{R}]^{0}$
---
$\Gamma\vdash$ shift1R e: $\mathbf{A}[\mathbf{R}]^{0}$
| for all $i$, $\Gamma$ $\vdash$ ei: Ti
---
$\Gamma\vdash$ $\langle$e1,$\ldots$,en$\rangle$: T${}_{1}\times\ldots\times$Tn

Figure 3. Grammar and type system of the target language.

###  3.3. Macro for pure reverse mode transformation

In Figure 5 we present our direct transformation from the source language to
the target language for pure reverse mode differentiation. Given a term
$\Gamma\vdash e:\mathbf{R}$, we can compute its gradient $\nabla_{\Gamma}e$
from a particular instance of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}(e)$. First,
$\rho,Y$ specifies if we want to compute the whole gradient regarding the
variables from $\Gamma$ or a subset of it. For a subset $\rho\subset\Gamma$,
one chooses $Y$ to be the projection function sending a variable $x_{i}:G$ of
$\Gamma$ to $x_{i}$ if it belongs to $\rho$ and to $0_{G}$ otherwise. In
particular, we take $Y=Id_{\Gamma}$ to compute the whole gradient. Next, the
gradient will be given by the second part of the pair
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}(e)$, and we
need to initialize the tangent variables. All of them are set to $0$, except
the one corresponding to the output value of e, which we initialize at $1$ to
run the backpropagation. All in all, we compute the gradient via
$\pi_{2}{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Id_{\Gamma}}(e)(0_{\Gamma},1)$.

let x1=e1,$\ldots$,xn=en | = | let x1=e1 in let x2=e2 in $\ldots$
---|---|---
in e | let xn = en in e
$Id_{\Gamma}$ ($\Gamma\,=\,x_{1}:A_{1},\ldots,x_{n}:A_{n})$ | = |  fun $(y_{1}:A_{1},\ldots,y_{n}:A_{n})$ $\rightarrow\;$$(y_{1},\ldots,y_{n})$
$\nabla_{\Gamma}(e)$ ($\Gamma\vdash e:\mathbb{R}$) | = | $\pi_{2}{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\Gamma}_{\Gamma;Id_{\Gamma}}(e)(0_{\Gamma},\underline{1})$
pos(x) (x$\in\Gamma=x_{:}A_{1},\ldots,x_{n}:A_{n})$ | = | position $i$ of x in $\Gamma$
[i]e (e of ground type $G_{i}$) | = | $(0_{G_{1}},\ldots,0_{G_{i-1}},e,0_{G_{i+1}},\ldots,0_{G_{n}})$
| ($G_{j}$ are ground types)
$\nabla_{\Gamma_{1}}$(e) ($\Gamma=\Gamma_{1},\Gamma_{2}$, $|\Gamma_{1}|=k$) | = | $(e_{1},\ldots,e_{k})$
$\nabla_{\Gamma_{2}}$(e) | = | $(e_{k+1},\ldots,e_{n})$
|  | when $\nabla_{\Gamma}($e$)=(e_{1},\ldots,e_{n})$
map (x. e) A | = | map2 (x y. e) A A
ZerosLike(A) | = | map (x.0) A
OnesLike(A) | = | map (x.1) A
ZerosLike(n) | = | map (x.0) (\_:$\mathbf{A}[\mathbf{R}]^{n}$)
OnesLike(n) | = | map (x.1) (\_:$\mathbf{A}[\mathbf{R}]^{n}$)
Figure 4. Notations used for reverse-mode AD transformation. $\\_$ represents
a dummy variable added to the context.

###### Notation 1.

We now introduce several notations which are useful when defining the
transformation for reverse-mode in Figure 4. Ground types are defined
inductively by

$G::=\mathbb{R}\mid G\times\ldots\times G\mid\mathbf{A}[\mathbb{R}]^{n}$

$(\mathbb{R},+,\underline{0})$ forms a monoid, and this monoid structure
extends canonically to a monoid structure $(G,\widehat{+},0_{G})$ for every
ground type $G$. It is defined inductively on $G$ as follows

$0_{\mathbb{R}}$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $\underline{0}$
---|---|---
$0_{G_{1}\times\ldots\times G_{n}}$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $<O_{G_{1}},\ldots,0_{G_{n}}>$
$0_{\mathbf{A}[\mathbb{R}]^{n}}$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | ZerosLike(n)
$a\widehat{+}_{\mathbb{R}}b$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $a+b$
---|---|---
$(a_{1},\ldots,a_{n})\widehat{+}_{G_{1}\times\ldots\times G_{n}}$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $(a_{1}\widehat{+}_{G_{1}}b_{1},\ldots,a_{n}\widehat{+}_{G_{n}}b_{n})$
$(b_{1},\ldots,b_{n})$ |  |
$A\widehat{+}_{\mathbf{A}[\mathbb{R}]^{n}}B$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ |  map2 \+ $A$ $B$

A ground context is a context only containing variables of ground type. The
previous monoid structure again extends canonically to ground contexts
$\Gamma$ by defining
$0_{x1:G1,\ldots,x_{n}:G_{n}}\stackrel{{\scriptstyle\mathrm{def}}}{{=}}0_{G_{1}},\ldots,0_{G_{n}}$
and
$a_{1},\ldots,a_{n}\widehat{+}_{x1:G_{1},\ldots,x_{n}:G_{n}}b_{1},\ldots,b_{n}\stackrel{{\scriptstyle\mathrm{def}}}{{=}}a_{1}\widehat{+}_{G_{1}}b_{1},\ldots,a_{n}\widehat{+}_{G_{n}}b_{n}$.

###### Example 3.1.

The reverse-mode transformation of the terms from the introduction are given
by

|  ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(let w1
= x1 * x2 in let w2 = w1 * x1 in w2)
---|---
= | let w1,Y1=
|  let y11,Y11= $\langle$x1, fun (y1,y2,y3,z) $\rightarrow\;$
Y(y1+z,y2,y3)$\rangle$ in
|  let y12,Y12= $\langle$x2, fun (y1,y2,y3,y4,z) $\rightarrow\;$
Y11(y1,y2+z,y3,y4)$\rangle$ in
|  $\langle$y11 * y12, fun (y1,y2,y3,z) $\rightarrow\;$
Y12(y1,y2,y3,y12*z,y11*z) $\rangle$ in
| let w2,Y2=
|  let y21,Y21= $\langle$w1, fun (y1,y2,y3,y4,z) $\rightarrow\;$
Y1(y1,y2,y3,y4+z)$\rangle$ in
|  let y22,Y22= $\langle$x1, fun (y1,y2,y3,y4,y5,z) $\rightarrow\;$
Y21(y1+z,y2,y3,y4,y5)$\rangle$ in
|  $\langle$y21 * y22, fun (y1,y2,y3,y4,z) $\rightarrow\;$
Y22(y1,y2,y3,y4,y22*z,y21*z) $\rangle$ in
| let y,Y3= $\langle$w, fun (y1,y2,y3,y4,z) $\rightarrow\;$
Y2(y1,y2,y3,y4+z)$\rangle$ in
| $\langle$y, fun (y1,y2,y3,z) $\rightarrow\;$ Y3(y1,y2,y3,0,z) $\rangle$

|
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(prod(A))
---|---
= | let y,Y1= $\langle$1, fun (X,z) $\rightarrow\;$ Y(X)$\rangle$ in
| let B,Y2= $\langle$A, fun (X,x,Z) $\rightarrow\;$ Y1(X+Z,x)$\rangle$ in
| let A0= shift1R (scanl * y B) in
| let A1= shift1L (map2 (a,b.b) A0 B) in
| let A2= map2 (a,b.a) A0 B in
| let A3= scanr * 1 A1 in
| $\langle$prod(B), fun (X,z) $\rightarrow\;$ Y2(X,0,map2 (a,b. a*z*b) A2
A3)$\rangle$

The idea is that $\rho$ represents the return type of the derivative part,
which should be A${}_{1}\times\ldots\times$ An if we want the whole gradient
of a term e in context $\Gamma=\;$x1:A1,$\ldots$,xn:An. The subscript $\Gamma$
denotes the current context, which is locally augmented, for instance when
differentiating a let rule. For non-unary operations, we differentiate the
arguments from the left to the right and add their derivatives to the current
stack, which is modeled by the continuation $Y$. Importantly for performance,
each continuation variable $Y$ is only used once.

We have the following typing lemma for
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$, routinely
proved by induction on derivation of $\Gamma\vdash$e:A.

###### Lemma 3.2 (Typing
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$).

If $\Gamma\vdash$ e: A, then
$\Gamma$,Y:$\Gamma\to\rho\vdash{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e):${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(A).

$\mathcal{D}_{\Gamma;Y}$(A) | = | A $\times$ ($\Gamma$ $\times$ A$\rightarrow\;$$\rho$)
---|---|---
$\mathcal{D}_{\Gamma;Y}$(c) | = | $\langle$c, fun ($\boldsymbol{x}$,z) $\rightarrow\;$
|  | Y($\boldsymbol{x}$) $\rangle$
$\mathcal{D}_{\Gamma;Y}$(x) | = | $\langle$x, fun ($\boldsymbol{x}$,z) $\rightarrow\;$
|  |  Y($\boldsymbol{x}$$\widehat{+}$[pos(x)]z) $\rangle$
$\mathcal{D}_{\Gamma;Y}$(let x:A = e1 in e2) | = |  let x,Y1 =$\mathcal{D}_{\Gamma;Y}$(e1) in
|  |  let y,Y2 =$\mathcal{D}_{\Gamma,x:A;Y_{1}}$(e2) in
|  |  $\langle$y, fun ($\boldsymbol{x}$,z) $\rightarrow\;$ Y2($\boldsymbol{x}$,$0_{A}$,z)$\rangle$
$\mathcal{D}_{\Gamma;Y}$($\langle$ e1, e2 $\rangle$) | = |  let y1,Y1 =$\mathcal{D}_{\Gamma;Y}$(e1) in
|  |  let y2,Y2 =$\mathcal{D}_{\Gamma,x_{1};Y_{1}}$(e2) in
|  | $\langle$ $\langle$y1,y2$\rangle$, fun ($\boldsymbol{x}$,z) $\rightarrow\;$
|  |  Y($\boldsymbol{x}$,$\pi_{1}$(z),$\pi_{2}$(z)) $\rangle$
$\mathcal{D}_{\Gamma;Y}$($\pi_{1}$(e:A$\times$B)) | = |  let x,Y1 =$\mathcal{D}_{\Gamma;Y}$(e) in
|  |  $\langle$$\pi_{1}$x, fun ($\boldsymbol{x}$,z) $\rightarrow\;$ Y($\boldsymbol{x}$,(z,$0_{B}$))$\rangle$
$\mathcal{D}_{\Gamma;Y}$($\pi_{2}$(e:A$\times$B)) | = |  let x,Y1 =$\mathcal{D}_{\Gamma;Y}$(e) in
|  |  $\langle$$\pi_{2}$x, fun ($\boldsymbol{x}$,z) $\rightarrow\;$ Y($\boldsymbol{x}$,($0_{A}$,z))$\rangle$
$\mathcal{D}_{\Gamma;Y}$(op1 e) | = |  let x,Y1 =$\mathcal{D}_{\Gamma;Y}$(e) in
|  | $\langle$op1 x, fun ($\boldsymbol{x}$,z) $\rightarrow\;$
|  |  Y($\boldsymbol{x}$,$\partial$op1(x)*z) $\rangle$
$\mathcal{D}_{\Gamma;Y}$(e1 op2 e2) | = |  let x1,Y1 =$\mathcal{D}_{\Gamma;Y}$(e1) in
|  |  let x2,Y2 =$\mathcal{D}_{\Gamma,x_{1};Y_{1}}$(e2) in
|  | $\langle$x1 op2 x2, fun ($\boldsymbol{x}$,z) $\rightarrow\;$
|  |  Y2($\boldsymbol{x}$,$\partial_{1}$op2(x1,x2)*z,$\partial_{2}$op2(x1,x2)*z$\rangle$
$\mathcal{D}_{\Gamma;Y}$(map2 (x,y.e1) e2 e3) | = |  let A,Y1 =$\mathcal{D}_{\Gamma;Y}$(e2) in
|  |  let B,Y2 =$\mathcal{D}_{\Gamma,A;Y_{1}}$(e3) in
|  | $\langle$map2 (x,y.e1) A B, fun ($\boldsymbol{x}$,Z) $\rightarrow\;$
|  | let G = (map2 * Z
|  |  (map2 (a,b.($\nabla_{\Gamma}$e1)[a/x.b/y]) A B)) in
|  | Y2( $\boldsymbol{x}$ $\widehat{+}$(reduce $\widehat{+}$ $\widehat{0}$ G),
|  |  map2 * (map2 (a,b.($\nabla_{\\{x\\}}$e1)[a/x,b/y]) A B) Z,
|  |  map2 * (map2 (a,b.($\nabla_{\\{y\\}}$e1)[a/x,b/y]) A B) Z)$\rangle$
$\mathcal{D}_{\Gamma;Y}$(reduce (x,y.e1) e2 e3) | = |  let y1,Y1 =$\mathcal{D}_{\Gamma;Y}$(e2) in
|  |  let A,Y2 =$\mathcal{D}_{\Gamma,y_{1};Y_{1}}$(e3) in
|  | let A0=shift1R (scanl (x,y.e1) y1 A) in
|  | let A1=shift1L (map2
|  |  (a,b.($\nabla_{\\{x\\}}$e1)[a/x,b/y]) A0 A) in
|  |  let A2=map2 (a,b.($\nabla_{\\{y\\}}$)e1[a/x,b/y]) A0 A in
|  | let A3=scanr * 1 A1 in
|  | $\langle$reduce (x,y.e1) y1 A, fun ($\boldsymbol{x}$,z) $\rightarrow\;$
|  | Y2($\boldsymbol{x}$, map2 (x,y. x*y*z) A2 A3)$\rangle$
Figure 5. Reverse-mode transformation from source to target language. We write
$\mathcal{D}_{\Gamma;Y}$ instead of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$, and
$\boldsymbol{x}$ instead of x1,$\ldots$,xn to alleviate notation.

Admittedly, the transformation presented in this section may be hard to read,
and it is not straightforward to show its correctness directly. Next, we
decompose this transformation into three simpler steps via a novel
intermediate representation.

##  4\. Unary Normal Form

Following the intuition highlighted in Section 2.4, we present a new language,
which we call Unary Normal Form (UNF). There is a Source UNF for the Source
language and a Target UNF for the Target language. Just as Target is an
extension of Source, Target UNF is an extension of Source UNF. They serve as
intermediate representations in the reverse-mode compilation pipeline (see
Figure 1).

One takeaway from this paper is that efficient pure reverse-mode is
complicated because of the several things it does. It goes through the term
and keeps a store for the gradients to be updated. It keeps track of the new
bound variables found while going through the term. The updates to the
gradient are done via pre-composition and not post-composition. Due to the
linearity requirement in the usage of the continuation variable, it needs to
process like a call-by-value evaluation and evaluate the arguments of
operators in a certain order before dealing with the operator itself.

By introducing UNF, we are decoupling some of the problems. UNF introduces the
call-by-value evaluation of the arguments before evaluating the arguments and
a good management of the environment. Differentiation on UNF deals with the
purely functional store and gradients update via pre-composition. Finally,
going back from UNF to Target uses the stored information in UNF to make sure
the Jacobians are computed efficiently.

###  4.1. Source UNF

Intuitively, a term consists of a composition of unary operators. To compile
our source language to this intermediate representation, we need to remember
some information about the context of the initial term. The grammar of Source
UNF is given in Figure 6.

Core Grammar | Description
---|---
T | ::= | [A1,$\ldots$,An] | Lists of types from source
e | ::= |  varT;i | Variable
| $\mid$ |  opT;n | Operations, for $0\leq n\leq 2$
| $\mid$ |  pairT;A×B | Pairing a pair of variables
| $\mid$ |  proj${}_{T_{1};T_{2};T_{3}}$ | Projection
| $\mid$ |  e;e | Sequential composition
| $\mid$ |  map2T;x,y.e | Map2
| $\mid$ |  reduceT;x,y.e;e | Reduce

|  T $\vdash$ varT;i: T,Ai
---
(T=A1,$\ldots$,An)  T,$\mathbf{R}^{\times(n)}$ $\vdash$ opT;n :
T,$\mathbf{R}^{\times(n+1)}$
---
(T=A1,$\ldots$,An)
|  T $\vdash$ e1: T,A T,A $\vdash$ e2: T,A,B
---
T $\vdash$ e1;e2: T,A,B
(T=A1,$\ldots$,An)
|  T,A,B $\vdash$ pairT;A×B : T,A$\times$B
---
(T=A1,$\ldots$,An)
|  T1,T2,T3 $\vdash$ proj${}_{T_{1};T_{2};T_{3}}$ : T1,T3
---
|  x1:A1,$\ldots$,xn:An,x:$\mathbf{R}$,y:$\mathbf{R}$ $\vdash$ e:$\mathbf{R}$
in Source Language
---
T,$\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n}$ $\vdash$
map2T;x,y.e:
T,$\mathbf{A}[\mathbf{R}]^{n}$,$\mathbf{A}[\mathbf{R}]^{n}$,$\mathbf{A}[\mathbf{R}]^{n}$
(T=A1,$\ldots$,An)
|  x:$\mathbf{R}$,y:$\mathbf{R}$ $\vdash$ e: $\mathbf{R}$ x1:A1,$\ldots$,xn:An
$\vdash$ e2:$\mathbf{R}$ in Source Language
---
T,$\mathbf{A}[\mathbf{R}]^{n}$ $\vdash$ reduce${}_{T;x,y.e_{1};e_{2}}$:
T,$\mathbf{A}[\mathbf{R}]^{n}$,$\mathbf{R}$
(T=A1,$\ldots$,An)

Figure 6. Grammar and type system of the source UNF.

There are a few notable things in this syntax. Every primitive is indexed by a
list of types from the source language, which corresponds to a context being
carried. We will often elide the brackets [,]. Every constant, unary, or
binary operator from Source has a corresponding $n$-ary operator opT,n in UNF.
Sequential composition is denoted by ; and e1;e2 means that e1 should be
performed, and then e2. The array operators map2 and reduce have extra indices
that represent well-formed terms in Source. The language also contains a
projection proj which forgets some of the elements, and a pairing operator
pair which takes the last pair of elements and returns the pairing of these
elements.

A judgment in Source UNF is a triple (T1,e,T2) where T1 and T2 are lists of
types of Source, and e is an expression of the Source UNF. The typing rules
are detailed in Figure 6. $\mathbf{R}^{\times n}$ is a notation for
$\mathbf{R},\ldots,\mathbf{R}$ with $n$ factors $\mathbf{R}$. When T1 and T2
are lists of types and A a type, the operator , denotes snoc in T1,A and
denotes append in T1,T2.

###### Example 4.1.

The term in Source UNF cosR,R ; pairR,R;R×R intuitively represents a term
x1:$\mathbf{R}$,x2:$\mathbf{R}$,x3:$\mathbf{R}$ $\vdash$ let
(x1,x2,x3,x4)=(x1,x2,x3,cos(x3)) in (x1,x2,$\langle$x3,x4$\rangle$).
The operator ; can more generally be understood as a let binding of a tuple,
and primitives only act on the last parts of the tuple, always also returning
their whole context.

###  4.2. Target UNF

Target UNF is an extension of Source UNF. Its types are lists of types of
Target. Higher-order types are needed for the continuation introduced by
reverse-mode. In addition, it contains several new primitives which represent
the transpose Jacobian of the primitives of Source UNF. Given a type
T$\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ [A1,$\ldots$,An], we write T for
the type A1 $\times\ldots\times$ An of Target. Target UNF has an internal
composition $\circ$. e1 $\circ$ e2 represents precomposition of a curried
function e1 by a term e2. It is used by the reverse-mode transformation to
pre-compose the continuation term by a transpose Jacobian. The grammar and
type system of target UNF is given in Figure 7.

Core Grammar | Description
---|---
T | ::= | [A1,$\ldots$,An] | Lists of types from target
e | ::= | $\ldots$ | Same as source UNF
| $\mid$ |  JTvarT;i | Jacobian for variable
| $\mid$ |  JTopT;n | Jacobian for operation, $0\leq n\leq 2$
| $\mid$ |  JTpairT;A×B | Jacobian for pairing
| $\mid$ |  JTproj${}_{T_{1};T_{2};T_{3}}$ | Jacobian for projection
| $\mid$ |  JTmap2T;x,y.e | Jacobian for map2
| $\mid$ |  JTreduceT;x,y.e;e | Jacobian for reduce
| $\mid$ | $\langle$e, e$\rangle$ | Term pairing
| $\mid$ |  e $\circ$ e | Internal function composition

|  T,Ai $\vdash$ JTvarT;i: T
---
(T=A1,$\ldots$,An)  T,$\mathbf{R}^{\times(n+1)}\vdash$ JTopT;n :
T,$\mathbf{R}^{\times(n)}$
---
(T=A1,$\ldots$,An)
|  T,A$\times$B $\vdash$ JTpairT;A×B : T,A,B
---
(T=A1,$\ldots$,An)
|  T1,T3 $\vdash$ JTproj${}_{T_{1};T_{2};T_{3}}$ : T1,T2,T3
---
|  x1:A1,$\ldots$,xn:An,x:$\mathbf{R}$,y:$\mathbf{R}$ $\vdash$ e: $\mathbf{R}$
in Source Language
---
T,$\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n}\vdash$
JTmap2T;x,y.e: T,$\mathbf{A}[\mathbf{R}]^{n}$,$\mathbf{A}[\mathbf{R}]^{n}$
(T=A1,$\ldots$,An)
|  x:$\mathbf{R}$,y:$\mathbf{R}$ $\vdash$ e1: $\mathbf{R}$
x1:A1,$\ldots$,xn:An $\vdash$ e2:$\mathbf{R}$ in Source Language
---
T,$\mathbf{A}[\mathbf{R}]^{n},\mathbf{R}\vdash$
JTreduce${}_{T;x,y.e_{1};e_{2}}$: T,$\mathbf{A}[\mathbf{R}]^{n}$
(T=A1,$\ldots$,An)
|  T $\vdash$ e1: T1 T $\vdash$ e1: T2
---
T $\vdash$ $\langle$e1, e2$\rangle$: T1,T2
T1 $\vdash$ e1: [T3 $\to$ B] T2 $\vdash$ e2: T3
---
T1 $\vdash$ e1 $\circ$ e2:[T2 $\to$ B]

Figure 7. Grammar and type system of Target UNF.

###  4.3. Simple reverse mode transformation

We are now able to present a simpler transformation for purely functional
reverse-mode from Source UNF to Target UNF. The differentiation transformation
is given in Figure 8. To a list of types, the transformation adds a higher-
order type. It is to be understood as the continuation for the pure storage of
the gradient. On primitive constants of Source UNF, it returns a pair. The
first part is the initial term (pre-composed by a projection because it is not
using the continuation variable). The second part returns the continuation
term pre-composed by the transpose Jacobian of the primitive. There is a sharp
contrast between the simplicity of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$ compared to
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ from
Figure 5.

###### Example 4.2.

Continuing with our previous example cosR,R ; pairR,R;R×R, we have

$\mathbf{R}$,$\mathbf{R}$,$\mathbf{R}$,$\mathbf{R}\times\mathbf{R}\times\mathbf{R}$$\rightarrow\;$$\rho\vdash{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(cos${}_{\mathbf{R},\mathbf{R}};$
pairR,R;R×R):
$\mathbf{R}$,$\mathbf{R}$,$\mathbf{R}\times\mathbf{R}$,$(\mathbf{R}\times\mathbf{R}\times(\mathbf{R}\times\mathbf{R}))$$\rightarrow\;$$\rho$
---
= $\langle$proj${}_{\mathbf{R}^{3};\mathbf{R}^{3}->\rho;[]};$ cosR,R,
proj${}_{[];\mathbf{R}^{3};\mathbf{R}^{3}->\rho}\circ$
(proj${}_{\mathbf{R}^{3};\mathbf{R}^{3}->\rho;[]}$ ; JTcosR,R)$\rangle$ ;
$\langle$proj${}_{\mathbf{R}^{4};\mathbf{R}^{4}->\rho;[]};$ pairR,R;R×R,
proj${}_{[];\mathbf{R}^{4};\mathbf{R}^{4}->\rho}\circ$
(proj${}_{\mathbf{R}^{4};\mathbf{R}^{4}->\rho;[]}$ ; JTpairR,R;R×R)$\rangle$

###### Lemma 4.3 (Well typedness of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$).

Let A1,$\ldots$,An $\vdash$ e: B1,$\ldots$,Bm be a term in Source UNF.
Then A1,$\ldots$,An,A${}_{1}\times\ldots\times$ An$\rightarrow\;$$\rho$
$\vdash$ ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(e):
B1,$\ldots$,Bm,B${}_{1}\times\ldots\times$ Bm$\rightarrow\;$$\rho$.

###### Proof.

By induction on derivation of A1,$\ldots$,An $\vdash$ e: B1,$\ldots$,Bm. ∎

|  ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(A1,$\ldots$,An) | = | A1,$\ldots$,An,(A${}_{1}\times\ldots\times$An)$\rightarrow\;$$\rho$
---|---|---
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(varT;i) | = |  $F$(varT;i, JTvarT;i)
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(opΓ;n) | = |  $F$(opT;n,JTopT;n)
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(pairT;A×B) | = |  $F$(pairT;A×B, JTpairT;A×B)
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(proj${}_{T_{1};T_{2};T_{3}}$) | = |  $F$(proj${}_{T_{1};T_{2};T_{3}}$,JTprojT1;T2;T3)
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(e${}_{1};$e2) | = |  ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(e1); ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(e2)
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(map2T;x,y.e) | = |  $F$(map2T;x,y.e, JTmap2T;x,y.e)
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(reduce${}_{T;x,y.e_{1};e_{2}}$) | = |  $F$(reduce${}_{T;x,y.e_{1};e_{2}}$, JTreduce${}_{T;x,y.e_{1};e_{2}}$)
where $F$(A,B)$\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$
$\langle$proj${}_{T;\underline{T}->\rho;[]};$ A,
proj${}_{[];T;\underline{T}->\rho}\circ$ (proj${}_{T;\underline{T}->\rho;[]}$
; B)$\rangle$
Figure 8. Reverse-mode differentiation from Source UNF to Target UNF

###  4.4. Transformations to and from UNF

We first give a translation from our Source to Source UNF in Figure 9, which
we call $\mathbf{UNF}$. As discussed above, $\mathbf{UNF}$ sends a constant or
an operator to a primitive in Source UNF. This term in Source UNF carries the
context of the original term in Source. $\mathbf{UNF}$ sequentializes a term,
mimicking a left-to-right call-by-value evaluation. Because of this, when we
are trying to return a pair, this needs to be witnessed, and this is the role
of pair. When we need to pass the result of type A of e1 through e2 which does
not take it as an input, we can transform e2 to accept and pass it. It’s the
equivalent in UNF of a weakening. This weakening is required when we transform
a non-unary operator, such as $+$ or pair, as we transform the arguments in
order, and we need to keep all of their results. We write this new term
$(\widetilde{e_{2}})_{A}$. It is defined by induction on e2 as follows. We
assume given a type A that needs to be passed, and drop the index A.

$\widetilde{var_{T;i}}$ | = | $var_{T,A;i}$
---|---|---
$\widetilde{op_{T;n}}$ | = | $op_{T,A;n}$
$\widetilde{pair_{T;B\times C}}$ | = | $pair_{T,A;B\times C}$
$\widetilde{proj_{T_{1};T_{2};T_{3}}}$ | = | $proj_{T_{1};T_{2};T_{3},A}$
$\widetilde{e_{1};e_{2}}$ | = | $\widetilde{e_{1}};\widetilde{e_{2}}$
---|---|---
$\widetilde{map2_{T;x,y.e}}$ | = | $map2_{T,A;x,y.e}$
$\widetilde{reduce_{T;x,y.e_{1};e_{2}}}$ | = | $reduce_{T,A;x,y.e_{1};e_{2}}$

We want to preserve the invariant that $\mathbf{UNF}$(e) represents the term e
which is also returning its context. Because x is not free in $\Gamma\vdash$
let x=e1in e2, we need to hide x after $\mathbf{UNF}$(e2). This explains the
projection proj in $\mathbf{UNF}$ of a let. More generally, in e${}_{1};$ e2
the return value of e1 acts as a new variable for e2, is not in the context of
e1, but will be returned by e2. This means we need to hide all these
intermediate values to preserve the invariant, and this explains the
projection proj in the $\mathbf{UNF}$ of other terms.

###### Lemma 4.4 (Well typedness of $\mathbf{UNF}$).

Let $\Gamma$=x1:A1,$\ldots$,xn:An $\vdash$ e:B be a term in Source. Then

A1,$\ldots$,An $\vdash$ $\mathbf{UNF}$(e): A1,$\ldots$,An,B.

###### Proof.

By induction on derivation of $\Gamma\vdash$e:B. ∎

$\mathbf{UNF}$($\Gamma\vdash$ c) | = |  cΓ,0 constant seen as a 0-ary operator
---|---|---
$\mathbf{UNF}$($\Gamma\vdash$ x) | = |  varΓ,i where x is the $i$-th variable in $\Gamma$
$\mathbf{UNF}$($\Gamma\vdash$ let x:A = e1 in e2:B) | = |  $\mathbf{UNF}$(e1) ; $\mathbf{UNF}$(e2) ; projΓ;A;B
$\mathbf{UNF}$($\Gamma\vdash$ $\langle$ e1, e2 $\rangle$:AxB) | = |  $\mathbf{UNF}$(e1) ; $\widetilde{\mathbf{UNF}(e_{2})}$ ; pairΓ;A×B
$\mathbf{UNF}$($\Gamma\vdash\pi_{i}$(e)) | = |  $\mathbf{UNF}$(e); $\pi_{i}$ ; proj${}_{\Gamma;A_{1}\times A_{2};A_{i}}$
|  |  $\pi_{i}$ seen as a unary operator
$\mathbf{UNF}$($\Gamma\vdash$ e1 op2 e2) | = |  $\mathbf{UNF}$(e1) ; $\widetilde{\mathbf{UNF}(e_{2})}$ ; opΓ;2 ; projΓ;R,R;R
$\mathbf{UNF}$($\Gamma\vdash$ op1 e) | = |  $\mathbf{UNF}$(e) ; opΓ;1 ; projΓ;R;R
$\mathbf{UNF}$($\Gamma\vdash$ map2 (x,y.e1) e2 e3) | = |  $\mathbf{UNF}$(e2) ; $\widetilde{\mathbf{UNF}(e_{3})}$ ; map2${}_{\Gamma;x,y.e_{1}}$ ;
|  |  proj${}_{\Gamma;\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n};\mathbf{A}[\mathbf{R}]^{n}}$
$\mathbf{UNF}$($\Gamma\vdash$ reduce (x,y.e1) e2 e3) | = |  $\mathbf{UNF}$(e3); reduce${}_{\Gamma;x,y.e_{1};e_{2}}$ ; proj${}_{\Gamma;\mathbf{A}[\mathbf{R}]^{n};\mathbf{A}[\mathbf{R}]^{n}}$
Figure 9. Transformation from Source to Source UNF

Next, the transformation from Target UNF to Target is presented in Figure 10.
We call this transformation $\mathbf{UNF}^{-1}$, but it is not a strict
inverse of $\mathbf{UNF}$. Doing $\mathbf{UNF}$ followed by
$\mathbf{UNF}^{-1}$ performs some version of the ANF transformation (Sabry and
Felleisen, 1993).

Target UNF does not have variables. A context type T=[A1,$\ldots$,An] is
transformed to a context
$\Gamma$=x1:A1,$\ldots$,xn:An. We use this convention when definiting
$\mathbf{UNF}^{-1}$. For operators like opΓ;m, we extend the notation above by
saying that the variables from the context are
x1:A1,$\ldots$,xn:An,xn+1:$\mathbf{R}$,$\ldots$,xn+m:$\mathbf{R}$.

$\mathbf{UNF}^{-1}$ does not have as a simple typing property as
$\mathbf{UNF}$, because it treats primitives with and without a $J$
differently and should be performed after
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$.

###### Lemma 4.5 (Well typedness of $\mathbf{UNF}^{-1}$).

Let A1,$\ldots$,An $\vdash$ e: B1,$\ldots$,Bm be a term in Source UNF. Then

x1:A1,$\ldots$,xn:An,xn+1:A${}_{1}\times\ldots\times$ An$\rightarrow\;$$\rho$
$\vdash$
$\mathbf{UNF}^{-1}$(${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(e)):
Bm $\times$ (B${}_{1}\times\ldots\times$ Bm$\rightarrow\;$$\rho$).

###### Proof.

By induction on derivation of A1,$\ldots$,An $\vdash$ e: B1,$\ldots$,Bm. ∎

$\mathbf{UNF}^{-1}$(varΓ;i) | = | xi: Ti
---|---|---
$\mathbf{UNF}^{-1}$(opΓ;m) | = | opn(xn,$\ldots$,xn+m)
$\mathbf{UNF}^{-1}$(e1 ; e2) | = | let xn+1=$\mathbf{UNF}^{-1}$(e1) in $\mathbf{UNF}^{-1}$(e2)
$\mathbf{UNF}^{-1}$(map2Γ;x,y.e) | = | map2 (x,y.e) xn+1 xn+2
$\mathbf{UNF}^{-1}$(reduce${}_{\Gamma;x,y.e_{1};e_{2}}$) | = | reduce (x,y.e1) $e_{2}$ xn+1
$\mathbf{UNF}^{-1}$($\langle$ e1, e2$\rangle$) | = | $\langle$$\mathbf{UNF}^{-1}$(e1), $\mathbf{UNF}^{-1}$(e2)$\rangle$
$\mathbf{UNF}^{-1}$(proj${}_{T_{1};T_{2};T_{3}}$) | = |  (xk,$\ldots$,yp) where the xi are the variables of $T_{3}$
$\mathbf{UNF}^{-1}$(pairT;A×B) | = | $\langle$xn+1, xn+2$\rangle$
$\mathbf{UNF}^{-1}$(e1$\circ$e2) | = | fun (y1,$\ldots$,ym) $\rightarrow\;$ $\mathbf{UNF}^{-1}$(e1)($\mathbf{UNF}^{-1}$(e2)[$\forall i$,yi/xi])
$\mathbf{UNF}^{-1}$(JTvarΓ;i) | = | (x1,$\ldots$,xi-i,xi+xn+1,xi+1,$\ldots$,xn)
$\mathbf{UNF}^{-1}$(JTopΓ;m) | = | (x1,$\ldots$,xn,xn+1+$\partial_{1}$opn*xn+m+1,$\ldots$,xn+m+$\partial_{m}$opn*xn+m+1)
$\mathbf{UNF}^{-1}$(JTproj${}_{T_{1};T_{2};T_{3}}$) | = | (x1,$\ldots$,xk,0,$\ldots$,0,xk+p,$\ldots$,xn)
$\mathbf{UNF}^{-1}$(JTpairT;A×B) | = | (x1,$\ldots$,xn-1,$\pi_{1}$xn,$\pi_{2}$xn)
$\mathbf{UNF}^{-1}$(JTmap2Γ;x,y.e) | = | let G = (map2 * xn+3
|  |  (map2 (a,b.($\nabla_{\Gamma}$e1)[a/x.b/y]) xn+1 xn+2)) in
|  | ((x1,$\ldots$,xn)$\widehat{+}$(reduce $\widehat{+}$ $\widehat{0}$ G),
|  |  map2 (a,b.($\nabla_{\\{x\\}}$e1)[a/x]*b) xn+1 xn+3,
|  |  map2 (a,b.($\nabla_{\\{y\\}}$e1)[a/x]*b) xn+2 xn+3 )
$\mathbf{UNF}^{-1}$(JTreduceΓ;x,y.e1;e2) | = | let A0 = shift1R (scanl (x,y.e1) y1 xn+1) in
|  | let A1 = shift1L (map2
|  |  (a,b.($\nabla_{\\{x\\}}$e1)[a/x,b/y]) A0 xn+1) in
|  |  let A2 = map2 (a,b.($\nabla_{\\{y\\}}$)e1[a/x,b/y]) A0 xn+1 in
|  | let A3 = scanr * 1 A1 in
|  | (x1,$\ldots$,xn, map2 (x,y. x*y*xn+2) A2 A3)
Figure 10. Transformation from Target UNF to Target

##  5\. Complexity analysis

In this section, we introduce a simple cost model for our language and show
that, after partial evaluation, our reverse-mode transformation satisfies a
version of the cheap gradient principle (Theorem 5.4).

###  5.1. Partial evaluation and optimization

As can be seen from the examples,
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ introduces
a lot of functions of several arguments. Following the insight from Section
2.4, we do not want to keep all these costly lambda abstractions. The
transformation is designed in such a way that all the lambda abstractions are
given arguments, and we can use partial evaluation to beta-reduce all these
lambda abstractions. By inspecting each rule in Figure 5, this allows us to
prove by induction on the judgment of e:

###### Lemma 5.1.

Let $\Gamma\vdash$e: A be a term in Source. Every variable Y in
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e) which
is not of ground type has exactly one occurrence in the term.

From the previous lemma, we see that there is a linear usage of each
continuation variable $Y$. This is one key property ensured by $\mathbf{UNF}$.
We also note that, apart from the map2 case, this continuation variable is
always applied to almost the identity. More precisely, we have applications of
the form fun (x1,$\ldots$,xn) $\rightarrow\;$ Y(e1,$\ldots$,en) where the ei
are xi except for at most $k$ (independent of $n$) terms. In fact, we have
$k=2$, except for the map2 and reduce.

We can first use inlining, the first optimization rule given in Figure 12.
Using the invariant above, most of the ei below are variables. Without loss of
generality, assume the only terms ei which are not variables are en-1 and en.
We can then use forward substitution, the third optimization rule in Figure
12, on the other ei. In brief:

fun (x1,$\ldots$,xn) $\rightarrow\;$ (fun (y1,$\ldots$,yn) $\rightarrow\;$
(f1,$\ldots$,fn))(e1,$\ldots$,en) $\leadsto$
---
fun (x1,$\ldots$,xn) $\rightarrow\;$ let y1,$\ldots$,yn = e1,$\ldots$,en in
(f1,$\ldots$,fn) $\leadsto$
fun (x1,$\ldots$,xn) $\rightarrow\;$ let yn-1,yn = e${}_{n-1},$e1 in
(f1[x1/y1],$\ldots$,f${}_{n-2}[$xn-2/y${}_{n-2}]$,fn-1,fn)

This rewriting does not change the evaluation cost of the fi for $1\leq i\leq
n-2$. The new evaluation cost is reduced to the sum of the cost of evaluating
the fi in addition to the cost of evaluating en-1 and en, gaining O($n$)
movement of variables.

###### Example 5.2.

After forward-substitution and inlining the inner Yi, the gradient of the
terms from the introduction reduces to

|  $\nabla_{\Gamma}$(let w1 = x1 * x2 in let w2 = w1 * x1 in w2)
---|---
= | let w1,Y1= $\langle$x1 * x2, fun (y1,y2,y3,z) $\rightarrow\;$ Y(y1+x2*z,y2+x1*z,y3) $\rangle$ in
| let w2,Y2= $\langle$w1 * x1, fun (y1,y2,y3,y4,z) $\rightarrow\;$
Y1(y1+w1*z,y2,y3,y4+x1*z) $\rangle$ in
| let y,Y3= $\langle$w2, fun (y1,y2,y3,y4,z) $\rightarrow\;$
Y2(y1,y2,y3,y4+z)$\rangle$ in
| $\langle$y, fun (y1,y2,y3,z) $\rightarrow\;$ Y3(y1,y2,y3,0,z) $\rangle$

After another simplification step, we obtain

|  $\nabla_{\Gamma}$(let w1 = x1 * x2 in let w2 = w1 * x1 in w2)
---|---
= | let w1= x1 * x2 in let w2= w1 * x1 in
| $\langle$w2, fun (y1,y2,y3,z) $\rightarrow\;$ let y’1=y1+w1*z in’
|  let z’=y4+x1*z in Y(y’+x2*z’,y2+x1*z’,y3) $\rangle$

Similarly, for the gradient of prod(A) we obtain

$\nabla_{\Gamma}$(prod(A))
---
let A0= shift1R (scanl * 1 A) in
let A1= shift1L (map2 (a,b.b) A0 A) in
let A2= map2 (a,b.a) A0 A in
let A3= scanr * 1 A1 in
$\langle$prod(A), fun (X,z) $\rightarrow\;$ Y(X+map2 (a,b. a*z*b) A2
A3)$\rangle$

We call this optimization step partial evaluation in the rest of the
complexity section. We have the following result.

###### Lemma 5.3.

Let $\Gamma\vdash$ e:A be a term in Source. After the partial evaluation step,
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e) does
not have lambda abstractions. Its only variable (bound and free) which does
not have a ground type is Y.

###### Proof.

By induction on the judgment of e. We inspect each rule in Figure 5 and note
that the partial evaluation step precisely allows us to conclude the inductive
step. ∎

###  5.2. Cost model

We follow a simple model similar to the one in (Griewank and Walther, 2008).
We assume the cost is divided into $4$ elementary measures, being the number
of MOVES, ADDS, MULTS, and NLOPS. MOVES assumes a flat memory and represents
moving fixed-size information (e.g. 64 bits). ADDS represents the number of
additions, MULTS the number of multiplications, and NLOPS the number of
elementary non-linear operations like cos or exp.

The gives a complexity function $Cost$ valued in $\mathbb{R}^{4}$. For
primitive operations, we have for instance

$Cost$(*)=$(3,0,1,0)$ |  $Cost($+$)=(3,1,0,0)$
---|---
$Cost$(c)=$(1,0,0,0)$ |  $Cost($sin$)=(2,0,0,1)$

More generally $Cost($op1$)=(2,0,0,1)$ for the other unary operations. For
simplicity, we will not address any parallelism in our cost model. So the cost
of map (x.e) on an array of size $n$ will be $n*Cost$(e). Following Lemma 5.3,
we do not need our cost model to deal with higher-order variables and lambda
abstractions. We can thus restrict our attention to the subset of the Target
language which does not contain lambdas, applications or variables which are
not of ground type. The cost function extends compositionally to the
restricted Target language. It is given in Figure 11.

$Cost$(c) | = | (1,0,0,0)
---|---|---
$Cost$(x) | = | (1,0,0,0)
$Cost$(op1 e) | = | $Cost$(op1)+$Cost$(e)
$Cost$(e1 op2 e2) | = | $Cost$(op2)+$Cost$(e1)+$Cost$(e2)
$Cost$($\langle$e1,e2$\rangle$) | = | $Cost$(e1)+$Cost$(e2)
$Cost$($\pi_{i}$(e)) | = | (1,0,0,0)+$Cost$(e)
$Cost$(let x=e1 in e2) | = | (1,0,0,0)+$Cost$(e1)+$Cost$(e2)
$Cost$(map2 (x,y.e1) e2 e3) | = | n*($Cost$(e1)+(2,0,0,0))+$Cost$(e2)+$Cost$(e3)
$Cost$(reduce (x,y.e1) e2 e3) | = | n*($Cost$(e1)+(2,0,0,0))+$Cost$(e2)+$Cost$(e3)
$Cost$(scanl (x,y.e1) e2 e3) | = | n*($Cost$(e1)+(3,0,0,0))+$Cost$(e2)+$Cost$(e3)
$Cost$(scanr (x,y.e1) e2 e3) | = | n*($Cost$(e1)+(3,0,0,0))+$Cost$(e2)+$Cost$(e3)
$Cost$(shift1L e) | = | $Cost$(e)+(n,0,0,0)
$Cost$(shift1R e) | = | $Cost$(e)+(n,0,0,0)
Figure 11. Cost model for the restricted target language

###  5.3. Cheap gradient principle

We define the Nesting of Array Operations $NAO$ of a term e of Source by
induction on e as follows.

$NAO$(c), $NAO$(x) | = | 0
---|---|---
$NAO$($\pi_{i}$(e)), $NAO$(op1 e) | = |  $NAO$(e)
$NAO$(let x:A = e1 in e2:B) | = | max($NAO$(e1), $NAO$(e2))
$NAO$($\langle$ e1, e2 $\rangle$:AxB) | = | max($NAO$(e1), $NAO$(e2))
$NAO$(e1 op2 e2) | = | max($NAO$(e1), $NAO$(e2))
$NAO$(map2 (x,y.e1) e2 e3) | = | max(1+$NAO$(e1), $NAO$(e2), $NAO$(e3))
$NAO$(reduce (x,y.e1) e2 e3) | = | max(1+$NAO$(e1), $NAO$(e2), $NAO$(e3))

We can now phrase our main complexity theorem.

###### Theorem 5.4.

Given a term $\Gamma\vdash$ e: $\mathbf{R}$ such that $NAO$(e)$\leq$ $p$.
Denote by $G$ the term $\nabla_{\Gamma}$e after the partial evaluation step
from Section 5.1. Then $Cost$(G) $\leq$ $4*3^{p}*Cost$(e).

The cheap gradient principle (see e.g. (Griewank and Walther, 2008)) for
reverse-mode asserts that evaluating the gradient of a function
$e:\mathbb{R}^{n}\to\mathbb{R}$ should be the same order of cost as evaluating
$e$. More precisely, there should be a constant $K$ such that for each program
$\Gamma\vdash$ e: $\mathbf{R}$ in the context
$\Gamma=\\{x_{1}:\mathbf{R},\ldots,x_{n}:\mathbf{R}\\}$, $Cost$($\nabla$e)
$\leq$ $K$*$Cost$(e).

###### Proof Sketch.

As ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ is
defined by induction on programs, it suffices to show locally that
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ verifies
the cheap gradient principle.

There is a first restriction that prevents our transformation from satisfying
the cheap gradient principle. If we try to show that the cheap gradient
principle holds by induction on terms, it fails for map2. When differentiating
map2, there are three series of calls to (sub)gradients of e1. The induction
hypothesis is then too weak to conclude. The problem comes from the fact that
e1 could itself use map2. So the constant $K$ can be independent of $n$ and of
the size of the term if we allow it to be dependent on the level of nesting of
map2. A similar phenomenon happens with reduce.

Another problem is that the continuation part of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ adds
$O(n)$ MOVES at each step. This overhead is precisely what that our inlining
and forward substitution removes, as was exemplified in Section 5.1.

After that, the proof is by routine induction on e. First, one computes the
cost of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e). It has
$n$ too many MOVES which are removed by inlining and forward substitution
using the invariant that at each step, at most 2 xi in the continuation are
not variables. ∎

Inlining and forward substitution |  |
---|---|---
(fun (x1,$\ldots$,xn) $\rightarrow\;$ e)(e1,$\ldots$,en) | $\leadsto$ | let x1=e1 in $\ldots$
| let xn=en in e
let x1=e1 in e2 (x1$\not\in$FV(e2)) | $\leadsto$ | e2
let x1=x2 in e | $\leadsto$ | e[x2/x1]
let x1= c in e (c= 0,1) | $\leadsto$ | e[c/x1]
Algebraic simplifications |  |
0*e | $\leadsto$ | 0
0+e, 1*e | $\leadsto$ | e
Array algebraic simplifications |  |
map2 * A OnesLike(B) | $\leadsto$ |
map2 + A ZerosLike(B) | A
map (x.x) A |
map2 * A ZerosLike(B) | $\leadsto$ | ZerosLike(B)
reduce * 1 OnesLike(A) | $\leadsto$ | 1
reduce + 0 ZerosLike(A) | $\leadsto$ | 0
shift1L OnesLike(n+1) | $\leadsto$ | OnesLike(n)
shift1R OnesLike(n+1) |
Classic array simplification |  |
map (x.e1) (map2 (y1,y2.e2) A B) | $\leadsto$ | map2 (y1,y2.let x=e2 in e1) A B
Tuple partial evaluation |  |
$\pi_{i}$$\langle$e1,$\ldots$,en$\rangle$ | $\leadsto$ | ei
Let normalisation |  |
let x=(let y=e1 in e2) in e3 | $\leadsto$ | let y=e1 in let x=e2 in e3
f(let x=e1 in e2) | $\leadsto$ | let x=e1 in f(e2)
Conditionals |  |
if e1 then e2 else e2 | $\leadsto$ | e2
if true then e2 else e3 | $\leadsto$ | e2
if false then e2 else e3 | $\leadsto$ | e3
f(if e1 then e2 else e3) | $\leadsto$ | if e1 then f(e2) else f(e3)
Figure 12. Optimizations for target language.

###  5.4. Optimizations

In Figure 12 we present a list of optimizations. These optimizations are all
obviously valid according to the semantic model. As our language is purely
functional, we can use these extra optimizations aggressively. A lot of
simplifications come from the ring structure of the reals, lifted to tuples
and arrays. After these optimizations, the output program is essentially a
sequence of let-bindings and resembles SSA-code (Cytron et al., 1989).

Even though our reverse-mode transformation has the right complexity after the
partial evaluation step, the constant factor is of huge importance in
practice. The purity of our transformation allows us to make the most out of
generic optimizations. In addition, hand-crafted efficient derivatives and
more optimizations can easily be added to our language.

###### Example 5.5.

As shown in the supplementary material, the optimizations from Figure 12 are
sufficient to show that the gradients of the terms of the introduction reduce
to the following.

$\nabla_{A}$prod(A) | = | map2 * (scanr * 1 (shift1L A)) (shift1R (scanl * 1 A))
---|---|---
$\nabla_{A}$sum(A) | = | map (x $\rightarrow\;$ 1) A
$\nabla_{A}$dot(A,B) | = | B

##  6\. Correctness

We now explain that the correctness of the three steps of the
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$
transformation from Section 3.3 can be understood in terms of translating
between different standard categorical languages, and verified by working in a
well-known category for differentiability: diffeological spaces.

In general, categories with products (and more generally monoidal categories)
can be equivalently presented in the following three ways. Although the three
styles of presentation are categorically equivalent, the choice of
presentation affects the internal language syntax.

  * •

categories, where each arrow has one source and one target. These are most
familiar, but do not directly match with a programming syntax (except the
categorical abstract machine).

  * •

multicategories, where each arrow has a list of source objects
A${}_{1},\ldots,$An and one target B, thought of as a map A1
$\times\ldots\times$A${}_{n}\to$B. These are very close to the syntax usually
used in typed programming languages (Lambek, 1968; Staton and Levy, 2013),
thinking of an arrow as a typed term x1:A${}_{1},\ldots,$xn:An $\vdash$ B.

  * •

concategories, aka coloured props (Bonchi et al., 2015; Fong et al., 2019),
where each arrow has a list of source objects A${}_{1}\ldots$An and a list of
target objects B${}_{1}\ldots$Bm, thought of as a map
A${}_{1}\times\ldots\times$A${}_{m}\to$ B${}_{1}\times\ldots\times$Bm. These
match the syntax of the UNF language.

Informally, notice that it makes sense to talk about the opposite of a
category or a concategory, but not the opposite of a multicategory. For
roughly this reason, concategories are a more natural place to consider
reverse derivatives than multicategories, even though they are syntactically
less familiar. The
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$
translation passes from the multicategory for the source syntax to the
concategory for UNF, where reverse-mode differentiation is easier, and then
back to the multicategory for the target syntax. In this way, the translations
between Source/Target and Source UNF/Target UNF are merely changing
perspective from multicategories to concategories. We work with a specific
multicategory and concategory built from diffeological spaces and smooth maps,
so that we can verify the construction of reverse derivatives. Since
diffeological spaces support products, they can be presented as a category, a
multicategory and a concategory. Mathematically, the difference between these
presentations is almost trivial. But in terms of programming syntax, the
difference between presentations has a big effect, as can be seen from the
difference between the complex translation in Section 3 and the simple
translations in Section 4, which are much more readily verified.

###  6.1. Denotational semantics

A multicategory generalizes a category by allowing multimorphisms, that is,
morphisms from a list of objects to an object. Most categorical structures
from category theory can be phrased similarly in multicategories. It is
standard to give a denotational semantics of a first-order language in a
Cartesian category, and alternatively in a multicategory.

A term $x_{1}:A_{1},\ldots,x_{n}:A_{n}\vdash e:A$ is interpreted in a
multicategory as a morphism $\llbracket e\rrbracket:[\llbracket
A_{1}\rrbracket,\ldots,\llbracket A_{n}\rrbracket]\to\llbracket A\rrbracket$.
Substitution is interpreted as composition. We first consider a syntactic
model for a language, which consists of a free multicategory on some base
types and primitives. Our source and target languages induce syntactic
multicategories as follows.

###### Definition 6.1 (Syntactic multicategory for Source).

Let $SynSource$ be the multicategory whose objects are types of Source, and
where a morphism $[A_{1},\ldots,A_{n}]\to A$ is a term
$x_{1}:A_{1},\ldots,x_{n}:A_{n}\vdash e:A$ of Source modulo the
$\eta\beta$-laws. Composition is by substitution.

We similarly define $SynTarget$, the syntactic multicategory on the target
language.

$SynSource$ satisfies the following universal property: for every Cartesian
multicategory $\mathcal{C}$, and every object $F(\mathbf{R})\in\mathcal{C}$,
morphisms $F(\underline{c})\in\mathcal{C}(1;F(\mathbf{R}))$,
$F(op1)\in\mathcal{C}(F(\mathbf{R});F(\mathbf{R}))$,
$F(op2)\in\mathcal{C}(F(\mathbf{R}),F(\mathbf{R});F(\mathbf{R}))$, there is a
unique multifunctor $F:SynSource\to\mathcal{C}$ respecting the interpretation
and preserving all the categorical structure.

This allows us to give a simple semantics of Source in the multicategory of
Cartesian spaces and smooth maps between them.

###### Definition 6.2 ($\mathbf{CartSp}$).

Let $\mathbf{CartSp}$ be the Cartesian multicategory whose objects are
Euclidean spaces and whose morphisms $[A_{1},\ldots,A_{n}]\to B$ are smooth
functions $A_{1}\times\ldots A_{n}\to B$.

We interpret the source language in $\mathbf{CartSp}$ as follows. A context
$\Gamma=\\{x_{1}:A_{1},\ldots,x_{n}:A_{n}\\}$ is interpreted as the product
$\prod_{1\leq i\leq n}\llbracket A_{i}\rrbracket$. Well typed terms
$\Gamma\vdash$e:$A$ are interpreted as functions
$\llbracket\Gamma\rrbracket\to\llbracket A\rrbracket$.

$\llbracket\mathbf{R}\rrbracket$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $\mathbb{R}$
---|---|---
$\llbracket$T1xT2$\rrbracket$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ |  $\llbracket$T1$\rrbracket\times\llbracket$T2$\rrbracket$
$\llbracket\mathbf{A}[\mathbf{R}]^{n}\rrbracket$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $\prod_{1\leq i\leq n}\mathbb{R}$
$\llbracket$op1$\rrbracket$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | op1: $\mathbf{R}\to\mathbf{R}$
---|---|---
$\llbracket$op2$\rrbracket$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | op2: $\mathbf{R}\times\mathbf{R}\to\mathbf{R}$
$\llbracket$c$\rrbracket$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $c\in\mathbb{R}$

Variables are interpreted as projections $\pi_{i}$, let binding as composition
in the multicategory. Pairs are interpreted using the Cartesian structure of
the multicategory. This interpretation extends to map2 (x,y.e1) e2 e3. It is
given by $\langle$$\Delta_{n},$,$\llbracket$ e${}_{2}\rrbracket$,$\llbracket$
e${}_{2}\rrbracket$$\rangle$;swap;$(\llbracket$ e${}_{1}\rrbracket)^{\times
n}$ where $\Delta_{n}$ is the $n$-copy map $\langle$id,$\ldots$,id$\rangle$,
and swap a permutation. Similarly, the semantics for reduce (x,y.e1) e2 e3 is
then given by first $\langle$$\Delta_{n},$,$\llbracket$
e${}_{2}\rrbracket$,$\llbracket$ e${}_{2}\rrbracket$$\rangle$, followed by a
permutation
$\Gamma^{n}\times\mathbb{R}\times\mathbb{R}^{n}\to\mathbf{R}\times(\mathbf{R}\times\Gamma)^{n}$.
Finally, we apply $\llbracket$ e${}_{1}\rrbracket\times id$ $n$ times, where
the identity is of the obvious type at each stage.

We can similarly interpret the target language in a multicategory of smooth-
like spaces and functions. However, Target is higher-order and
$\mathbf{CartSp}$ is not Cartesian Closed. Instead, we can interpret Target in
the category of Diffeological spaces, as in (Huot et al., 2020). Diffeological
spaces ((Iglesias-Zemmour, 2013)) are a conservative extension of
$\mathbf{CartSp}$. The key idea will be that a higher-order function is called
smooth if it sends smooth functions to smooth functions, meaning that we can
never use it to build first-order functions that are not smooth.

###### Definition 6.3.

A _diffeological space_ $(X,\mathcal{P}_{X})$ consists of a set $X$ together
with, for each $n$ and each open subset $U$ of $\mathbb{R}^{n}$, a set
$\mathcal{P}_{X}^{U}\subseteq[U\to X]$ of functions, called _plots_ , such
that

  * •

all constant functions are plots;

  * •

if $f:V\to U$ is a smooth function and $p\in\mathcal{P}_{X}^{U}$, then
$f;p\in\mathcal{P}_{X}^{V}$;

  * •

if $\left(p_{i}\in\mathcal{P}_{X}^{U_{i}}\right)_{i\in I}$ is a compatible
family of plots $(x\in U_{i}\cap U_{j}\Rightarrow p_{i}(x)=p_{j}(x))$ and
$\left(U_{i}\right)_{i\in I}$ covers $U$, then the gluing $p:U\to X:x\in
U_{i}\mapsto p_{i}(x)$ is a plot.

We call a function $f:X\to Y$ between diffeological spaces _smooth_ if, for
all plots $p\in\mathcal{P}_{X}^{U}$, we have that $p;f\in\mathcal{P}_{Y}^{U}$.
We write $\mathbf{Diff}(X,Y)$ for the set of smooth maps from $X$ to $Y$.
Smooth functions compose, and so we have a category $\mathbf{Diff}$ of
diffeological spaces and smooth functions.

A diffeological space is thus a set equipped with structure. Many
constructions of sets carry over straightforwardly to diffeological spaces.
For instance, given a family $\left(X_{i}\right)_{i\in I}$ of diffeological
spaces, we can equip the product $\prod_{i\in I}X_{i}$ of sets with the
_product diffeology_ in which $U$-plots are precisely the functions of the
form $\left(p_{i}\right)_{i\in I}$ for $p_{i}\in\mathcal{P}_{X_{i}}^{U}$.
Cartesian spaces $\mathbb{R}^{n}$ can be given the structure of a
diffeological space by taking all the smooth functions $U\to\mathbb{R}^{n}$ as
$\mathcal{P}_{\mathbb{R}^{n}}^{V}$. We can equip the set $\mathbf{Diff}(X,Y)$
of smooth functions between diffeological spaces with the _functional
diffeology_ in which $U$-plots consist of functions $f:U\to\mathbf{Diff}(X,Y)$
such that $(u,x)\mapsto f(u)(x)$ is an element of $\mathbf{Diff}(U\times
X,Y)$. We can thus interpret function types $\llbracket$A $\rightarrow\;$
B$\rrbracket$ =
$\mathbf{Diff}$($\llbracket$A$\rrbracket$,$\llbracket$B$\rrbracket$).

###  6.2. Semantics for UNF languages with concategories

One main reason for introducing UNF is to have a better handle over the
computation flow of the term, in the same vein as the ANF or CPS
transformations. A convenient categorical setting for this is to use string
diagrams. To better fit the standard denotational semantics of languages, we
use concategories instead. Similar to the settings of categories and
multicategories, most categorical constructions used for the semantics of
functional languages have equivalent in concategories.

In particular, one can form a syntactic concategory on some base types and
primitives, which satisfies a similar universal property as syntactic
multicategories. Two particular examples are of interest to us, as they will
allow us to interpret Source UNF and Target UNF but they will also help us
explain the $\mathbf{UNF}$ and $UNF^{-1}$ transformations.

###### Definition 6.4 ($Concat1$).

Let $Concat1$ be the syntactic Cartesian concategory whose types are those of
Source and with primitives given by $op1:\mathbf{R}\to\mathbf{R}$,
$op2:\mathbf{R},\mathbf{R}\to\mathbf{R}$,
$map2_{x,y.e}:\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n}\to\mathbf{A}[\mathbf{R}]^{n}$
and $reduce_{x,y.e1;e2}:\mathbf{A}[\mathbf{R}]^{n}\to\mathbf{R}$.

One may notice that the syntax of $Concat1$ is somewhere in between the syntax
of Source and of Source UNF. Given two morphisms e1:T1$\rightarrow\;$T2 and
e2:T3$\rightarrow\;$T4, we denote by
e${}_{1}\mathmakebox[2ex][s]{+\kern-3.87498pt+\kern 3.09999pt}$
e2:T1,T3$\rightarrow\;$T2,T4 their parallel composition. We denote by $u_{T}$
the unique morphism from T to the terminal object $[]$. $pair_{T;A\times B}$
is the canonical isomorphism pairing A,B into A$\times$B in the context T. We
can interpret Source UNF in $Concat1$ as follows.

$\llbracket$var${}_{T;i}\rrbracket$ | = | $<id_{\Gamma},\pi_{i}>$
---|---|---
$\llbracket$op${}_{T;n}\rrbracket$ | = | $id_{\Gamma}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}op_{n}$
$\llbracket$pair${}_{T;A\times B}\rrbracket$ | = | $pair_{T;A\times B}$
$\llbracket$proj${}_{T1;T2;T3}\rrbracket$ | = | $id_{T1}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}u_{T2}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}id_{T3}$
$\llbracket$e1;e2$\rrbracket$ | = |  $\llbracket$e1$\rrbracket;\llbracket$e2$\rrbracket$
---|---|---
$\llbracket$e1$\widetilde{;}$e2$\rrbracket$ | = |  $\llbracket$e1$\rrbracket;(\llbracket$e2$\rrbracket\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}id);swap$
$\llbracket$map2${}_{T;x,y.e}\rrbracket$ | = | $id_{\Gamma}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}map2_{x,y.e}$
$\llbracket$reduce${}_{T;x,y.e;e}\rrbracket$ | = | $id_{\Gamma}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}reduce_{x,y.e1;e2}$

Similarly, we introduce a second concategory for the Target part of our
transformations.

###### Definition 6.5 ($Concat2$).

Let $Concat2$ be the syntactic Cartesian concategory whose types are those of
Target and whose primitives are given by $op1:\mathbf{R}\to\mathbf{R}$,
$op2:\mathbf{R},\mathbf{R}\to\mathbf{R}$,
$map2_{x,y.e}:\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n}\to\mathbf{A}[\mathbf{R}]^{n}$,
$reduce_{x,y.e1;e2}:\mathbf{A}[\mathbf{R}]^{n}\to\mathbf{A}[\mathbf{R}]^{n}$,
$J^{T}op1:\mathbf{R}\to\mathbf{R}$,
$J^{T}op2:\mathbf{R}\to\mathbf{R},\mathbf{R}$,
$J^{T}map2_{x,y.e}:\mathbf{A}[\mathbf{R}]^{n}\to\mathbf{A}[\mathbf{R}]^{n},\mathbf{A}[\mathbf{R}]^{n}$,
$J^{T}reduce_{x,y.e1;e2}:\mathbf{A}[\mathbf{R}]^{n}\to\mathbf{A}[\mathbf{R}]^{n}$

We can interpret Target UNF in $Concat2$ as follows. The part common with
Source UNF is interpreted in the same way as for the case of Source UNF.

$\llbracket$JTvar${}_{T;i}\rrbracket$ | = | $\pi_{T}\widehat{+}(0,\pi_{A},0)$
---|---|---
$\llbracket$JTmap2${}_{T;x,y.e}\rrbracket$ | = | $id_{T}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}map2_{T;x,y.e}$
$\llbracket$e1$\circ$e2$\rrbracket$ | = |  $\Lambda$($id_{\Gamma}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}\llbracket$e2$\rrbracket$; $\Lambda^{-1}$($\llbracket$e1$\rrbracket$))
$\llbracket$JTreduce${}_{T;x,y.e_{1};e_{2}}\rrbracket$ | = | $id_{T}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}reduce_{T;x,y.e_{1};e_{2}}$
$\llbracket$JTproj${}_{T1;T2;T3}\rrbracket$ | = | $id_{T_{1}}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}\widehat{0}_{T_{2}}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}id_{T_{3}}$
$\llbracket$JTop${}_{T;n}\rrbracket$ | = | $id_{T}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}J^{T}op_{n}$
---|---|---
$\llbracket$$\langle$e1, e2$\rangle$$\rrbracket$ | = |  $\langle$$\llbracket$e1$\rrbracket$, $\llbracket$e2$\rrbracket$$\rangle$
$\llbracket$JTpair${}_{T;A\times B}\rrbracket$ | = |  $id_{T}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern 3.44444pt}$$\langle$$\pi_{1},\pi_{2}$$\rangle$

where T,A $\vdash$ (0,$\pi_{A}$,0): T and the projection $\pi_{A}$ lands in
the $i$-th element of the list T.

###  6.3. Semantics of $\mathbf{UNF}$ transformations

We interpret the source language in a new multicategory, whose morphisms are
particular morphisms of $Concat1$. As Source UNF is itself interpreted in
$Concat1$, this gives us a way to compare terms in Source with terms in Source
UNF. This comparison of morphisms of $Concat1$ gives the $\mathbf{UNF}$
transformation.

###### Definition 6.6 (Mutlicategory from concat).

We define $\mathcal{C}_{Source}$ to be the multicategory with the same objects
as Source and
$\mathcal{C}_{Source}([A_{1},\ldots,A_{n}],B)=\\{f\in
Concat1([A_{1},\ldots,A_{n}],[A_{1},\ldots,A_{n},B],\forall
i.f;\pi_{A_{i}}=id_{A_{i}})$
The composition of $f_{i}:\underline{A}_{i}\to B_{i}$ with $g:\underline{B}\to
C$ is given $f_{1}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern
3.44444pt}\ldots\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern
3.44444pt}f_{n};(u_{\underline{B}_{1}}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern
3.44444pt}A_{1}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern
3.44444pt}\ldots\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern
3.44444pt}u_{\underline{B}_{n}}\mathmakebox[2ex][s]{+\kern-4.30554pt+\kern
3.44444pt}A_{n});g$. In other words each term $f_{i}$ forgets about its output
$\underline{B}_{i}$, and then we use the composition in the concategory.

We can interpret Source in $\mathcal{C}_{Source}$ as follows. The functor is
an identity on types. On morphisms, we interpret them as morphisms in the
concategory as follows. The terminal map is interpreted as the identity. The
operators op1,op2 by themselves. Crucially, the semantics of a let is simple
composition in the concategory. The semantics of a variable is the pairing of
the identity with a projection. The operation map2 (x,y.e) is interpreted as
the pairing of the identity and itself, and reduce (x,y.e1) $e_{2}$ as itself.

Interpreting Source in $\mathcal{C}_{Source}$ allows us to see terms of Source
as morphisms in the concategory $Concat1$, and to compare them to terms of
Source UNF which are already interpreted in $\mathcal{C}_{Source}$. The
following proposition can be shown by induction on the structure of the terms.

###### Proposition 6.7 (construction above gives UNF).

Let $\Gamma\vdash$ e:A be a term in Source. Seen as morphisms in
$\mathcal{C}_{Source}$,
$\llbracket$e$\rrbracket$=$\llbracket\mathbf{UNF}$(e)$\rrbracket$.

Dually, we form a concategory from the syntactic multicategory for Target.
Then we use the universal property of $Concat2$ to construct a functor from
$Concat2$ to this concategory. This allows us to compare the terms of Target
UNF and the terms of Target, and $\mathbf{UNF}^{-1}$ arises in this way.

###### Definition 6.8 (Concat from multicat).

A multicategory $\mathcal{C}$ naturally defines a concategory with the same
objects as $\mathcal{C}$ and with morphisms
$\underline{A}\to[B_{1},\ldots,B_{n}]$ being $n$ morphisms $\underline{A}\to
B_{i}$ of $\mathcal{C}$.

We thus consider the concategory $\mathcal{C}_{Target}$ from the syntactic
multicategory for Target. We can interpret $Concat2$ in
$\mathcal{C}_{Target}$. The functor is identity-on-objects, sends operations
op1,op2 to themselves. It sends Jacobian operations to terms of Target as
given by $\mathbf{UNF}^{-1}$ in Figure 10.

This interpretation is in essence $\mathbf{UNF}^{-1}$. The difference is that
to preserve typing, the semantic $\mathbf{UNF}^{-1}$ sends a non-Jacobian
primitive to a tuple, as in Example 4.1. This is highly inefficient, and the
syntactic $\mathbf{UNF}^{-1}$ additionally projects to the last element, where
the non-trivial information of the term is.

What remains to explain now is the reverse mode transformation between Source
UNF and Target UNF. We construct a functor
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}:Concat1\to Concat2$.
This functor computes reverse-mode derivatives. Because terms of Source UNF
are interpreted in $Concat1$, we observe the effect of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$ on them and show that
it matches the syntactic
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$ from Section 4.3.

###### Definition 6.9
(${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$ as a lax functor).

We define ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}:Concat1\to
Concat2$ as follows.
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([A1,$\ldots$,An])=[A1,$\ldots$,An,A${}_{1}\times\ldots\times$A${}_{n}\to\rho$].
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(op1)=$\langle$op1,$\pi_{last}\circ$
JTop1$\rangle$,
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(op2)=$\langle$op2,$\pi_{last}\circ$
JTop2$\rangle$,
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(map2
(x,y.e))=$\langle$map2 (x,y.e),$\pi_{last}\circ$ JTmap2 (x,y.e)$\rangle$,
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$(reduce${}_{x,y.e_{1},e_{2}}$)=$\langle$reduce${}_{x,y.e_{1},e_{2}}$,$\pi_{last}\circ$
JTreduce${}_{x,y.e_{1},e_{2}}$$\rangle$, where $\pi_{last}$ is the projection
to the last element (the continuation variable). It naturally extends to
sequential composition. It does not automatically extend to a multi-functor
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([A1,$\ldots$,Ak]$\mathmakebox[2ex][s]{+\kern-3.87498pt+\kern
3.09999pt}$[Ak+1,$\ldots$,An]) $\neq$
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([A1,$\ldots$,Ak])$\mathmakebox[2ex][s]{+\kern-3.87498pt+\kern
3.09999pt}{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([Ak+1,$\ldots$,An]).
Still, there is a map
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([A1,$\ldots$,Ak])$\mathmakebox[2ex][s]{+\kern-3.87498pt+\kern
3.09999pt}{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([Ak+1,$\ldots$,An])$\to{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$([A1,$\ldots$,Ak]$\mathmakebox[2ex][s]{+\kern-3.87498pt+\kern
3.09999pt}$[Ak+1,$\ldots$,An]). Internally, as a lambda term, it is given by
$(x_{1},\ldots,x_{k},Y_{1},(x_{k+1},\ldots,x_{n},Y_{s})\mapsto(x_{1},\ldots,x_{n},\lambda(y_{1},\ldots,y_{n})\to
Y_{1}(y_{1},\ldots,y_{k})\widehat{+}Y_{2}(y_{k+1},\ldots,y_{n})$. Here,
$\widehat{+}$ is the reverse-derivative of the copy map, which is known to be
fanout . It is not surprising to see it appear as
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$ is a semantic
functor, and it does not need to be efficient.

The design of the syntactic
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$ on the UNF language
is inspired by the semantic one between concategories, and it’s routine to
check that they match.

###### Proposition 6.10 (semantic of syntactic D matches D lax functor).

Given a well typed term T1 $\vdash$ e: T2 in Source UNF, we have

$\llbracket{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$e$\rrbracket$
=
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}\llbracket$e$\rrbracket$

In summary, we have the following picture (not a proper categorical diagram):

${SourceUNF}$${Concat1}$${Concat2}$${C_{Target}}$${Source}$${\mathcal{C}_{Source}}$${TargetUNF}$${Target}$$\scriptstyle{\llbracket-\rrbracket}$$\scriptstyle{{\scalebox{0.8}{$\scriptstyle\overleftarrow{\mathcal{D}}$}}_{\rho}}$$\scriptstyle{\mathbf{UNF}^{-1}}$$\scriptstyle{\mathbf{UNF}}$$\scriptstyle{\llbracket-\rrbracket}$

###  6.4. Correctness theorem

First, we start from the correctness of the syntactic
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$: Source UNF $\to$
Target UNF, which is easy to establish, and then propagate this information to
Source and Target via the $\mathbf{UNF}$ and $\mathbf{UNF}^{-1}$
transformations. The semantics brackets $\llbracket-\rrbracket$ in this
section are in diffeological spaces.

###### Proposition 6.11 (Correctness
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$).

For every term $\mathbf{R}^{\times n}\vdash$ e: $\mathbf{R}^{\times n+1}$ in
Source UNF,

$\pi_{2}\llbracket{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$e$\rrbracket(x_{1},\ldots
x_{n},Id_{\mathbb{R}^{n}})$=$J^{T}_{(x_{1},\ldots
x_{n})}\llbracket$e$\rrbracket$

This is routinely proved by induction, as the language is first-order. This
uses the fact that for every primitive constant $A$, $\llbracket
J^{T}A\rrbracket=J^{T}\llbracket A\rrbracket$.

Recall that the intuition from Source UNF is that it consists of terms of
Source that also return their context. From there, the intuition for the
Jacobian of a primitive in Target UNF is that it should be the Jacobian of the
corresponding term in Target. This is easily checked for scalar operations.
For the non-trivial cases, we have

###### Proposition 6.12.

$\mathbf{UNF}^{-1}$ preserves the semantics of Jacobians of map2 and reduce.

$\llbracket\mathbf{UNF}^{-1}(J^{T}$map2${}_{x,y.e})\rrbracket$ | = | $J^{T}\llbracket\mathbf{UNF}^{-1}($map2${}_{x,y.e})\rrbracket$
---|---|---
$\llbracket\mathbf{UNF}^{-1}(J^{T}$reduce${}_{x,y.e_{1};e_{2}})\rrbracket$ | = | $J^{T}\llbracket\mathbf{UNF}^{-1}($reduce${}_{x,y.e_{1};e_{2}})\rrbracket$

This is proved in the supplementary material.

From this, we now deduce that the composite transformation $\mathbf{UNF}$,
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}$, $\mathbf{UNF}^{-1}$
is correct in the sense that it produces a term that computes the gradient of
the original term.

###### Proposition 6.13.

If x${}_{1}:\mathbf{R}$,$\ldots$,x${}_{n}:\mathbf{R}$ $\vdash$ e: $\mathbf{R}$
then
$\llbracket\pi_{2}\mathbf{UNF}^{-1}({\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}(\mathbf{UNF}($e$)))\rrbracket(x_{1},\ldots
x_{n},Id_{\mathbb{R}^{n}})$=$J^{T}_{(x_{1},\ldots
x_{n})}$$\langle$$Id_{\mathbb{R}^{n}}$,$\llbracket$e$\rrbracket$$\rangle$

By inspecting what that composition of transformations does on the terms of
Source, we show that this indeed computes the same as the transformation
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ from
Section 3.3. This should not come as a surprise because the design of
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ was in
fact guided via this decomposition and intermediate representation.

###### Proposition 6.14.

$\llbracket\mathbf{UNF}^{-1}({\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}_{\rho}(\mathbf{UNF}($e$)))\rrbracket$
=
$\llbracket{\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}($e$)\rrbracket$

Combining the previous propositions, we have shown that
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$ is
correct.

###### Theorem 6.15.

For every $\Gamma\vdash$ e:$\mathbf{R}$, we have
$\llbracket\nabla_{\Gamma}e\rrbracket=\nabla_{\Gamma}\llbracket e\rrbracket$.

##  7\. Beyond the Source language

In this section, we generalize our language by relaxing the imposed
restrictions. First, we allow free variables for the function argument of
reduce in Section 7.1. Then we show how to support conditionals in Section
7.2. Finally, we provide a recipe for adding more constructs to the language
in Section 7.3. Further generalizations for more array operations, non-smooth
scalar operators, and general array support can be found in the supplementary
materials.

###  7.1. Lifting the restriction on reduce

Assume we allow $\Gamma$, x: $\mathbf{R}$, y: $\mathbf{R}$ $\vdash$ e:
$\mathbf{R}$ to be the function argument in reduce (x,y.e) v A. We need to add
a term depending on $\nabla_{\\{xi\\}}e$ to yi in the continuation. As writing
the derivative becomes quite cumbersome, we use a more compact notation using
arrays of tuples. Similarly to the monoid $\widehat{+},0_{\Gamma}$ defined in
Section 3.3, we use $\widehat{\times},1_{\Gamma}$ for obvious extension of the
monoid $(\mathbf{R},\times,1)$. The modified reverse-mode transformation is
shown in Figure 7.1.

${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(reduce (x,y.e1) e2 e3) | = |  $\ldots$ same as Figure 5 up until A_$3$
---|---|---
|  |  let B0=map2 (a,b $\rightarrow\;$($\nabla_{\Gamma}$e)[a/x,b/y] A0 A) in
|  | let B1=scanr $1_{\Gamma}$ $\widehat{\times}$ B0 in
|  | let B2=reduce $\widehat{+}$ $0_{\Gamma}$ B1 in
|  | $\langle$reduce (x,y.e1) y1 A, fun ($\boldsymbol{x}$),z) $\rightarrow\;$
|  | Y($\boldsymbol{x}\widehat{+}(z*$B2),0,map2 (a,b. a*b*z) A2 A3)

Figure 13. The reverse-mode AD transformation of reduce without restrictions
on its function argument.

###  7.2. Conditionals

Non-Solution 1. Adding conditionals to the source language can be done easily,
for instance via defining

${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(if e1 then e2 else e3) | = | if e1 then ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e2) else ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e3)
---|---|---

The problem is that this could break the complexity of reverse mode because of
the non-linear usage of $Y$, and makes everything harder to optimize.

Non-Solution 2. A slightly better option would be to define

${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(if e1 then
e2 else e3) =
---
let b=e1 in $\langle$if b then e2 else e3, fun (x1,$\ldots$,xn,z)
$\rightarrow\;$
Y(b*${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e2)(x1,$\ldots$,xn,z)+(1-b)*${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e3)(x1,$\ldots$,xn,z))$\rangle$

Now both derivatives of e2 and e3 are put together, and this might unlock some
optimizations, but there is still a non-linear usage of the continuation
variable Y.

Solution. If we know we want to compute the whole gradient of the expression,
we can define the translation as follows:

${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(if e1 then
e2 else e3) = let b=e1 in $\langle$if b then e2 else e3,
---
fun (x1,$\ldots$,xn,z) $\rightarrow\;$
Y((x1,$\ldots$,xn)$\widehat{+}$($\nabla_{\Gamma}$(e2)*b$\widehat{+}\nabla_{\Gamma}$(e3)*(1-b))*z)$\rangle$

This time, there is a linear usage of the continuation variable Y. We note
that adding conditionals does not break differentiability as long as there are
no non-trivial primitives $\mathbf{R}\to\mathbb{B}$. Non-smooth functions such
as the Rectified Linear Unit (ReLU) in machine learning, which can be defined
as $ReLU(x)\stackrel{{\scriptstyle\mathrm{def}}}{{=}}max(0,x)$ requires a non-
smooth primitive such as $>0:\mathbf{R}\to\mathbb{B}$. In several AD systems
such as TensorFlow, the condition is evaluated before differentiation is
applied, and so conditionals are never directly differentiated.

###  7.3. Recipe for Adding More Constructs

To add more constructs to our language, there is no need to go through UNF;
most of the backend of our work can be used as a black box. Non-smooth
variables (typically booleans or integers) should be considered as external to
the language, similarly to the way we treat indices $n$ of arrays.

Reverse AD. If the operator has type $A\to B$, then one should provide its
transpose Jacobian, a term $B\to A$. Such operators can often be unrolled to
first-order programs.

Correctness. To check the correctness of the given Jacobian, it suffices to
check that the semantics of the transpose Jacobian matches the Jacobian of the
unrolled program.

Complexity. To ensure the complexity guarantee of the whole transformation,
one needs to check that there is a linear usage of the continuation variable
$Y$ and that the cost of the proposed Jacobian is at most $k$ times the
complexity of the operator.

##  8\. Discussion and future work

Design space First, there is a tradeoff to reach between a general expressive
language and a domain specific one. The latter usually has more static
information and a specific representation that lends itself to better
optimizations. Then, many optimizations performed on AD implementations
consist in hand-crafted derivatives for useful operations like matrix-matrix
multiplication, dot-product, etc. They don’t seem to arise from theoretical
justifications, are error-prone, and can hardly fit with more general
optimizations. This makes these systems harder to prove correct. The problem
is thus to ensure provable correctness and pureness, while not compromising on
efficiency. In addition, we would like something easily, provably efficient. A
real-world implementation based on our work would of course optimize further,
potentially using hand-crafted operations as well, but it would be based on
solid grounds.

Higher-order functions Some recent work (Vákár, 2021; Sherman et al., 2021)
present reverse-mode in a higher-order language. (Vákár, 2021) uses
categorical semantics to show correctness of the reverse-mode transformation
and (Sherman et al., 2021) uses sophisticated higher-order primitives such as
root finding, max, argmax or integral. Their work focuses on computable reals,
which is hard to compare in terms of efficiency with our more standard
approach of AD. It is however quite difficult to prove a complexity result in
the higher-order setting. In addition, standard techniques of
defunctionalization struggle when higher-order is combined with recursion. Our
reverse-mode transformation does not support recursion, so it is currently
always possible to partially evaluate a higher-order program to our Source
language, seen as intermediate representation, then perform our reverse-mode
transformation.

Array primitives We focused on giving reverse derivatives for a small set of
array primitives. Other common primitives include filter, flatten, gather.
These functions can easily be added to our Source language once we provide a
reverse derivative for them. This is reminiscent of hand-crafted derivatives
present in large AD frameworks (usually hundreds), and these can already be
added in our language. As shown in 7.3, one mostly only needs to make sure
that the provided transpose Jacobian is correct.

##  9\. Related Work

Correctness of AD in functional languages. Several recent works (Huot et al.,
2020; Vákár, 2021, 2020; Brunel et al., 2019; Barthe et al., 2020; Mazza and
Pagani, 2021; Lee et al., 2020; Abadi and Plotkin, 2020) have focused on
correctness of AD in a purely functional setting, often leaving efficiency on
the side, especially for reverse-mode differentiation. We see our work as a
complement and a first bridge between these works and more practical
considerations of efficiency, which often require a lot more care than is
acknowledged in more theoretical works.

Usage of iteration mechanics. An immense effort in machine learning for the
past decade has been in finding good architectures, to limit computational
costs, avoid vanishing and exploding gradients, and have better building
blocks for large complicated systems than traditional layers of a neural
network. Different approaches such as Dynamic neural networks (Jin et al.,
2017; Wu et al., 2016), Recursive NN (Socher et al., 2011; Biancofiore et al.,
2017), Reccurent NN (Bahdanau et al., 2014; Luong et al., 2015), Tree LSTM
(Tai et al., 2015; Chen et al., 2016), Dynamic Recursive NN (Guo et al.,
2019), Top-down Tree LSTM (Zhang et al., 2015), and Recursion in DNN (Jeong et
al., 2018) have found that recursive data structures such as trees are good
candidates. We have emphasized here on differentiating fold-based recursion on
arrays for efficiency, but one should be able to adapt this to any algebraic
data type. It will be interesting to see if and how we recover efficient
purely functional backpropagation (as opposed to the imperative version of
(Wang et al., 2019)) on the proposed architectures, which is usually derived
by hand and one main goal of these papers.

Array Languages and AD. Given the enormous computation needs for state-of-the-
art large scale machine learning applications, which require extremely
efficient tensor computations and automatic differentiation for
backpropagation, combining array languages and automatic differentiation
(tensor calculus) in the best-fitted intermediate representation for
optimizations is of key interest and active research (Bernstein et al., 2020;
Laue et al., 2018, 2020). Advanced array programming is considered an
orthogonal problem to AD, and we focused our work on the differentiation
aspect.

Comparison to other recent papers. Table 1 does not reflect all the aspects of
AD. For instance, (Lee et al., 2020) studies in more detail non
differentiability, and (Sherman et al., 2021) the differentiability of highly
non-trivial higher-order and partial functions. Our work is somewhat close in
spirit to the idea of (Elliott, 2018) of compiling to categories. The idea of
using closures as back-propagators is receiving recent attention, as is
highlighted in (Vytiniotis et al., 2019; Wang et al., 2019). These ideas are
used in Julia Zygote (Innes et al., 2019), Swift AD (Wei, 2018), and recently
in (Paszke et al., 2021a). These seem closer to using control mechanisms than
having purely functional reverse-derivatives. Other aspects of AD are
discussed in recent surveys (van Merrienboer et al., 2018; Baydin et al.,
2017). (Krawiec et al., 2022) show how to do efficient reverse-mode AD for a
higher-order purely functional language, but at the cost of requiring a
monadic translation.

##  10\. Conclusion

We introduced a transformation on programs to compute provably efficient (§5)
gradients via reverse derivatives in a purely functional way (§3.3) on a
simple yet expressive language with functions on arrays (§3.1), combined with
standard functional optimizations (§12). We introduced a novel intermediate
representation, Unary Normal Form (§4) to decompose our translation into
simpler ones. We gave denotational semantics to our languages (§6), and we
proved the correctness of the reverse-mode translation (§6.4). We showed (§7)
how to lift the restrictions that we introduced on arrays and how to extend
our approach with other constructs such as conditionals.

###### Acknowledgements.

We have benefited from discussing this work with many people, including
Younesse Kaddar, Jesse Sigal, Matthijs Vákár, Emmanuel Arrighi, Sam Staton and
others. The first author is supported by a Royal Society University Research
Fellowship. The second author thanks Huawei for their support of the
distributed data management and processing laboratory at the University of
Edinburgh.

## References

  * (1)
  * Abadi and Plotkin (2020) Martín Abadi and Gordon D Plotkin. 2020.  A Simple Differentiable Programming Language. In _Proc. POPL 2020_. ACM.
  * Bahdanau et al. (2014) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2014\.  Neural machine translation by jointly learning to align and translate.  _arXiv preprint arXiv:1409.0473_ (2014).
  * Barthe et al. (2020) Gilles Barthe, Raphaëlle Crubillé, Ugo Dal Lago, and Francesco Gavazzo. 2020.  On the Versatility of Open Logical Relations: Continuity, Automatic Differentiation, and a Containment Theorem. In _Programming Languages and Systems_ , Peter Müller (Ed.). Springer, Springer International Publishing, Cham, 56–83.
  * Baydin et al. (2017) Atılım Günes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. 2017\.  Automatic differentiation in machine learning: a survey.  _The Journal of Machine Learning Research_ 18, 1 (2017), 5595–5637.
  * Baydin et al. (2016) Atılım Güneş Baydin, Barak A Pearlmutter, and Jeffrey Mark Siskind. 2016.  Diffsharp: An AD library for. NET languages.  _arXiv preprint arXiv:1611.03423_ (2016).
  * Bernstein et al. (2020) Gilbert Bernstein, Michael Mara, Tzu-Mao Li, Dougal Maclaurin, and Jonathan Ragan-Kelley. 2020\.  Differentiating a Tensor Language.  _arXiv preprint arXiv:2008.11256_ (2020).
  * Biancofiore et al. (2017) Fabio Biancofiore, Marcella Busilacchio, Marco Verdecchia, Barbara Tomassetti, Eleonora Aruffo, Sebastiano Bianco, Sinibaldo Di Tommaso, Carlo Colangeli, Gianluigi Rosatelli, and Piero Di Carlo. 2017.  Recursive neural network model for analysis and forecast of PM10 and PM2. 5.  _Atmospheric Pollution Research_ 8, 4 (2017), 652–659.
  * Bonchi et al. (2015) Filippo Bonchi, Pawel Sobocinski, and Fabio Zanasi. 2015\.  Full abstraction for signal flow graphs.  _ACM SIGPLAN Notices_ 50, 1 (2015), 515–526.
  * Brunel et al. (2019) Aloïs Brunel, Damiano Mazza, and Michele Pagani. 2019\.  Backpropagation in the simply typed lambda-calculus with linear negation.  _Proceedings of the ACM on Programming Languages_ 4, POPL (2019), 1–27.
  * Carpenter et al. (2015) Bob Carpenter, Matthew D Hoffman, Marcus Brubaker, Daniel Lee, Peter Li, and Michael Betancourt. 2015\.  The Stan math library: Reverse-mode automatic differentiation in C++.  _arXiv preprint arXiv:1509.07164_ (2015).
  * Chen et al. (2016) Qian Chen, Xiaodan Zhu, Zhenhua Ling, Si Wei, Hui Jiang, and Diana Inkpen. 2016\.  Enhanced lstm for natural language inference.  _arXiv preprint arXiv:1609.06038_ (2016).
  * Cockett et al. (2019) Robin Cockett, Geoffrey Cruttwell, Jonathan Gallagher, Jean-Simon Pacaud Lemay, Benjamin MacAdam, Gordon Plotkin, and Dorette Pronk. 2019.  Reverse derivative categories.  _arXiv preprint arXiv:1910.07065_ (2019).
  * Cruttwell et al. (2019) Geoff Cruttwell, Jonathan Gallagher, and Ben MacAdam. 2019\.  Towards formalizing and extending differential programming using tangent categories. In _Proc. ACT 2019_.
  * Cytron et al. (1989) Ron Cytron, Jeanne Ferrante, Barry K Rosen, Mark N Wegman, and F Kenneth Zadeck. 1989.  An efficient method of computing static single assignment form. In _Proceedings of the 16th ACM SIGPLAN-SIGACT symposium on Principles of programming languages_. 25–35.
  * Elliott (2017) Conal Elliott. 2017\.  Compiling to categories.  _Proceedings of the ACM on Programming Languages_ 1, ICFP (2017), 27.
  * Elliott (2018) Conal Elliott. 2018\.  The Simple Essence of Automatic Differentiation.  _Proc. ACM Program. Lang._ 2, ICFP, Article 70 (July 2018), 70:1–70:29 pages.
  * Fong et al. (2019) Brendan Fong, David Spivak, and Rémy Tuyéras. 2019\.  Backprop as functor: A compositional perspective on supervised learning. In _2019 34th Annual ACM/IEEE Symposium on Logic in Computer Science (LICS)_. IEEE, 1–13.
  * Griewank and Walther (2008) Andreas Griewank and Andrea Walther. 2008.  _Evaluating derivatives: principles and techniques of algorithmic differentiation_. Vol. 105.  Siam.
  * Guo et al. (2019) Qiushan Guo, Zhipeng Yu, Yichao Wu, Ding Liang, Haoyu Qin, and Junjie Yan. 2019\.  Dynamic Recursive Neural Network. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_. 5147–5156.
  * Huot et al. (2020) Mathieu Huot, Sam Staton, and Matthijs Vákár. 2020\.  Correctness of Automatic Differentiation via Diffeologies and Categorical Gluing.  _arXiv preprint arXiv:2001.02209_ (2020).
  * Iglesias-Zemmour (2013) Patrick Iglesias-Zemmour. 2013\.  _Diffeology_.  American Mathematical Soc.
  * Innes et al. (2019) Mike Innes, Alan Edelman, Keno Fischer, Chris Rackauckus, Elliot Saba, Viral B Shah, and Will Tebbutt. 2019.  A Differentiable Programming System to Bridge Machine Learning and Scientific Computing.  _arXiv preprint arXiv:1907.07587_ (2019).
  * Jeong et al. (2018) Eunji Jeong, Joo Seong Jeong, Soojeong Kim, Gyeong-In Yu, and Byung-Gon Chun. 2018.  Improving the expressiveness of deep learning frameworks with recursion. In _Proceedings of the Thirteenth EuroSys Conference_. 1–13.
  * Jin et al. (2017) Long Jin, Shuai Li, Hung Manh La, and Xin Luo. 2017\.  Manipulability optimization of redundant manipulators using dynamic neural networks.  _IEEE Transactions on Industrial Electronics_ 64, 6 (2017), 4710–4720.
  * Krawiec et al. (2022) Faustyna Krawiec, Simon Peyton Jones, Neel Krishnaswami, Tom Ellis, Richard A Eisenberg, and Andrew W Fitzgibbon. 2022.  Provably correct, asymptotically efficient, higher-order reverse-mode automatic differentiation.  _Proc. ACM Program. Lang._ 6, POPL (2022), 1–30.
  * Lambek (1968) Joachim Lambek. 1968\.  Deductive systems and categories.  _Mathematical Systems Theory_ 2, 4 (1968), 287–318.
  * Laue et al. (2018) Sören Laue, Matthias Mitterreiter, and Joachim Giesen. 2018\.  Computing higher order derivatives of matrix and tensor expressions. In _Advances in Neural Information Processing Systems_. 2750–2759.
  * Laue et al. (2020) Sören Laue, Matthias Mitterreiter, and Joachim Giesen. 2020\.  A simple and efficient tensor calculus. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 34. 4527–4534.
  * Lee et al. (2020) Wonyeol Lee, Hangyeol Yu, Xavier Rival, and Hongseok Yang. 2020\.  On correctness of automatic differentiation for non-differentiable functions.  _arXiv preprint arXiv:2006.06903_ (2020).
  * Liao et al. (2019) Hai-Jun Liao, Jin-Guo Liu, Lei Wang, and Tao Xiang. 2019\.  Differentiable programming tensor networks.  _Physical Review X_ 9, 3 (2019), 031041.
  * Luong et al. (2015) Minh-Thang Luong, Hieu Pham, and Christopher D Manning. 2015\.  Effective approaches to attention-based neural machine translation.  _arXiv preprint arXiv:1508.04025_ (2015).
  * MacLane (1965) Saunders MacLane. 1965\.  Categorical algebra.  _Bull. Amer. Math. Soc._ 71, 1 (1965), 40–106.
  * Mak and Ong (2020) Carol Mak and Luke Ong. 2020.  A Differential-form Pullback Programming Language for Higher-order Reverse-mode Automatic Differentiation.  _arXiv preprint arXiv:2002.08241_ (2020).
  * Manzyuk (2012) Oleksandr Manzyuk. 2012\.  A Simply Typed $\lambda$-Calculus of Forward Automatic Differentiation. In _Proc. MFPS 2012_.
  * Mazza and Pagani (2021) Damiano Mazza and Michele Pagani. 2021.  Automatic differentiation in PCF.  _Proceedings of the ACM on Programming Languages_ 5, POPL (2021), 1–27.
  * Paszke et al. (2017) Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. 2017\.  Automatic differentiation in pytorch.  (2017).
  * Paszke et al. (2021a) Adam Paszke, Daniel Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Matthew Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin. 2021a.  Getting to the Point. Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programming.  _arXiv preprint arXiv:2104.05372_ (2021).
  * Paszke et al. (2021b) Adam Paszke, Matthew J Johnson, Roy Frostig, and Dougal Maclaurin. 2021b.  Parallelism-preserving automatic differentiation for second-order array languages. In _Proceedings of the 9th ACM SIGPLAN International Workshop on Functional High-Performance and Numerical Computing_. 13–23.
  * Pearlmutter and Siskind (2008) Barak A Pearlmutter and Jeffrey Mark Siskind. 2008.  Reverse-mode AD in a functional framework: Lambda the ultimate backpropagator.  _ACM Transactions on Programming Languages and Systems (TOPLAS)_ 30, 2 (2008), 7.
  * Sabry and Felleisen (1993) Amr Sabry and Matthias Felleisen. 1993.  Reasoning about programs in continuation-passing style.  _Lisp and symbolic computation_ 6, 3 (1993), 289–360.
  * Shaikhha et al. (2019) Amir Shaikhha, Andrew Fitzgibbon, Dimitrios Vytiniotis, and Simon Peyton Jones. 2019.  Efficient differentiable programming in a functional array-processing language.  _Proceedings of the ACM on Programming Languages_ 3, ICFP (2019), 97.
  * Sherman et al. (2021) Benjamin Sherman, Jesse Michel, and Michael Carbin. 2021\.  $\backslash\lambda\\_S$: computable semantics for differentiable programming with higher-order functions and datatypes.  _Proceedings of the ACM on Programming Languages_ 5, POPL (2021), 1–31.
  * Socher et al. (2011) Richard Socher, Cliff C Lin, Chris Manning, and Andrew Y Ng. 2011\.  Parsing natural scenes and natural language with recursive neural networks. In _Proceedings of the 28th international conference on machine learning (ICML-11)_. 129–136.
  * Staton and Levy (2013) Sam Staton and Paul Blain Levy. 2013.  Universal properties of impure programming languages.  _ACM SIGPLAN Notices_ 48, 1 (2013), 179–192.
  * Tai et al. (2015) Kai Sheng Tai, Richard Socher, and Christopher D Manning. 2015\.  Improved semantic representations from tree-structured long short-term memory networks.  _arXiv preprint arXiv:1503.00075_ (2015).
  * Team (2017) The XLA Team. 2017\.  XLA – TensorFlow compiled. https://www.tensorflow.org/xla.
  * Vákár (2020) Matthijs Vákár. 2020\.  Denotational Correctness of Foward-Mode Automatic Differentiation for Iteration and Recursion.  _arXiv preprint arXiv:2007.05282_ (2020).
  * Vákár (2021) Matthijs Vákár. 2021\.  Reverse AD at Higher Types: Pure, Principled and Denotationally Correct.. In _ESOP_. 607–634.
  * Van Der Walt et al. (2011) Stefan Van Der Walt, S Chris Colbert, and Gael Varoquaux. 2011\.  The NumPy array: a structure for efficient numerical computation.  _Computing in science & engineering_ 13, 2 (2011), 22–30.
  * van Merrienboer et al. (2018) Bart van Merrienboer, Olivier Breuleux, Arnaud Bergeron, and Pascal Lamblin. 2018.  Automatic differentiation in ML: Where we are and where we should be going. In _Advances in neural information processing systems_. 8757–8767.
  * Vytiniotis et al. (2019) Dimitrios Vytiniotis, Dan Belov, Richard Wei, Gordon Plotkin, and Martin Abadi. 2019.  The Differentiable Curry.  (2019).
  * Wang et al. (2019) Fei Wang, Daniel Zheng, James Decker, Xilun Wu, Grégory M. Essertel, and Tiark Rompf. 2019.  Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator.  _Proc. ACM Program. Lang._ 3, ICFP, Article 96 (July 2019), 31 pages.
  * Wei (2018) Richard Wei. 2018\.  First-Class Automatic Differentiation in Swift: A Manifesto.
  * Wu et al. (2016) Di Wu, Lionel Pigou, Pieter-Jan Kindermans, Nam Do-Hoang Le, Ling Shao, Joni Dambre, and Jean-Marc Odobez. 2016.  Deep dynamic neural networks for multimodal gesture segmentation and recognition.  _IEEE transactions on pattern analysis and machine intelligence_ 38, 8 (2016), 1583–1597.
  * Zhang et al. (2015) Xingxing Zhang, Liang Lu, and Mirella Lapata. 2015.  Top-down tree long short-term memory networks.  _arXiv preprint arXiv:1511.00060_ (2015).

##  Appendix A Appendix

###  A.1. Operational semantics

In Figure 15 a small step call-by-value operational semantics for the
language. The evaluation contexts aregiven in Figure 14.

Evaluation contexts
---
$E$ | ::= | [] $\mid$ let x = E in e $\mid$ $\langle$E, e$\rangle$ $\mid$ $\langle$v, E$\rangle$ $\mid$ $\pi_{i}$(E) $\mid$ E op2 e $\mid$ v op2 e $\mid$ op1 E
|  |  $\mid$ map (x.e) E $\mid$ map2 (x,y.e) E e $\mid$ map2 (x,y.e) v E
|  |  $\mid$ foldl (x,y.e) E e $\mid$ foldl (x,y.e) v e
|  |  $\mid$ reduce (x,y.e) E e $\mid$ reduce (x,y.e) v e
|  |  $\mid$ scanl (x,y.e) E e $\mid$ scanl (x,y.e) v e
|  |  $\mid$ scanr (x,y.e) E e $\mid$ scanr (x,y.e) v e
|  |  $\mid$ shift1L E $\mid$ shift1R E
|  |  $\mid$ E(e$\ldots$e) $\mid$ e(v$\ldots$vEe$\ldots$e) $\mid$ if E then e else e
|  |  $\mid$ [v,$\ldots$,v,E,e$\ldots$,e]
Values
v | ::= |  c $\mid$ $\langle$v, v$\rangle$ $\mid$ true $\mid$ false $\mid$ fun (x1,$\ldots$,xn) $\rightarrow\;$ e $\mid$ [v,$\ldots$,v]
Figure 14. Evaluation contexts and values op1 c | $\leadsto$ | op1(c)
---|---|---
c op2 c’ | $\leadsto$ | op2(c,c’)
let x=v in e | $\leadsto$ | e[v/x]
$\pi_{i}$$\langle$v1, v2$\rangle$ | $\leadsto$ | vi
(fun (x1,$\ldots$,xn) $\rightarrow\;$ e)(v1$\ldots$vn) | $\leadsto$ | e[v1/x1,$\ldots$,vn/xn]
shift1L [v1,$\ldots$,vn+1] | $\leadsto$ | [v2,$\ldots$,vn+1]
shift1L [] | $\leadsto$ | []
shift1R [v1,$\ldots$,vn+1] | $\leadsto$ | [v1,$\ldots$,vn]
shift1R [] | $\leadsto$ | []
scanl (x,y.e) v [] | $\leadsto$ | [v]
scanl (x,y.e) v [v1,$\ldots$,vn] | $\leadsto$ | v::(scanl (x,y.e) e[v/x,v1/y]
|  | [v2,$\ldots$,vn])
scanr (x,y.e) v [] | $\leadsto$ | [v]
scanr (x,y.e) v [v1,$\ldots$,vn] | $\leadsto$ | (scanl (x,y.e) e[v/x,v1/y]
|  | [v2,$\ldots$,vn])::v
reduce (x,y.e) v [] | $\leadsto$ | v
reduce (x,y.e) v [v1,$\ldots$,vn] | $\leadsto$ | reduce (x,y.e) e[v/x,v1/y] [v2,$\ldots$,vn])
map2 (x,y.e) [v11,$\ldots$,v1n] | $\leadsto$ | [e[v11/x,v21/y],$\ldots$,e[v1n/x,v2n/y]]
[v21,$\ldots$,v2n] |
if true then e1 else e2 | $\leadsto$ | e1
if false then e1 else e2 | $\leadsto$ | e2
Figure 15. Operational semantics of the source and target languages

###  A.2. Reverse derivative of array operations

We now prove that the reverse mode transformation is correct on array
operations.

###### Proposition A.1.

The reverse derivative of map2 is correct.

###### Proof.

Let us denote by t the term map2 (x,y.e) A B. Without loss of generality,
assume A=$[a_{1},\ldots,a_{n}]$ and B=$[b_{1},\ldots,b_{n}]$. By chosing Z to
a hot vector at i, call it Zi, we’re back to showing the result for the term
$\Gamma\vdash$e[ai/x,bi/y]: $\mathbf{R}$. In
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(t)(Zi),
the term G = (map2 * Z (map2 (a,b.($\nabla_{\Gamma}$e1)[a/x.b/y]) A B))
reduces to $\frac{\partial e}{\partial x_{i}}$[ai/x,bi/y]. As ai,bi are
independant of $x_{i}$, this term is equal to $\frac{\partial
e[a_{i}/x,b_{i}/y]}{\partial x_{i}}$, as expected. Similarly, the term map2 *
(map2 (a,b.($\nabla_{\\{x\\}}$e1)[a/x,b/y]) A B) Z reduces to $\frac{\partial
e}{\partial x}$[ai/x,bi/y]. As ai is a variable, and independant of bi, this
is equal to $\frac{\partial e[a_{i}/x,b_{i}/y]}{\partial a_{i}}$. ∎

###### Lemma A.2.

if op2 is an associative binary operation with unit $\Gamma\vdash$ v:
$\mathbf{R}$, then $\frac{\partial op2(v,e)}{\partial
y_{1}}\times\frac{\partial v}{\partial x_{i}}=0$ and $\frac{\partial
op2(e,v)}{\partial y_{2}}\times\frac{\partial v}{\partial x_{i}}=0$ for all
$x_{i}$.

###### Proof.

For any $\Gamma\vdash$ e: $\mathbf{R}$, we have v op2 e=e. Differentiating and
using the chain rule we get

$\frac{\partial op2(v,e)}{\partial y_{1}}\times\frac{\partial v}{\partial
x_{i}}+\frac{\partial op2(v,e)}{\partial y_{2}}\times\frac{\partial
e}{\partial x_{i}}=\frac{\partial e}{\partial x_{i}}$

As $\frac{\partial e}{\partial x_{i}}$ is arbitrary, this shows that
$\frac{\partial op2(v,e)}{\partial y_{2}}=1$ and $\frac{\partial
op2(v,e)}{\partial y_{1}}\times\frac{\partial v}{\partial x_{i}}=0$. Similarly
for the other case. ∎

###### Proposition A.3.

The reverse derivative of reduce is correct.

###### Proof.

The notation for the general case are cumbersome and non-insightful. We will
exemplify the proof on the case of an array of size 3. We use infix notation
for the binary operation $e$. The output term unrolls to
$v_{3}\stackrel{{\scriptstyle\mathrm{def}}}{{=}}((v~{}e~{}a_{1})~{}e~{}a_{2})~{}e~{}a_{3}$,
where $v$ is the unit of $e$. By the lemma above we know that the derivative
w.r.t $v$ is 0, and focus on the partial derivatives w.r.t $a_{i}$. Write
$v_{1}\stackrel{{\scriptstyle\mathrm{def}}}{{=}}(v~{}e~{}a_{1})$,
$v_{2}\stackrel{{\scriptstyle\mathrm{def}}}{{=}}v_{1}~{}e~{}a_{2}$. By
inspection we have

  * •

$\frac{\partial v_{3}}{\partial a_{3}}=\frac{\partial e}{\partial
x_{2}}(v_{2},a_{3})$

  * •

$\frac{\partial v_{3}}{\partial a_{2}}=\frac{\partial v_{3}}{\partial
x_{1}}\times\frac{\partial v_{2}}{\partial x_{2}}=\frac{\partial e}{\partial
x_{1}}(v_{2},a_{3})\times\frac{\partial e}{\partial x_{2}}(v_{1},a_{2})$

  * •

$\frac{\partial v_{3}}{\partial a_{1}}=\frac{\partial v_{3}}{\partial
x_{1}}\times\frac{\partial v_{2}}{\partial x_{2}}\times\frac{\partial
v_{1}}{\partial x_{2}}=\frac{\partial e}{\partial
x_{1}}(v_{2},a_{3})\times\frac{\partial e}{\partial
x_{1}}(v_{1},a_{2})\times\frac{\partial e}{\partial x_{2}}(v,a_{1})$

We thus need the following intermediate results

  * •

$A_{0}=[v,v_{1},v_{2}]$=shift1R(scanl v e A)

  * •

$A_{1}=[\nabla_{\\{x_{1}\\}}e(v_{1},a_{2}),\nabla_{\\{x_{1}\\}}e(v_{2},a_{3})]=$
shift1L (map2 (a,b.$\nabla_{\\{x_{1}\\}}$e(a/x,b/y)) $A_{0}$ A)

  * •

$A_{2}=[\nabla_{\\{x_{2}\\}}e(v,a_{1}),\nabla_{\\{x_{2}\\}}e(v_{1},a_{2}),\nabla_{\\{x_{2}\\}}e(v_{2},a_{3})]=$
map2 (a,b.$\nabla_{\\{x_{2}\\}}$e(a/x,b/y)) $A_{0}$ A

  * •

$A_{3}=[\frac{\partial e}{\partial x_{1}}(v_{2},a_{3})\times\frac{\partial
e}{\partial x_{1}}(v_{1},a_{2}),\frac{\partial e}{\partial
x_{1}}(v_{2},a_{3}),1]=$ scanr 1 * $A_{1}$

And we return $A_{4}=$ map2 $A_{2}$ $A_{3}$. ∎

###  A.3. Adding more array operators

There are two main differences with fold left foldl compared to reduce. First,
in foldl (x,y.e) v A the starting accumulation element v is not a unit for
(x,y.e), so it will have non-trivial derivatives in general and we need to
account for that. Second, we will need a more general scanl which allows
(x,y.e) as a function argument. In other words, we need the general scan left
computing the intermediate values of a fold left. This should be a different
primitive but we will still call it scanl in this section.

Finally, the former point implies we need a few more array manipulations.
These can be elegantly dealt with by changing the semantics of scanl and
scanr. Let’s assume that they now return a pair of an array of size $n$ of the
intermediate computations and the final result. The reverse derivative of
foldl is then as shown in Figure A.3. In this figure, we have shown the
translation of the map operator as well, which is very similar to map2.

${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(foldl (x,y.e1) e2 e3) | = | let v,Y1 = ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e2) in
---|---|---
|  | let A,Y2 = ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma,v;Y_{1}}$(e3) in
|  | let A0,r1 = (scanl (x,y.e1) v A) in
|  |  let A1 = map2 (a,b.($\nabla_{\\{x\\}}$e1)[a/x,b/y]) A0 A
|  |  let A2 = map2 (a,b.($\nabla_{\\{y\\}}$)e1[a/x,b/y]) A0 A
|  | let r2, A3 = scanr * 1 A1
|  | $\langle$r1, fun (x1,$\ldots$,xn,z) $\rightarrow\;$
|  | let y1,B = r2*z ,map2 (x,y. x*y*z) A2 A3 in
|  | Y2(x1,$\ldots$,xn,y1,B)$\rangle$
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(map (x.e1) e2) | = | let A,Y1 = ${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e2) in
|  | $\langle$map (x.e1) A, fun (x1,$\ldots$,xn,Z) $\rightarrow\;$
|  | let G = (map2 * Z
|  |  (map (a.($\nabla_{\Gamma}$e1)[a/x]) A)) in
|  | Y1( (x1,$\ldots$,xn)$\widehat{+}$(reduce $\widehat{+}$ $\widehat{0}$ G),
|  |  map2 (a,b.($\nabla_{\\{x\\}}$e1)[a/x]*b) A Z )$\rangle$

Figure 16. The reverse-mode AD transformation of foldl and map operators.

###  A.4. Adding non-smooth scalar operators

We assumed the unary and binary operators were denoted by smooth functions
$\mathbb{R}^{n}\to\mathbb{R}$. There is no additional difficulty in
considering operators which are partial functions like division or operators
which are not smooth at a point like square root.

These functions are then given intentional derivatives which provide valid
derivatives on the domain of definition and differentiability of the operator.
These functions are well known to be the bete noire of AD (Griewank and
Walther, 2008) and we do not provide novel solutions to these. Several recent
work have shown how to give semantics to such operators in the context of AD
(Vákár, 2020; Mazza and Pagani, 2021; Sherman et al., 2021; Lee et al., 2020).

###  A.5. General arrays

We now show how to generalize our reverse-mode transformation to be defined on
arrays over any ground type G. That is, we need to adapt the reverse
derivatives of map2 and reduce when they have more general function arguments.

A ground type G is interpreted as an Euclidean space $A$. It is in particular
a real vector space. Similarly, a ground context $\Gamma$=x1:G1,$\ldots$,xn:Gn
is interpreted as $\bigoplus_{1\leq i\leq n}A_{i}$, where
$A_{i}\stackrel{{\scriptstyle\mathrm{def}}}{{=}}\llbracket$(Gi)$\rrbracket$.
The denotation of the gradient of a term $\Gamma\vdash$ e: G at a point is
then a matrix, more precisely an element of $(\bigoplus_{1\leq i\leq
n}A_{i})\otimes A$ where $\otimes$ is the tensor product of real vector
spaces. This space is isomorphic to $\bigoplus_{1\leq i\leq n}A_{i}\otimes A$.

We can define $\otimes$ on the types of our language inductively by

$\mathbf{R}\otimes A$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $A$
---|---|---
$(A_{1}\times\ldots\times A_{n})\otimes A$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $(A_{1}\otimes A)\times\ldots\times(A_{n}\otimes A)$
$\mathbf{A}[A_{1}]^{n}\otimes A$ | $\stackrel{{\scriptstyle\mathrm{def}}}{{=}}$ | $\mathbf{A}[A_{1}\otimes A]^{n}$

With this definition, we recover that the gradient of
$\Gamma\vdash$e:$\mathbf{R}$ is a tuple of type $A_{1}\times\ldots\times
A_{n}$ as expected. For map2, we need a generalization of
$*:\mathbf{R}\to\mathbf{R}$. If e1: A, then $\nabla_{\\{x\\}}$e1: $A\otimes A$
and we need a new primitive $\widehat{*}:(A\otimes A)\times A\to A$. If we
represent $A\otimes A$ as a matrix, then $\widehat{*}$ is matrix-vector
multiplication. Similarly, for A3 in reduce we need a new primitive
$\widetilde{*}:(A\otimes A)\times(A\otimes A)\to(A\otimes A)$. $\widetilde{*}$
corresponds to matrix-matrix multiplication. Then, using this notation, there
are only minimal changes to the reverse derivatives of map2 and reduce, as can
be seen in Figure A.5.

${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(map2 (x,y.e1: G) e2 e3) | = |  let A,Y1 =${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e2) in
---|---|---
|  |  let B,Y2 =${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma,A;Y1}$(e3) in
|  | $\langle$map2 (x,y.e1) A B, fun (x1,$\ldots$,xn,Z) $\rightarrow\;$
|  | let G = (map2 $\widehat{*}$ Z
|  |  (map (a.($\nabla_{\Gamma}$e1)[a/x]) A)) in
|  |  Y2( (x1,$\ldots$,xn)$\widehat{+}$(reduce $\widehat{+}$ $\widehat{0}$ G),
|  |  map2 (a,b.($\nabla_{\\{x\\}}$e1)[a/x]$\widehat{*}$b) A Z,
|  |  map2 (a,b.($\nabla_{\\{y\\}}$e1)[a/x]$\widehat{*}$b) B Z )$\rangle$
${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(reduce (x,y.e1) e2 e3) | = |  let y1,Y1 =${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma;Y}$(e2) in
|  |  let A,Y2 =${\scalebox{0.8}{$\overleftarrow{\mathcal{D}}$}}^{\rho}_{\Gamma,y_{1};Y_{1}}$(e3) in
|  | let A0 = shift1R (scanl (x,y.e1) y1 A) in
|  | let A1 = shift1L (map2
|  |  (a,b.($\nabla_{\\{x\\}}$e1)[a/x,b/y]) A0 A) in
|  | let A2 = map2
|  |  (a,b.($\nabla_{\\{y\\}}$)e1[a/x,b/y]) A0 A in
|  | let A3 = scanr $\widetilde{*}$ 1 A1 in
|  | $\langle$reduce (x,y.e1) y1 A, fun (x1,$\ldots$,xn,z) $\rightarrow\;$
|  | Y2(x1,$\ldots$,xn, map2 (x,y. x$\widehat{*}$(y$\widehat{*}$z)) A2 A3$\rangle$

Figure 17. The reverse-mode AD transformation of map2 and reduce for general
arrays.

Evidently, one can combine all the generalizations from the previous
subsections. Even though this transformation has the correct complexity, it is
open for future research to find even better representations to allow for more
optimizations. In particular, representations looking like Einsum (Van Der
Walt et al., 2011) could be of interest and has been recently studied in the
context of AD (Laue et al., 2018, 2020). More generally, there is growing
interest in tensor calculus (Liao et al., 2019; Bernstein et al., 2020).

###  A.6. Gradient from the introduction

We show that the gradients from Section 1 are obtained as instances of our
general construction. The proofs consist in instantiating the general
derivatives to these cases and showing that each rewrite step is a simple
known optimization.

Similarly to numpy, we use the notation OnesLike(A) to mean map (x -¿ 1) A and
ZerosLike(A) to mean map (x -¿ 0) A. Finding these constant arrays is key to a
lot of optimizations that leverage the ring algebraic structure of the reals