#  $\nabla$SD: Differentiable Programming for Sparse Tensors

Amir Shaikhha<EMAIL_ADDRESS>University of EdinburghUnited Kingdom ,
Mathieu Huot<EMAIL_ADDRESS>University of OxfordUnited Kingdom and
Shideh Hashemian<EMAIL_ADDRESS>Amirkabir University of
TechnologyIran

###### Abstract.

Sparse tensors are prevalent in many data-intensive applications, yet existing
differentiable programming frameworks are tailored towards dense tensors. This
presents a significant challenge for efficiently computing gradients through
sparse tensor operations, as their irregular sparsity patterns can result in
substantial memory and computational overheads. In this work, we introduce a
novel framework that enables the efficient and automatic differentiation of
sparse tensors, addressing this fundamental issue. Our experiments demonstrate
the effectiveness of the proposed framework in terms of performance and
scalability, outperforming state-of-the-art frameworks across a range of
synthetic and real-world datasets. Our approach offers a promising direction
for enabling efficient and scalable differentiable programming with sparse
tensors, which has significant implications for numerous applications in
machine learning, natural language processing, and scientific computing.

Sparse Tensor Algebra, Automatic Differentiation, Semi-Ring Dictionaries

††copyright: none

##  1\. Introduction

Sparse tensors are essential in many scientific and engineering applications,
such as natural language processing, computer vision, and graph analytics.
Unlike dense tensors, which store all of their elements regardless of their
value, sparse tensors only store non-zero values, resulting in significant
memory savings and computational efficiency. Sparse tensors also enable
efficient representation and manipulation of high-dimensional data structures,
which are often encountered in modern machine learning and scientific
computing, such as sparse tensors representing the frequency of words in a
document or corpus in natural language processing, adjacency matrices of large
and sparse graphs in network/relational analysis, or sparse user-item
interaction matrices for collaborative filtering in recommender systems. This
has inspired recent interest in developing better support for sparse tensors
(Strout et al., 2018; Kjolstad et al., 2017; Tang et al., 2020).

Automatic differentiation (AD) is a fundamental technique in modern machine
learning and scientific computing that enables efficient computation of the
gradient of a function. This is crucial for optimization, parameter
estimation, and many other applications in which gradient-based optimization
methods are employed. While AD tools for dense tensors are abundant and well-
established, the lack of efficient AD tools for sparse tensors hinders the
wider adoption of these techniques and represents a major research challenge
in this field. Libraries such as TensorFlow, PyTorch, and JAX provide
efficient and scalable implementations of gradient computation for dense
tensor operations, but their support for sparse operations is limited
(TensorFlow, 2023; JAX, 2023; PyTorch, 2023). As a result, there have been
various efforts to manually provide differentiation for particular sparse
tensor kernels (Nytko et al., 2022).

AD for sparse tensor algebra is more challenging than for dense tensor algebra
for several reasons. Firstly, the structure of sparse tensors is more complex
than that of dense tensors, and their sparsity patterns are often irregular
and vary across different operations. This can make it challenging to
efficiently propagate gradients through the computation graph and to identify
which elements of the sparse tensor are relevant for the gradient computation.
Secondly, sparse tensor operations often require the use of specialized data
structures and algorithms, such as compressed sparse row (CSR) or compressed
sparse column (CSC) formats, which are not natively supported by most AD-
enabled frameworks or have incomplete AD rules.

This paper presents $\nabla$SD, a novel differentiable programming framework
that supports the automatic differentiation of arbitrary sparse computations.
To the best of our knowledge, $\nabla$SD is the first framework that provides
this capability. As opposed to the existing frameworks that offer AD support
for a limited number of sparse kernels (Nytko et al., 2022; Abadi et al.,
2015; Paszke et al., 2019; Bradbury et al., 2018), $\nabla$SD allows the AD of
an arbitrary sparse computation expressible in tensor algebra.

The key insight is to perform differentiation over a logical representation of
a sparse tensor. This means that there is a clear separation of concerns
between the semantics of differentiation over a program on the one hand, and
optimizations and data layout representations on the other (cf. Figure 1). To
see the simple example of the dot product of two vectors, see Example 1
(cont.) for its logical form before differentiation in §2.2, after
differentiation in §4.2, its physical coordinate format (COO) representation
in §5.3, and in §5.5 for the generated code.

The physical representation of sparse tensors (e.g., COO/CSR/CSC) involves
multiple arrays storing a compressed representation of the matrix (cf. Figure
3(a)). The computations over such representations involve imperative nested
loops over these arrays. However, our logical representation uses a nested
dictionary, where sparse computations are expressed functionally as nested
summations over them. The logical representation can be later fused with a
concrete physical storage format (cf. Section 5.3).

In more detail, the contributions of this paper are as follows:

Figure 1. The high-level overview of transformations in $\nabla$SD.

  * •

We present $\nabla$SD, the first framework with systematic support for the
automatic differentiation of sparse tensors. $\nabla$SD is based on SDQLite
(Shaikhha et al., 2022; Schleich et al., 2022), a recently introduced
intermediate language that can express sparse tensor workloads by separating
the tensor computations from the storage specifications (Section 3).

  * •

We introduce a novel tensorized forward-mode AD that computes the gradients in
a batch (Section 4). Our automatic differentiation transform is over the
logical part of the language, which we call Logical SDQLite, without worrying
about the physical storage formats.

  * •

The differentiated program is then optimized by leveraging the following AD-
agnostic transformations: (1) sparsity-aware rewrite rules, (2) composing with
the physical storage formats, and (3) algebraic rewrite rules applied in a
cost-based manner using equality saturation (Tate et al., 2009; Willsey et
al., 2021) (Section 5). In addition, $\nabla$SD performs low-level
transformations for removing unnecessary intermediate tensors appearing in
nested loops before generating low-level C++ code.

  * •

We show the correctness of our approach (Section 6). That is, AD is a well-
typed transformation, computes derivatives of programs, and our optimizations
are sound with respect to our denotational semantics.

  * •

We experimentally validate the effectiveness of $\nabla$SD in comparison with
the state-of-the-art frameworks (Section 7). We demonstrate that $\nabla$SD
scales the gradient computation to large matrices with many zero elements over
both real-world and synthetic datasets.

##  2\. Background

###  2.1. Automatic Differentiation

Automatic differentiation (AD) is a powerful and widely used technique in
machine learning and scientific computing that enables efficient computation
of the gradient of a function. The gradient is a crucial quantity in many
optimization, parameter estimation, and machine learning algorithms, and its
computation is often a bottleneck in the training process. AD provides a
computationally efficient and accurate way of computing the gradient by
breaking down a function into a series of elementary operations and applying
the chain rule to compute the derivatives of each operation. The result is an
exact gradient that is computed with a similar computational cost as the
original function, with no need for approximate methods such as finite
differences, or manual derivation.

Forward Mode. One method of computing the gradient of a function is the
forward-mode AD which involves computing the derivatives of each operation in
the forward direction through the computational graph. The method starts with
the input variables and propagates the values and their derivatives through
the graph, one operation at a time, until the output variable is reached. At
each operation, the derivative of the output variable with respect to each
input variable is computed using the chain rule, and these derivatives are
stored in a computational graph that can be used to compute the gradient of
the function.

Example 1 - Vector Dot Product. Consider the following function:

$f([x_{1},x_{2}],[y_{1},y_{2}])=x_{1}y_{1}+x_{2}y_{2}$

which takes two pairs of input variables $x_{1}$, $x_{2}$ and $y_{1}$, $y_{2}$
and computes the dot product of the vectors $x=[x_{1},x_{2}]$ and
$y=[y_{1},y_{2}]$. To compute the gradient of $f$ with respect to the inputs
using forward-mode AD, we start by converting the program into ANF (Flanagan
et al., 1993):

$f([x_{1},x_{2}],[y_{1},y_{2}])=$ |  let $t_{1}=x_{1}y_{1}$
---|---
|  let $t_{2}=x_{2}y_{2}$
|  let $t_{3}=t_{1}+t_{2}$
| $t_{3}$

The forward-mode AD lifts every variable to a dual number by associating a
tangent variable $v^{\prime}$ to each input and intermediate variable $v$.
Then, each intermediate tangent variable is computed by following the chain
rule. In the previous example, the function $f$ is transformed as follows:

$f^{\prime}([x_{1},x_{2}],[y_{1},y_{2}],[x_{1}^{\prime},x_{2}^{\prime}],[y_{1}^{\prime},y_{2}^{\prime}])=$ |  let $t_{1}=x_{1}y_{1}$
---|---
|  let $t_{1}^{\prime}=x_{1}^{\prime}y_{1}+x_{1}y_{1}^{\prime}$
|  let $t_{2}=x_{2}y_{2}$
|  let $t_{2}^{\prime}=x_{2}^{\prime}y_{2}+x_{2}y_{2}^{\prime}$
|  let $t_{3}=t_{1}+t_{2}$
|  let $t_{3}^{\prime}=t_{1}^{\prime}+t_{2}^{\prime}$
| $t_{3}^{\prime}$

To compute the partial derivative of $f$ with respect to each input we need to
set the corresponding tangent variable to $1$ and the other input tangent
variables to $0$. For example, the gradient of $f$ with respect to the first
vector is computed by the following partial derivative computations:

$f^{\prime}([a_{1},a_{2}],[b_{1},b_{2}],[1,0],[0,0])\to^{*}$ | $b_{1}=$ | $\frac{\partial\llbracket f\rrbracket}{\partial x_{1}}(a_{1},a_{2},b_{1},b_{2})$
---|---|---
$f^{\prime}([a_{1},a_{2}],[b_{1},b_{2}],[0,1],[0,0])\to^{*}$ | $b_{2}=$ | $\frac{\partial\llbracket f\rrbracket}{\partial x_{2}}(a_{1},a_{2},b_{1},b_{2})$

Reverse-mode AD. Forward mode AD is computationally expensive for the
derivative computation of scalar-valued functions with tensor inputs, which
among other use cases appear in training machine learning models by optimizing
an objective function. This is due to the fact that when differentiating a
program representing a function $\mathbb{R}^{n}\to\mathbb{R}$, as is often the
case in these contexts, one needs $n$ runs of the program transformed by
forward-mode to obtain the whole gradient. The reverse-mode technique, which
computes the gradient of such functions in one run, is then more appropriate,
and is massively used in deep learning frameworks (Abadi et al., 2015; Paszke
et al., 2019; Bradbury et al., 2018). There has been recent interest on
bridging the gap between theoretical correctness guarantees and more
practical, efficient implementations with complexity guarantees (Krawiec et
al., 2022; Smeding and Vákár, 2023; Huot and Shaikhha, 2022).

Example 1 (cont.). Consider a generalization of the previous function, where
$\cdot$ denotes the dot product of two vectors:

$f(V_{1},V_{2})=V_{1}\cdot V_{2}$

If each input vector has $m$ elements, then the cost of forward-mode AD is
$O(m^{2})$ as it requires $m$ forward passes, each costing $O(m)$. However,
reverse-mode AD can compute the gradient by one forward pass to compute the
primal values and one reverse pass to compute the gradient values, resulting
in an $O(m)$ overall complexity.

Vector Forward-mode AD. Inspired by the use cases that require the computation
of the full Jacobian matrix (Moré, 2006), there have been efforts on batch
computations of forward-mode AD (Khan and Barton, 2015; Shaikhha et al.,
2019). It has been experimentally demonstrated that by leveraging rewrite
rules, one can even recover the asymptotic performance of reverse-mode AD on
vectorized forward-mode AD (Shaikhha et al., 2019).

Figure 2. The bag-of-words representation of documents using a matrix with
documents as rows and words as columns. Note that the matrix has many zero
elements, i.e., it is a sparse matrix.

###  2.2. Sparse tensors and semi-ring dictionaries

Sparse Tensors. Sparse tensors are a type of data structure that is commonly
used to represent high-dimensional data that have a majority of zero values.
Sparse tensors have a compact representation that only stores the non-zero
values and their corresponding indices, which makes them more memory-efficient
than dense tensors for large-scale data. Sparse tensors are used in many
domains, including natural language processing, computer vision, and
scientific computing. For example, in natural language processing, sparse
tensors can be used to represent text data as a bag-of-words or term
frequency-inverse document frequency (TF-IDF) matrix, where the rows
correspond to documents and the columns correspond to words (Figure 2).

Sparse tensors can be manipulated using a variety of specialized algorithms
and data structures, such as compressed sparse row (CSR) and compressed sparse
column (CSC) formats, which enable efficient matrix-vector multiplication and
other operations. However, the irregular sparsity patterns of sparse tensors
pose significant challenges for automatic differentiation.

Example 1 (cont.). In the previous example, if the majority of the elements of
the input vectors of sizxe $m$ are zeros (the number of non-zero elements,
denoted by $nnz$, is such that $nnz<<m$), one can use the CSR representation
shown in Figure 3(a). In this representation, the array pos is a compressed
representation of rows, whereas idx and val show the columns and values of
non-zero elements. For example, pos(0)=0,pos(1)=7 depicts that the indices 0
to 6 of idx/val correspond to the column/value of the elements in the row=0 of
the matrix, and 7 to 14 for the positions of the row=1. However, existing
linear algebra frameworks do not support gradients over this representation.
Thus, rather than computing the gradient in $O(nnz)$, they compute it over the
dense representation in $O(m)$.

(a) Compressed sparse row (CSR) representation.

(b) The logical semi-ring dictionary representation.

Figure 3. Sparse representations of the previous bag-of-words matrix.

Semi-ring Dictionaries. Semi-ring dictionaries are data structures that
subsume sets, multisets, and dense/sparse tensors (Shaikhha et al., 2022). A
semi-ring is a set with two binary operations that satisfy certain axioms,
such as associativity, distributivity, and commutativity. For example, the set
of non-negative integers with addition and multiplication forms a semi-ring,
and the set of Booleans with logical $\vee$ and $\wedge$ forms another semi-
ring. Semi-ring dictionaries are designed to represent sparse tensors as key-
value pairs. In sparse vectors, the keys correspond to the vector indices and
the values correspond to the non-zero elements. In sparse matrices, the keys
correspond to the row indices and the values correspond to the sparse vector
associated with that row. The semi-ring operations are then defined in terms
of the corresponding operations on the values, such as addition or
multiplication. The multiplication operator for semi-ring dictionaries has a
semantics of tensor outer product, as can be observed in the next example.

Example 2 - Scalar-Vector Product. Consider the scalar-vector product between
a scalar value s and a vector value V represented using a semi-ring
dictionary. The equivalent semi-ring dictionary representation is s * V, where
* has the semantics of tensor outer product.

SDQL. SDQL (Shaikhha et al., 2022) is a functional language for querying
against semi-ring dictionaries. SDQL is expressive enough to capture database
queries and linear algebra expressions; this makes it appropriate as an
intermediate language for hybrid database and machine learning workloads. SDQL
provides the following constructs for manipulating semi-ring dictionaries:

  1. (1)

dict(k) accesses the value associated with the key k in dict, and if the key
does not exist, it returns the addition identity of semi-ring ($0$ in the case
of real and natural numbers).

  2. (2)

{k -> v} constructs a dictionary with a single key-value pair of k and v.

  3. (3)

sum(<k,v> in dict) f(k,v) folds over each key-value pair of dict and computes
the summation of f(k,v) starting from the addition identity of the
corresponding semi-ring.

Example 1 (cont.). The equivalent SDQL expression for $V1\cdot V2$ can be one
of the following two:

⬇ sum(<i, a> in V1) a * V2(i) |  ⬇ sum(<i, a> in V2) V1(i) * a
---|---

The preferred choice depends on the number of non-zero elements of V1 and V2.
If V1 (resp. V2) has fewer non-zero elements, the left (resp. right) variant
is more efficient. Otherwise, if both have the same number of non-zero
elements (e.g., both are dense), both variants have the same performance.

SDQLite. SDQLite (Schleich et al., 2022) is a dialect of SDQL tailored for
sparse tensor processing; it restricts SDQL to the types required for sparse
tensors while extending it with constructs required for different sparse
storage formats (e.g., CSR, CSC). The following constructs are central to
SDQLite:

  1. (1)

(st:en) specifies a dense array holding the range of numbers from st to en
(excluding).

  2. (2)

arr(st:en) specifies the sub-array of arr ranging from st to en (excluding).

  3. (3)

Additional annotations for guiding rewrite rules such (Schleich et al., 2022).

##  3\. Languages

In this section, we give an overview of the languages used in $\nabla$SD. We
divide SDQLite in two. The smaller fragment, Logical SDQLite, on which AD will
be performed, and the full fragment, Physical SDQLite, which augments the
logical subset of the language with the constructs for expressing the
different sparse storage formats. The grammar and most important typing rules
of these languages are shown in Figure 4.

Construct | Description | Language
---|---|---
e | ::= | sum(<x,y> in e) e | Dictionary Aggregation |  Logical SDQLite
| $\mid$ |  { e -> e } $\mid$ { } $\mid$ e(e) | Singleton/Empty Dictionary, Lookup
| $\mid$ |  let x = e in e $\mid$ x | Variable Binding & Access
| $\mid$ |  not e $\mid$ if e then e | Negation, Conditional
| $\mid$ |  e + e $\mid$ e * e | Addition, Multiplication
| $\mid$ |  n $\mid$ r $\mid$ false $\mid$ true | Numeric and Boolean Constants
| $\mid$ | op(e) | Real unary operation such as $\cos,\exp,\sin$
| $\mid$ | e = e | Equality of discrete types: bool and int
T | ::= |  I $\mid$ bool $\mid$ D | Scalar and Tensor Types
D | ::= |  real $\mid$ { I -> D } | Dictionary (Tensor) Type
I | ::= | int | Index Types (Only Integer) |
e | ::= | ... | Source SDQLite Expressions |  Physical SDQLite
| $\mid$ |  e:e $\mid$ e(e:e) | Range and Sub-Array
| $\mid$ | cf. (Schleich et al., 2022) | Other Constructs
T | ::= | ... | Source SDQLite Types
D | ::= |  ... $\mid$ int | Tensor Types with Integer Value
I | ::= |  ... $\mid$ dense_int | Index Types (+ Dense Integer) |
Definition of zero[D]: | Definition of $\otimes$:
zero[real] $\triangleq$ 0 |  real $\otimes$ D $\triangleq$ D
zero[{I -> D}] $\triangleq$ { } |  {I -> D1} $\otimes$ D2 $\triangleq$ {I -> D1 $\otimes$ D2}
Definition of let-tupling: | Definition of tensor n:
let <v1, ..., vn> = <e1, ..., en> in e |  tensor 0 $\triangleq$ real
$\triangleq$ let v1=e1 in ... let vn=en in e |  tensor n $\triangleq$ { int -> tensor (n-1) }

|  $\Gamma\vdash$ e1: tensor n+1 $\Gamma$, k: int, v: tensor n $\vdash$ e2: D
---
$\Gamma\vdash$ sum(<k,v> in e1) e2: D
$\Gamma\vdash$ e1: tensor n+1 $\Gamma\vdash$ e2: int
---
$\Gamma\vdash$ e1(e2): tensor n
|  $\Gamma\vdash$ k: int $\Gamma\vdash$ v: tensor n
---
$\Gamma\vdash$ { k -> v }: tensor (n+1)
$\Gamma\vdash$ e1: D1 $\Gamma\vdash$ e2: D2
---
$\Gamma\vdash$ e1 * e2: D1 $\otimes$ D2

Figure 4. Grammar and a subset of typing rules of the languages used in
$\nabla$SD.

Logical SDQLite. Initially, the program is expressed in a subset of SDQLite
that is sufficient for expressing tensor programs at the logical level, i.e.,
without worrying about the storage format. Thus, at this stage, we do not
require the support for dense arrays. Furthermore, there is no need for
expressing tuples. Nevertheless, for convenience, we use tupled let-binding as
a syntactic sugar for multiple let bindings.

Logical SDQLite is expressive enough to capture Einstein summations (Shaikhha
et al., 2022). For example, the matrix-matrix product for two matrices M1 and
M2 (both represented as nested dictionaries) is expressed as:

⬇

sum(<i,row> in M1) { i ->

sum(<j, v1> in row)

sum(<k, v2> in M2(j)) { k ->

v1 * v2 } }

However, the expressiveness goes beyond Einstein summations. As an example,
map of function f over the values of a tensor e is expressed as sum(<k,v> in
e) { k -> f(v) }.

Physical SDQLite. The storage specifications require dense-array-related
constructs such as range (st:end) and subarray (arr(st:end)). Thus, after
combining the differentiated program with the storage specification, we need
to include these constructs for the Physical SDQLite. We go back to this
intermediate language in Section 5. The next section focuses primarily on
Logical SDQLite and the differentiation rules over its constructs.

##  4\. Differentiation

In this section, we present the differentiation transformations applied to
Logical SDQLite expressions. First, for exposition purposes, we present a
variant of traditional forward-mode AD (FAD). Then, we show a tensorized FAD
that not only subsumes the traditional FAD, but also computes gradients more
efficiently. Finally, we show the high-level API exposed to the programmers.

###  4.1. Scalar Forward-Mode Transformation

Traditional FAD uses dual numbers to compute the tangent (derivative)
component along with the actual (original) computation. We refer to it as
scalar FAD because for each scalar expression, it stores a scalar tangent
component.

Similar to other FAD frameworks, $\nabla$SD precedes the differentiation
transformation with an ANF conversion (Flanagan et al., 1993). This allows for
sharing sub-expressions and avoids duplication of computations for non-unary
constructs such as multiplication. Logical SDQLite does not allow function
definitions nor higher-order functions; all functions need to be inlined
(Paszke et al., 2021).

Scalar Constructs. Figure 5 shows the forward-mode transformation rules. As
opposed to existing functional AD systems, $\nabla$SD does not use explicit
pair construction and projection for dealing with dual numbers. Instead, the
$\mathcal{F}\llbracket\rrbracket$ construct only computes the tangent part of
differentiation and refers to the expressions in the ANF transformed program
for primal components (cf. the rule for let binding). This avoids the need to
extend the target language of differentiation with pairing constructs.
Furthermore, this makes the differentiation rules simpler. For every unary
real operation op, we assume that the language has a unary real operation op’
representing its derivative. Finally, the differentiation for all discrete
types (int and bool) is 0.

|  |  $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$ Forward mode on Types
---
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} D }\rrbracket$ | = | D
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{} bool }\rrbracket$ | = | real
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{} int }\rrbracket$ | = | real
|  $\mathcal{F}\llbracket\Gamma\rrbracket$ Forward mode on Context
---
$\mathcal{F}\llbracket\emptyset\rrbracket$ | = | $\emptyset$
$\mathcal{F}\llbracket\Gamma,\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}x:T}\rrbracket$ | = |  $\mathcal{F}\llbracket\Gamma\rrbracket$$,\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}x:T},\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}x':}$$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$
Forward mode on Expressions
Invariant: If $\Gamma\vdash$ e: T, then
$\mathcal{F}\llbracket\Gamma\rrbracket$ $\vdash$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$:
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}sum}(<k,v> {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}in} e1) e2 }\rrbracket$ | = | sum(<k,v> in e1)
|  | let <k’, v’> = <0, $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1(k)}\rrbracket$> in
|  | $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}let} x = e1 {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}in} e2 }\rrbracket$ | = | let <x, x’> = <e1, $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$> in
|  | $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}if} e1 {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}then} e2 }\rrbracket$ | = | if e1 then $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} e1(e2) }\rrbracket$ | = |  $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$(e2)
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} \\{ e1 -> e2 \\} }\rrbracket$ | = | { e1 -> $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$ }
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} e1 * e2 }\rrbracket$ | = | e1 * $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$ \+ $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$ * e2
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} e1 + e2 }\rrbracket$ | = |  $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$ \+ $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} op(e) }\rrbracket$ | = | op’(e) * $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} x }\rrbracket$ | = | x’
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} r }\rrbracket$ | = | 0
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} n }\rrbracket$ = $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}false}\rrbracket$ = $\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}true}\rrbracket$ | = | 0
Figure 5. Forward-mode AD transformation rules for SDQLite expressions. In
order to maximize sharing, all sub-expressions need to be let-bound, i.e., ANF
transformation should be applied before differentiation.

Example 1 (cont.). Consider again the case of the dot-product of two unrolled
vectors of size two initially used in Section 2. Applying differentiation over
the ANF transformed program in SDQLite is as follows:

⬇ let t1 = x1*y1 in let t2 = x2*y2 in let t3 = t1+t2 in t3 | $\Biggl{\rrbracket}$
---|---

After applying the rules in Figure 5, we obtain the following program:

⬇

let <t1, t1’> = <x1*y1, x1*y1’ + x1’*y1> in

let <t2, t2’> = <x2*y2, x2*y2’ + x2’*y2> in

let <t3, t3’> = <t1+t2, t1’+t2’> in

t3’

Note that the let-binding constructs are syntactic sugar; there is no pair
created (cf. Figure 4).

Tensor Constructs. By choosing not to incorporate pairs in the language, we
have eliminated the option of differentiating vectors as vectors of pairs
(i.e., arrays of structs). Instead, an expression of type tensor n is
differentiated as an expression with the same type. One of the interesting
tensor-based differentiation rules is our rule for summation, where we need to
access the corresponding element from the differentiated range, as we can
observe in the following example.

Example 1 (cont.). Let us go back to the dot-product for two vectors V1 and
V2. The differentiation transformation is expressed as follows:

⬇ sum(<i,a> in V2) V1(i) * a | $\rrbracket$
---|---

Applying differentiation rules results in the following program:

⬇

sum(<i,a> in V2)

let <i’,a’> = <0,V2’(i)>

V1(i) * a’ + V1’(i) * a

In order to compute the gradient of this function with respect to one of its
vector inputs, say V1, we need to repeatedly set V1’ into a one-hot vector
that is 1 at index j and 0 everywhere else, and set V2’ to be the zero vector.
This requires multiple rounds of running the forward-mode AD for different
one-hot vectors, which is computationally expensive. Previous research
(Shaikhha et al., 2019) has shown how this can be optimized by wrapping the
vector construction around the forward-mode AD and applying loop
optimizations.

Example 2 (cont.). Consider the case of scalar-vector product, represented as
s * V in SDQLite. Applying the scalar FAD rules on this program results in:

⬇

s * V’ + s’ * V

In the case of differentiation with respect to V, similar to the dot-product
example, one needs to repeatedly pass all different one-hot vectors as V’.
However, the differentiation with respect to s can be done by setting s’ to 1
and V’ to zero[tensor 1].

Next, we show an alternative differentiation transformation that enables
native tensorized forward-mode AD.

###  4.2. Tensorized Forward-Mode Transformation

As tensors are first-class citizens in SDQLite, one can directly express the
differentiation with respect to tensor variables of type $\tau$, represented
by $\mathcal{D}_{\tau}\llbracket\rrbracket$. This means that the derivative of
an expression of type tensor n with respect to a tensor of type tensor m, will
be a tensor (n + m), which is the same as the outer tensor product type (cf.
Figure 6, the rules for types).

Figure 6 shows the rules for tensorized FAD. They generalize the rules for
scalar forward-mode AD, which one recovers by setting $\tau$ to be real. The
key differences are in the rules for constant reals and multiplication. Rather
than just returning a real-valued 0, tensorized FAD returns the zero value of
type $\tau$ represented as zero[$\tau$]. For multiplication, if $\tau$ is a
tensor type with a non-zero order, the first term still computes the
multiplication of e1 and the differentiation of e2. However, the second term
requires re-arranging the indices of the tensors. This complication can be
avoided by only allowing for the multiplication of real numbers in the input
program. This is achieved by applying multiplication normalization (cf.
Section 5.4).

|  |  $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$ Tensorized FAD on Types
---
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} D }\rrbracket$ | = | $\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}D}~{}~{}~{}\textrm{\color[rgb]{0,0.4,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0.4,0}{}$\otimes$}~{}~{}~{}\tau$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{} bool }\rrbracket$ | = | real
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{} int }\rrbracket$ | = | real
|  $\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$ Tensorized FAD on Context
---
$\mathcal{D}_{\tau}\llbracket\emptyset\rrbracket$ | = | $\emptyset$
$\mathcal{D}_{\tau}\llbracket\Gamma,\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}x:T}\rrbracket$ | = |  $\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$$,\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}x:T},\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}x':}$$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$
Tensorized FAD on Expressions
Invariant: If $\Gamma\vdash$ e: T, then
$\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$ $\vdash$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$:
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}sum}(<k,v> {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}in} e1) e2 }\rrbracket$ | = | sum(<k,v> in e1)
|  | let <k’, v’> = <0, $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1(k)}\rrbracket$> in
|  | $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}let} x = e1 {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}in} e2 }\rrbracket$ | = | let <x, x’> = <e1, $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$> in
|  | $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}if} e1 {\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}then} e2 }\rrbracket$ | = | if e1 then $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} e1(e2) }\rrbracket$ | = |  $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$(e2)
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} \\{ e1 -> e2 \\} }\rrbracket$ | = | { e1 -> $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$ }
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} e1 * e2 }\rrbracket$ | = | e1 * $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$ \+ $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$ *T[$\tau$] e2
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} e1 + e2 }\rrbracket$ | = |  $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$ \+ $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} op(e) }\rrbracket$ | = | op’(e) * $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} x }\rrbracket$ | = | x’
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} r }\rrbracket$ | = | zero[$\tau$]
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{} n }\rrbracket$ = $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}false}\rrbracket$ = $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}true}\rrbracket$ | = | 0
e1*T[real] e2 | $\triangleq$ | e1 * e2
e1*T[tensor n] e2 | $\triangleq$ | sum(<i_1,r_2> in e1) ...
e1: tensor (m + n) |  | sum(<i_m,v> in r_m))
|  | {i_1 -> ... { i_m -> 1 } ... } * e2 * v
Figure 6. Tensorized Forward-mode Automatic Differentiation (FAD) rules for
SDQLite expressions. The type $\tau$ needs to be a tensor type, i.e., it
follows the grammar of D. The rules of Figure 5 are the special case of
$\mathcal{F}\llbracket\rrbracket$=$\mathcal{D}_{\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}real}}\llbracket\rrbracket$.

Example 1 (cont.). In our running example, tensorized FAD transformation is
represented as $\mathcal{D}_{\tau}\llbracket$sum(<i,a> in V2) V1(i) *
a$\rrbracket$.

By applying the rules in Figure 6 we have:

⬇

sum(<i,a> in V2)

let <i’,a’> = <0,V2’(i)>

V1(i) * a’ + V1’(i) * a

As one can see, this program looks identical to the version generated by
scalar FAD. The difference is in the type of V1’ and V2’. In scalar FAD, their
type is the same as V1 and V2, i.e., tensor 1. In tensorized FAD, their type
is tensor 2. Thus, the multiplications in the last expression are now scalar-
vector multiplications.

As opposed to scalar FAD, we need to assemble the zero and one-hot vectors of
all iterations together; instead of passing them vector by vector, we pass
them as an entire matrix in which each row represents one of the one-hot
vectors. The definition of onehot[tensor 1] in Figure 7 specifies how one can
build such a matrix for variable x.

Example 3. Consider the case of computing the trace of the matrix M. The
tensorized FAD over this expression in SDQLite is represented as
$\mathcal{D}_{\tau}\llbracket$ sum(<i,r> in M) r(i)$\rrbracket$.

For each row r of matrix M at index i, we compute the summation of diagonal
elements specified by r(i). After our tensorized FAD transformation, we obtain
the following program:

⬇

sum(<i,r> in M)

let <i’,r’> = <0,M’(i)>

r’(i)

To compute the one-hot input for a matrix, we need to generalize the case of a
vector. In the case of differentiating with respect to a vector (tensor 1), we
were passing a matrix (tensor 2) for the one-hot inputs. Here, in the case of
differentiating w.r.t. a matrix (tensor 2), we need to pass an order-4 tensor
(tensor 4). The definition of onehot[tensor 2] can also be found in Figure 7.

gradient e (x: $\tau$) | $\triangleq$ | let <v1’,..., vn’> = <ingrad v1 x,...,ingrad vn x> in
---|---|---
|  | $\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$
where |  |  $\\{$v1, ..., vn$\\}=FVS($e$)$
ingrad (v: D) x | $\triangleq$ | zero[D]
ingrad (x: D) x | $\triangleq$ | onehot[D] x
onehot[real] x | $\triangleq$ | 1
onehot[tensor 1] x | $\triangleq$ | sum(<i,_> in x) { i -> {i -> 1} }
onehot[tensor 2] x | $\triangleq$ | sum(<i,v> in x) { i ->
|  | sum(<j,_> in v) {j ->
|  | {i -> {j -> 1} } } }
onehot[tensor n] x | $\triangleq$ | sum(<i_1,v_2> in x) { i_1 ->
if n > 2 |  | sum(<i_2,v_3> in v_2) { i_2 ->
|  | ... sum(<i_n,_> in v_n)) { i_n ->
|  | {i_1 -> ... { i_n -> 1 } ... }
Figure 7. The gradient API exposed by $\nabla$SD to the programmer.

###  4.3. Putting it All Together

Programmers who use machine learning frameworks need not concern themselves
with the source-to-source transformations employed in the background. To
accomplish this, $\nabla$SD provides a high-level API through the gradient
macro, which accepts two inputs. The first input is the SDQLite expression to
be differentiated, while the second input is the free variable with respect to
which we perform the differentiation.

Consider the case of computing the gradient of the expression e with respect
to x. This is represented as gradient e x. The gradient macro generates let-
bindings for the differentiation components of all free variables of the
expression e
($FVS(\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e})$).
The RHS of let-binding for a free variable v is ingrad v x. The ingrad macro
is responsible for computing the input zeros or one-hots, depending on whether
variable v is different than x or is the same. In the former case, the zero
macro is used by passing the type of v. In the latter case, the onehot macro
is used. For an input variable of type tensor n, the onehot assembles a sparse
tensor (2*n) where the n indices of the input variable x are repeated twice so
that the diagonals of the hypercube are set to 1.

Example 1 (cont.). Our running example is specified by the high-level API as
follows:

⬇

gradient (sum(<i, a> in V1) a * V2(i)) V2

This is expanded to the following expression:

⬇

let V1’ = zero[{int -> real}]

let V2’ = onehot[{int -> real}] V2

sum(<i, a> in V1)

let <i’, a’> = <{}, V1’(i)> in

a * V2’(i) + a’ * V2(i)

Which is, in turn, expanded to:

⬇

let V1’ = {}

let V2’ = sum(<i,_> in V2) {i -> {i -> 1}}

sum(<i, a> in V1)

let <i’, a’> = <{}, V1’(i)> in

a * V2’(i) + a’ * V2(i)

In the next section, we see how this expression will be optimized. Note that
by simple constant propagation, the asymptotic complexity will already be
$O(n)$ on this example. In fact, we will only perform general optimizations
which are not AD-specific.

##  5\. Performance

In this section, we review the techniques we employ to improve performance.
While Figure 1 showed the simplified overview of our approach, Figure 8
presents in more detailed the series of transformations applied in $\nabla$SD.

Figure 8. The transformations in $\nabla$SD.

###  5.1. Pre-differentiation

Equality Saturation. First, $\nabla$SD applies algebraic rewrite rules over
the input program to leverage optimizations such as factorization, loop
fusion, etc (Schleich et al., 2022). We use equality saturation in order not
to worry about phase ordering problems and making sure that the rewrite rules
are applied globally (Tate et al., 2009). We use EGG (Willsey et al., 2021), a
state-of-the-art implementation of equality saturation that has been
successfully used for SDQLite in the context of flexible storage specification
for tensor programs (Schleich et al., 2022). We rely on the cost models and
the algebraic rewrite rules specified in (Schleich et al., 2022); there is no
need to specify any AD-specific cost model or rewrite rule.

ANF. Then, $\nabla$SD applies an ANF transformation (Flanagan et al., 1993),
which ensures that sub-expressions are simple expressions, i.e. constant
values or variable references. This is achieved by using a let-binding for the
sub-expressions if they are not already simple expressions. The ANF
transformed program is then fed into the differentiation transformation
presented in Section 4.

Rewrite Rule | Condition
---|---
zero[D1] * e2 | $\leadsto$ |  zero[D1 $\otimes$ D2] | if e2: D2
e1 * zero[D2] | $\leadsto$ |  zero[D1 $\otimes$ D2] | if e1: D1
e1 + zero[D1] | $\leadsto$ | e1 | must have e1: D1
zero[D1] + e2 | $\leadsto$ | e2 | must have e2: D1
zero[tensor n](e2) | $\leadsto$ | zero[tensor (n-1)] | must have n > 0
let x=zero[D] in e2 | $\leadsto$ | e2[x$\rightarrow$zero[D]] |
sum(<k,v> in zero[D])e2 | $\leadsto$ | zero[D2] | if e2:D2
e1 * e2 | $\leadsto$ | sum(<i_1,v_2> in e1) ... | if e1: tensor n
|  | sum(<i_n,v> in v_n)) | and n > 0
|  | {i_1->...{i_n-> v*e1}...} |
e1 * e2 | $\leadsto$ | sum(<i_1,v_2> in e2) ... | if e1: real
|  | sum(<i_n,v> in v_n)) | and e2: tensor n
|  | {i_1->...{i_n-> e1*v}...} | and n > 0
Figure 9. Transformations used in $\nabla$SD. The first group of rewrite rules
are related to sparsity propagation. The second group correspond to
multiplication normalization.

###  5.2. Post-Differentiation

Sparsity Propagation. After differentiation, many intermediate highly sparse
values (e.g., zero tensors) are constructed. Even though expressing local
rewrite rules for simplifying them is possible, these programs are very large
and optimizing them with equality saturation requires a large search space.
Thus, we applied these sparsity propagation rules as a separate pass. The
first group of rewrite rules in Figure 9 correspond to the sparsity
propagation rules. Finally, we again pass the program to equality saturation.

Example 1 (cont.). In the dot-product example, by assuming the second variant
(iteration over the second vector), the differentiated program with respect to
the first vector is as follows:

⬇

let V1’ = sum(<i, a> in V1) {i -> {i -> 1}}

let V2’ = {}

sum(<i, a> in V2)

let <i’, a’> = <{}, V2’(i)> in

V1(i) * a’ + V1’(i) * a

The sparsity propagation optimization propagates V2’ and simplifies the
relevant expressions (e.g., V1(i) * a’) as follows:

⬇

let V1’ = sum(<i, a> in V1) {i -> {i -> 1}}

sum(<i, a> in V2)

V1’(i) * a

After applying equality saturation, if V1 is a dense vector, the program is
optimized and reduces to simply V2, a slightly less optimized version of which
is:

⬇

sum(<i, a> in V2)

{ i -> a }

| Storage Format
---
(Inputs)
Physical SDQLite Definition
⬇ sum(<_, i> in (0:len)) { unique(row(i)) -> val(i) }
⬇ sum(<_, i> in (0:len)) { unique(i) -> V(i) }
⬇ sum(<_, i> in (0:len)) { unique(i) -> sum(<p, j> in idx(pos(i):pos(i+1))) {
unique(j) -> val(p) } }
⬇ sum(<_, i> in (0:rows)) { unique(i) -> sum(<_, j> in (0:cols)) { unique(j)
-> M(i*cols+j) } }
Figure 10. The functional implementation of various sparse/dense formats in
Physical SDQLite. Matrix CSC and Dense Column Major are expressed similarly to
Matrix CSR and Dense Row Major, respectively.

###  5.3. Storage Composition

After optimizing the differentiated expression, $\nabla$SD uses the storage
specifications as the definition for each input tensor. Figure 10 shows the
specification of several storage formats in SDQLite. $\nabla$SD generates let-
bindings that bind every input tensor to the expression specifying its
physical storage format. However, this makes the performance even worse, due
to unnecessary intermediate tensor allocations. Luckily, previous research
(Schleich et al., 2022) showed how equality saturation can not only remove
these intermediate tensors but also recover well-known algorithms for sparse
tensor computations.

Example 1 (cont.). In the previous example, if we assume that V1 is a dense
vector and V2 is a sparse vector using the COO representation, combining the
storage specification with the differentiated program results in:

⬇

let V1 = sum(<_, i> in (0:V1_len)) { unique(i) -> V1_V(i) } in

let V2 = sum(<_, i> in (0:V2_len)) { unique(V2_row(i)) -> V2_val(i) } in

sum(<i, a> in V2)

{ i -> a }

By applying equality saturation, $\nabla$SD returns the following program that
does not introduce unnecessary intermediate tensors, as expected from
(Schleich et al., 2022):

⬇

sum(<_, i> in (0:V2_len))

{ V2_row(i) -> V2_val(i) }

###  5.4. Normalization

$\nabla$SD performs additional lower-level transformations on the optimized
storage-format-aware program. First, multiplication normalization rewrites
tensor outer products into summation expressions with scalar multiplications.
The rewrite rules for multiplication normalization are shown in the second
group of rules of Figure 9. Once combined with loop fusion rules,
multiplication normalization removes unnecessary intermediate tensors. The
second normalization involves ANF transformations which have been previously
cancelled by equality saturation and other transformations.

Example 5 - BATAX. The BATAX kernel (Nelson et al., 2015) is represented as
$\beta A^{T}Ax$ where $\beta$ is a scalar value, $A$ a matrix, and $x$ a
vector. This kernel is expressed in SDQLite as follows:

⬇

sum(<i, r> in A)

sum(<j, v1> in r)

sum(<k, v2> in r)

{ j -> ((beta * v1) * v2) * (x(k)) }

After differentiating with respect to x, post-differentiation optimizations,
$\nabla$SD produces:

⬇

beta * (sum(<i, r> in A) r * r)

Considering a CSR representation for A and a dense representation for x, after
applying storage composition and equality saturation we derive:

⬇

beta * sum(<_, i> in (0:A_len))

let r = sum(<p, j> in A_VCol(A_VRow(i):A_VRow((i + 1))))

{ j -> A_VVal(p) } in

r * r

The multiplication of beta with the result of summation is a scalar-vector
product. Also, the last expression r * r corresponds to a vector outer
product. The multiplication normalization rewrites both expressions, the
result of which is as follows:

⬇

sum(<_, i> in (0:A_len))

let r = sum(<p, j> in A_VCol(A_VRow(i):A_VRow((i + 1))))

{ j -> A_VVal(p) } in

sum(<i1, v1> in r) { i1 ->

sum(<i2, v2> in r)

{ i2 -> beta * v1 * v2 } }

###  5.5. Destination-Passing Style & C++ Code Generation

As final step, $\nabla$SD generates C++ code. The index and scalar types int
and real translate into size_t and double, respectively. The expressions of
type { int -> real } are translated to objects of type dict_type<size_t,
double>. Nested dictionaries are recursively translated, e.g., { int -> { int
-> real } } $\leadsto$ dict_type<size_t, dict_type<size_t, double>>. We use
the robinhood dictionary (rbd, 2023) for the C++ runtime. Dictionaries with a
key type dense_int are translated to arr_type instead of dict_type. For
example, { dense_int -> real } translates into arr_type<double> and {
dense_int -> int } into arr_type<size_t>.

The C++ code generation for most constructs of SDQLite is straightforward. The
construct dict(key) is translated into dict[key], which corresponds to a
dictionary or array lookup in C++. The addition and multiplication constructs
are converted to the same primitives in C++. This requires the support for +
and * over dictionary types, which is provided by our C++ runtime. The sum
construct is translated into a for-loop. Similarly, nested summations are
translated into nested for-loops. If the range expression of a summation is
the range construction (st:en) or sub-array arr(st:en), the translation
produces a standard for-loop. However, if it is a dictionary, the translation
generates a for-each construct over the dictionary. Another interesting case
is the code generation for the singleton dictionary construct. This construct
is mostly used inside a summation. In this case, it is translated to a
dictionary update as follows:

⬇ let res = sum(<k,v> in dict) { f(k,v) -> g(k,v) } | $\leadsto$ |  ⬇ dict_type<size_t, double> res; for(auto& kv : dict) { res[f(kv.first, kv.second)] += g(kv.first, kv.second); }
---|---|---

Note that the above translation works if g(k,v) outputs a dictionary. This is
because += is also overloaded by dictionaries.

Example 1 (cont.). The generated C++ code for the differentiated program is as
follows:

⬇

void VVD_wrt_V1(

// inputs for the dense vector V1

arr_type<double>& V1_V, size_t V1_len,

// inputs for the sparse vector V2

arr_type<size_t>& V2_VRow, arr_type<double>& V2_VVal, size_t V2_len,

// destination dictionary

dict_type<size_t, double>& result) {

for(size_t i = 0; i < V2_len; i++) {

result[V2_VRow[i]] += V2_VVal[i];

} }

In order to generate efficient C++, $\nabla$SD leverages the destination-
passing style (DPS) technique (Minamide, 1998) in two ways. First, the
generated function is provided with the destination object to store the final
results (Shaikhha et al., 2017). Second, The destination-passing style
transformation removes intermediate tensors created in inner loops. This is
achieved by pushing all the singleton dictionary constructions into the inner
loops to avoid unnecessary intermediate dictionary constructions (Shaikhha,
2022).

Example 4 (cont.). In the BATAX kernel, the C++ code generation produces the
following program:

⬇

void BATAX_stg_wrt_X(

// input scalar

double& beta_S,

// input for the CSR matrix A

arr_type<size_t>& A_VRow, arr_type<size_t>& A_VCol, arr_type<double>& A_VVal,
size_t A_len,

// input for the dense vector X

arr_type<double>& X_V, size_t X_len,

// destination dictionary

dict_type<size_t, dict_type<size_t, double>>& result) {

for(size_t i = 0; i < A_len; i++) {

dict_type<size_t, double> r;

for(size_t p = A_VRow[i]; p < A_VRow[(i + 1)]; p++) {

size_t j = A_VCol[p];

r[j] += A_VVal[p];

}

⬇

dict_type<size_t, dict_type<size_t, double>> tmp1;

for(auto& iv1 : r) {

dict_type<size_t, double> tmp2;

for(auto& iv2 : r) {

tmp2[iv2.first] += beta_S * (iv1.second * iv2.second);

}

tmp1[iv1.first] += tmp2;

}

result += tmp1;

⬇

} }

Note that the last inner loop (highlighted) requires constructing two
intermediate dictionaries tmp1 and tmp2. One can use the associativity of
semi-ring dictionaries in order to reorder the insertions so that all the
insertions happen in the most-inner loop directly into result:

⬇

for(auto& iv1 : r) {

for(auto& iv2 : r) {

result[iv1.first][iv2.first] += beta_S * (iv1.second * iv2.second);

} }

This removes the need to allocate the intermediate dictionaries tmp1 and tmp2.
We will see the impact of this optimization in Section 7.

##  6\. Proofs

In this section, first, we prove the type soundness of the automatic
differentiation. Then, we show that programs in Logical SDQLite denote
differentiable functions for a generalization of differentiability that
conservatively extends the one between Euclidean spaces. Finally, we show the
correctness of both scalar and tensorized forward-mode AD.

###  6.1. Type Soundness of Differentiation

###### Lemma 6.1 (Well-typedness of $\mathcal{D}_{\tau}\llbracket\rrbracket$).

If $\Gamma\vdash$ e:T then
$\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$$\vdash$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$:
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}T}\rrbracket$.

###### Proof.

By induction on the typing tree of e and case analysis. The interesting cases
are:

Case *: Let us consider e1 * e2: tensor m and $\tau=$ tensor n. We need to
prove that
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1
* e2}\rrbracket$: tensor (m+n). Based on the induction hypothesis and typing
rules, we have the following:

  * •

e1: tensor m1 $\Rightarrow$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$:
tensor (m1+n)

  * •

e2: tensor m2 $\Rightarrow$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$:
tensor (m2+n)

  * •

e1: tensor m1, e2: tensor m2 $\Rightarrow$ e1 * e2: tensor (m1+m2)
$\Rightarrow$ m1+m2 = m

By considering the sub-expressions of
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1
* e2}\rrbracket$ we have:

  * •

e1
*$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$:
tensor (m1+(m2+n)) $\Rightarrow$ e1
*$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$:
tensor ((m1+m2)+n)

  * •

$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$
*T[$\tau$] e2: tensor ((m1+m2)+n)

Thus, the addition of these two sub-expressions also has the same type, i.e.,
tensor ((m1+m2)+n). As m1+m2 = m, the output type is tensor (m+n).

Case sum: Let us consider $\Gamma\vdash$ sum(<k, v> in e1) e2: tensor m and
$\tau=$ tensor n. We need to prove that
$\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$ $\vdash$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}sum}\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}(<k,
v> in e1) e2}\rrbracket$: tensor (m+n). Based on the induction hypothesis and
typing rules, we have the following:

  * •

$\Gamma\vdash$ e1: tensor (m1+1) $\Rightarrow$
$\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$ $\vdash$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$:
tensor (m1+1+n)

  * •

$\Gamma^{\prime}=\Gamma,$k:int,v:tensor m1

  * •

$\Gamma^{\prime\prime}=$
$\mathcal{D}_{\tau}\llbracket\Gamma\rrbracket$,k:int,v:tensor
m1,k’:real,v’:tensor (n+m1)

  * •

$\Gamma^{\prime}\vdash$ e2: tensor m $\Rightarrow$
$\Gamma^{\prime\prime}\vdash$
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e2}\rrbracket$:
tensor (m+n)

Let us name the body of sum in its differentiation as e’:

  * •

$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}sum}\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}(<k,
v> in e1) e2}\rrbracket$ =
sum(<k,v> in e1) let <k’,v’> =
<0,$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$(k)>
in$\mathcal{D}_{\tau}\llbracket e2\rrbracket$ =
sum(<k,v> in e1) e’ $\Rightarrow$
e’ = let <k’,v’> =
<0,$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e1}\rrbracket$(k)>
in$\mathcal{D}_{\tau}\llbracket e2\rrbracket$

By the induction hypothesis, we know that $\Gamma^{\prime}\vdash$ e’: tensor
(m+n).

As a result, we derive $\Gamma\vdash$ sum(<k,v> in e1) e’: tensor (m+n), which
we needed to prove. ∎

The rest of this section goes into more mathematical depth and can be skipped
by non-expert readers. Its aim is to formalize the argument that tensorized
FAD computes gradients. We first interpret programs of Logical SDQLite as
functions and define what it means to compute derivatives and gradients for
such functions. We then extend the definition of differentiability and
differentials beyond the usual one for functions between Cartesian spaces. We
then use a logical relations argument to show that
$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$
denotes the differential of e and therefore produces correct gradients.

###  6.2. Semantics and Differentiability

Non-smooth semantics. We can give a denotational semantics to Logical SDQLite
as a special case of the one given in (Shaikhha et al., 2022). The type real
is interpreted as reals $\mathbb{R}$, and programs x: real $\vdash$ e: real as
functions $\mathbb{R}\to\mathbb{R}$. We use semantic brackets
$\llbracket-\rrbracket$ for the denotational semantics of types and terms. We
interpret bool as the booleans $\mathbb{B}$, and int as the natural numbers
$\mathbb{N}$. The semantics of tensor n is defined inductively, with
$\llbracket${ int -> D }$\rrbracket$ given by the tensor product of the free
$\mathbb{R}$ vector space on the countable set
$\mathbb{N}=\llbracket$int$\rrbracket$ with the real vector space
$\llbracket$D$\rrbracket$. This space is only defined up to isomorphism of
real vector spaces, and we are free to choose $\llbracket$tensor n
$\rrbracket$ to be the free real vector space on I$\times...\times$I, the $n$
fold Cartesian product of I with itself. By interpreting contexts
$\llbracket\Gamma\rrbracket=\\{$x1: T1, $\ldots$, xn:Tn$\\}$ as
$\llbracket\Gamma\rrbracket=\llbracket$T1$\rrbracket\times\ldots\times\llbracket$Tn$\rrbracket$,
programs $\Gamma\vdash$ e: T denote functions
$\llbracket$e$\rrbracket:\llbracket\Gamma\rrbracket\to\llbracket$T$\rrbracket$.

The correctness of the rewrite rules from Figure 9 is immediate, as they can
all be seen as special cases of optimizations given in related work (Shaikhha
et al., 2022; Schleich et al., 2022), and they rely on the semi-ring structure
of our dictionaries.

Smooth semantics. The limitation of the non-smooth semantics is that it does
not inform us of the differentiability of our programs. A function
$f:\mathbb{R}^{n}\to\mathbb{R}^{m}$ is smooth if it is differentiable and its
derivative is smooth. We can refine our semantics of Logical SDQLite by
interpreting our types as generalizations of Cartesian spaces, called
diffeological spaces (Iglesias-Zemmour, 2013), and our functions as morphisms
of diffeological spaces, which can be thought of as generalized smooth
functions between these spaces, closely following (Huot et al., 2020). More
precisely, a diffeological space is a pair $(X,\mathcal{P}^{X})$ of a set $X$
and sets of functions $\mathcal{P}_{U}^{X}\subseteq X^{U}$ for all subsets $U$
of Euclidean spaces $\mathbb{R}^{n}$, called plots. Each such set of plots
$\mathcal{P}_{U}^{X}$ must contain all constant functions, be closed under
pre-composition by smooth functions, and be closed under countable glueing.
The latter means that if $(f_{i}:U_{i}\subseteq\mathbb{R}^{n}\to
X)_{i\in\mathbb{N}}$ are a countable family of plots such that for each
$i,j\in\mathbb{N}$, $f_{i}$ and $f_{j}$ agree on $U_{i}\cap U_{j}$, then
$f:\bigcup_{i}U_{i}\to X,x\in U_{i}\mapsto f_{i}(x)$ is a plot. Morphisms
$f:(X,\mathcal{P}^{X})\to(Y,\mathcal{P}^{Y})$ between diffeological spaces are
functions $f:X\to Y$ such that for all $p\in\mathcal{P}_{U}^{X}$, $f\circ
p\in\mathcal{P}_{U}^{Y}$. $\mathbb{N}$ or $\mathbb{B}$ with constant functions
as plots form diffeological spaces. $\mathbb{R}$ with smooth functions
$U\to\mathbb{R}$ as plots $\mathcal{P}_{U}^{\mathbb{R}}$, is a diffeological
space. Morphisms of diffeological spaces
$(\mathbb{R},\mathcal{P}^{\mathbb{R}})\to(\mathbb{R},\mathcal{P}^{\mathbb{R}})$
are exactly smooth functions $\mathbb{R}\to\mathbb{R}$. Diffeological spaces
are closed under products, where a plot $f:U\to X\times Y$ is given by a pair
of plots $(f_{1},f_{2})$, with
$f_{1}\in\mathcal{P}_{U}^{X},f_{2}\in\mathcal{P}_{U}^{Y}$. We can now revise
our denotational semantics by interpreting types as diffeological spaces and
terms as morphisms. For the interpretation of types, the only remaining case
is that of dictionary types. The underlying set is the same as the set
interpretation, and the plots are given by

$\mathcal{P}_{U}^{\llbracket\\{I\to D\\}\rrbracket}=\\{f:U\to\llbracket\\{I\to
D\\}\rrbracket~{}\mid~{}\exists g\in\mathcal{P}_{U}^{\llbracket
D\rrbracket^{n}},i:\llbracket D\rrbracket^{n}\to\llbracket\\{I\to
D\\}\rrbracket\text{ linear injective},f=i\circ g\\}$

One can easily show that all the semantics of all the constructs of the
language preserve plots. As a corollary, we obtain the following result:

###### Proposition 6.2.

Well-typed programs x1: real, $\ldots$, xn: real $\vdash$ e: real denote
smooth functions $\mathbb{R}^{n}\to\mathbb{R}$.

Differentials. For functions $f:\mathbb{R}\to\mathbb{R}$, the differential of
a differentiable function $f$ is a function
$df:\mathbb{R}\times\mathbb{R}\to\mathbb{R}$ such that for any
$x\in\mathbb{R}$, $df(x,-):\mathbb{R}\to\mathbb{R}$ is the best linear
approximation of $f$ at $x$, which is characterized by its slope
$f^{\prime}(x)$. Thus, $df(x,v)=f^{\prime}(x)\cdot v$. More generally, the
differential of a differentiable function $f:\mathbb{R}^{n}\to\mathbb{R}^{m}$
is a function $df:\mathbb{R}^{n}\times\mathbb{R}^{n}\to\mathbb{R}^{m}$ such
that $df(x,v)=J_{x}f\cdot v$, where $\cdot$ is now a matrix-vector product.
$J_{x}f$ is the Jacobian of $f$ at each, i.e. the $n\times m$ matrix of
partial derivatives of $f$ at $x$: $J_{x}f=\Big{(}\frac{\partial}{\partial
x_{j}}f_{i}(x)\Big{)}_{1\leq j\leq n,1\leq i\leq m}$, where $f_{i}$ is the
projection of $f$ on its $i$-th component.

###  6.3. Correctness of Scalar Forward-Mode

Correctness on a first-order sublanguage. Consider the simple first-order
sublanguage of Logical SDQLite that does not involve dictionaries. On this
sublanguage, semantically, the forward-mode algorithm sends $f$ to its
differential $df$, up to a reordering of the inputs. More precisely, if x1:
real, $\ldots$, xn: real $\vdash$ e: real, and x1: real, x1’: real, $\ldots$,
xn: real, xn’: real $\vdash$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$:
real, then $\llbracket$
$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$$\rrbracket(x_{1},x_{1}^{\prime},\ldots,x_{n},x_{n}^{\prime})=J_{(x_{1},\ldots,x_{n})}\llbracket$e$\rrbracket\cdot(x_{1}^{\prime},\ldots,x_{n}^{\prime})^{T}$,
where $(-)^{T}$ is transposition. This can be proved by routine induction on
the structure of terms.

Beyond differentials. Because of tensors, the semantics of our programs do not
represent functions $\mathbb{R}^{n}\to\mathbb{R}^{m}$, and thus lack a
definition of differentials. Even if we were only interested in programs of
the form x1: real, $\ldots$, xn: real $\vdash$ e: real, e may use tensors
internally, and a correctness proof by simple induction would fail. Thus, we
need to define differentials of functions whose input or outputs can be
tensors. We call $\pi:\llbracket$tensor m$\rrbracket\to\mathbb{R}^{k}$ a
canonical projection if it is a projection onto the free-vector space
$\mathbb{R}^{k}$ generated by a subset of size $k$ of I. The adjoint linear
maps $i:\mathbb{R}^{k}\to\llbracket$tensor m$\rrbracket$ are called canonical
injections. A canonical injection embodies the intuition of tensors with a
fixed shape, i.e., a guarantee on their size if they were to be stored as
dense tensors. When we differentiate with respect to a tensor, we implicitly
mean for a tensor of fixed shape. Likewise, when a program outputs a tensor in
our language, its size cannot depend on real inputs, as we do not have non-
continuous primitives from the reals to integers. Therefore, if all the
integers inputs of a program are fixed, its output tensor will have fixed
shape, which can be captured by a canonical projection. Another way to say it
is that for programs x: tensor m $\vdash$ e: tensor n, there is a canonical
projection that is injective on the image of $\llbracket$e$\rrbracket$.
Therefore, we define differentials for functions outputting tensors as
follows:

###### Definition 6.3.

The function $dg:\mathbb{R}^{n}\times\mathbb{R}^{n}\to\llbracket$tensor
m$\rrbracket$ is the differential of $f:\mathbb{R}^{n}\to\llbracket$tensor
m$\rrbracket$ if, for all canonical projections $\pi:\llbracket$tensor
m$\rrbracket\to\mathbb{R}^{k}$, $\pi\circ dg$ is the differential of $\pi\circ
f$.

We can extend this definition to functions taking a tensor as input:

###### Definition 6.4.

The function $df:\llbracket$tensor n$\rrbracket\times\llbracket$tensor
n$\rrbracket\to\llbracket$tensor m$\rrbracket$ is the differential of
$f:\llbracket$tensor n$\rrbracket\to\llbracket$tensor m$\rrbracket$ if, for
all canonical injections $i:\mathbb{R}^{k}\to\llbracket$tensor n$\rrbracket$,
$dg\circ(i\times i)$ is the differential of $f\circ i$, where $i\times
i:(x,y)\mapsto(i(x),i(y))$.

Similarly to the case of usual differentiability, we can extend the definition
to functions with multiple inputs and outputs in a pointwise style. That is:

###### Definition 6.5.

A function $dg:\llbracket T\rrbracket\times\llbracket T\rrbracket\to\llbracket
T1\rrbracket\times\llbracket T2\rrbracket$ is the differential of a function
$f:\llbracket T\rrbracket\to\llbracket T1\rrbracket\times\llbracket
T2\rrbracket$ iff $\pi_{i}\circ dg$ is the differential of $\pi_{i}\circ f$. A
function $dg:\llbracket T1\rrbracket\times\llbracket
T1\rrbracket\times\llbracket T2\rrbracket\times\llbracket T2\rrbracket$ is the
differential of a function $f:\llbracket T1\rrbracket\times\llbracket
T2\rrbracket\to\llbracket T\rrbracket$ iff for all $x\in\llbracket
T1\rrbracket$, $dg(x,0,-,-)$ is the differential of $f(x,-):\llbracket
T2\rrbracket\to\llbracket T\rrbracket$, and for all $y\in\llbracket
T2\rrbracket$, $dg(-,-,y,0)$ is the differential of $f(-,y):\llbracket
T1\rrbracket\to\llbracket T\rrbracket$.

We can now state that FAD computes differentials:

###### Theorem 6.6.

Let $\Gamma\vdash$ e: D, where $\Gamma$ only contains tensor types. Then,
$\llbracket$$\mathcal{F}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$$\rrbracket$
is the differential of $\llbracket$e$\rrbracket$.

The proof now follows a simple logical relations argument, as in (Huot et al.,
2020). The logical predicates are defined by

$\mathcal{R}_{\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}real}}$ | $=$ | $\\{(f:\mathbb{R}\to\mathbb{R},g:\mathbb{R}\to\mathbb{R})~{}\mid~{}f\in\mathcal{P}_{\mathbb{R}}^{\mathbb{R}},g=f^{\prime}\\}$
---|---|---
$\mathcal{R}_{\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}int}}$ | $=$ | $\\{(f:\mathbb{R}\to\mathbb{N},g:\mathbb{R}\to\mathbb{R})~{}\mid~{}f\text{ constant},g=\lambda x.0\\}$
$\mathcal{R}_{\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}bool}}$ | $=$ | $\\{(f:\mathbb{R}\to\mathbb{B},g:\mathbb{R}\to\mathbb{R})~{}\mid~{}f\text{ constant},g=\lambda x.0\\}$
$\mathcal{R}_{\text{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}\\{I -> D\\}}}}$ | $=$ | $\big{\\{}(f:\mathbb{R}\to\llbracket\text{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}\\{I -> D\\}}}\rrbracket,g:\mathbb{R}\to\llbracket\text{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}\\{I -> D\\}}}\rrbracket~{}\mid~{}f\in\mathcal{P}_{\mathbb{R}}^{\llbracket\\{I\to D\\}\rrbracket},g=df(-,1),$
|  |  $df$ differential of $f\\}$

After showing the fundamental lemma for logical relations, where one needs to
show that the interpretation of the forward differentiation of every construct
of the language preserves the logical predicate (by post-composition), we
conclude that the interpretation of every program preserves the logical
predicate. We conclude the theorem as a corollary of the fundamental lemma in
the same way as in (Huot et al., 2020). More precisely, for the predicate
$\mathcal{R}_{\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}real}}$,
we choose $f=id$. For the predicate
$\mathcal{R}_{\text{{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}\\{I
-> D\\}}}}$, for any linear injection $\mathbb{R}^{k}\to\llbracket${I ->
D}$\rrbracket$, and any $x_{1},\ldots,x_{k-1}\in\mathbb{R}^{k-1}$, we take
$f_{i}=i\circ g_{i}$ for $g_{i}=\lambda x.(x_{1},\ldots,x,\ldots,x_{k-1})$,
where $x$ is at position $i$.

###  6.4. Correctness of Tensorized Forward-Mode AD

A change of perspectives: $\tau$-differentials. Given a differentiable
function $f:\mathbb{R}^{n}\to\mathbb{R}^{m}$, we can equivalently see its
differential as a higher-order function
$df:\mathbb{R}^{n}\times(\mathbb{R}\multimap\mathbb{R}^{n})\to(\mathbb{R}\multimap\mathbb{R}^{m})$,
where $A\multimap B$ is the real vector space of $\mathbb{R}$-linear functions
from $A$ to $B$. Given bases for the spaces $A,B$, we identify linear
functions $A\multimap B$ with matrices. Seen this way, we have
$df(x,v)=Jf_{x}\cdot v$, where $v$ is an $n\times 1$ matrix, and $\cdot$ is
now matrix-matrix multiplication. Let $\tau=\mathbb{R}^{k}$. We can generalize
the previous definition by seeing the differential of
$f:\mathbb{R}^{n}\to\mathbb{R}^{m}$ as a function
$df:\mathbb{R}^{n}\times(\tau\multimap\mathbb{R}^{n})\to(\tau\multimap\mathbb{R}^{m})$,
still defined by $df(x,v)=Jf_{x}\cdot v$. We call these functions
$\tau$-differentials. The functions $f\in\tau\multimap\mathbb{R}^{n}$ can be
seen as the $(k,1)$ velocities from (Betancourt, 2018). For finite dimensional
vector spaces $A,B$, the vector space $A\multimap B$ is isomorphic to
$A\otimes B$, where $A\otimes B$ is the tensor product of vector spaces.
Therefore, an equivalent way to define the $\tau$-differential of
$f:\mathbb{R}^{n}\to\mathbb{R}^{m}$ is as a function
$df:\mathbb{R}^{n}\times(\mathbb{R}^{n}\otimes\tau)\to(\mathbb{R}^{m}\otimes\tau)$.
This is, in essence, what our tensorized forward-mode AD computes. For
instance, if $\tau=\mathbb{R}^{n}$, by choosing $v$ to be the identity matrix,
we have that $df(x,v)=J_{x}f$ is the entire Jacobian of $f$, whereas usual
forward-mode only computes a column of the Jacobian matrix, by choosing $v$ to
be a hot vector. Such a view on differentials extends to functions $f$ with a
tensor input or output, as follows:

###### Definition 6.7.

The function
$df:\mathbb{R}^{n}\times(\tau\multimap\mathbb{R}^{n})\to(\tau\multimap\llbracket$tensor
m$\rrbracket)$ is the $\tau$-differential of
$f:\mathbb{R}^{n}\to\llbracket$tensor m$\rrbracket$ if, for all canonical
projections $\pi:\llbracket$tensor m$\rrbracket\to\mathbb{R}^{k}$, and
canonical injections $i:\mathbb{R}^{p}\to\tau$, $\lambda(x,l).\pi\circ
df(x,l\circ i)$ is the $\mathbb{R}^{p}$-differential of $\pi\circ f$.

Likewise, The function $df:\llbracket$tensor
n$\rrbracket\times(\tau\multimap\llbracket$tensor
n$\rrbracket)\to(\tau\multimap\llbracket$tensor m$\rrbracket)$ is the
$\tau$-differential of $f:\llbracket$tensor n$\rrbracket\to\llbracket$tensor
m$\rrbracket$ if, for all canonical projections $\pi:\llbracket$tensor
m$\rrbracket\to\mathbb{R}^{k}$, and canonical injections
$i:\mathbb{R}^{p}\to\tau,j:\mathbb{R}^{q}\to\llbracket$tensor n$\rrbracket$,
$\lambda(x,l).\pi\circ df(j(x),i^{*}\circ l\circ i)$ is the
$\mathbb{R}^{p}$-differential of $\pi\circ f\circ j$, where $i^{*}$ is the
linear adjoint of $i$.

It then further extends to functions with multiple inputs or outputs,
following Definition 6.5.

We can now state that tensorized forward-mode AD computes differentials:

###### Theorem 6.8.

Let $\tau$ be a tensor type. Let $\Gamma\vdash$ e: D, where $\Gamma$ only
contains tensor types. Then,
$\llbracket$$\mathcal{D}_{\tau}\llbracket\textrm{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}{}e}\rrbracket$$\rrbracket$
is the $\tau$-differential of $\llbracket$e$\rrbracket$.

The proof follows a similar logical relations argument as Theorem 6.6, where
the logical predicates are slightly changed to take the type $\tau$ into
account. For instance

$\displaystyle\mathcal{R}_{\textrm{\color[rgb]{0,0,0.5}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0.5}{}real}}^{\tau}=$
$\displaystyle\big{\\{}f:\llbracket\tau\rrbracket\to\mathbb{R},g:\mathbb{R}\times(\llbracket\tau\rrbracket\multimap\llbracket\tau\rrbracket)\to(\llbracket\tau\rrbracket\multimap\mathbb{R})~{}\mid~{}f\text{
morphism of diffeological spaces},$ $\displaystyle g\text{ $\tau$-differential
of }f\big{\\}}$

##  7\. Experimental Results

In this section, we see the effectiveness of $\nabla$SD in practice by using
real-world and synthetic tensors. We answer the following research questions:

  * •

How does $\nabla$SD perform for tensor kernels over real-world sparse matrices
in comparison with the state-of-the-art AD frameworks? Does $\nabla$SD scale
to large sparse matrices?

  * •

What is the impact of different physical storage formats on the performance of
$\nabla$SD?

  * •

What is the impact of different optimizations on the performance of
$\nabla$SD?

  * •

For which densities does it make sense to use sparse representations and
compute the gradients using $\nabla$SD?

###  7.1. Experimental Setup

The experiments were conducted on a MacBook Pro featuring a dual-core Intel
Core i7 CPU clocked at 3.5 GHz with 16 GB of LPDDR3 RAM (2133MHz), running
macOS Big Sur 11.5.2. To compile the generated C++ code, we employed CLang
1205.0.22.11 and Python 3.7.4 for executing the Python code. We used
TensorFlow 2.11.0 (XLA-enabled) and PyTorch 1.13.1.

Table 1. Real-world matrices used in the end-to-end experiments. Matrix | Dimensions | Density | # NNZ | Description
---|---|---|---|---
bcspwr10 |  5.3K $\times$ 5.3K | $2^{-9.3}$ | 22K | Power Network
nopoly |  10K $\times$ 10K | $2^{-9.5}$ | 70K | Undirected Weighted Graph
pdb1HYS |  36K $\times$ 36K | $2^{-8.2}$ | 2.19M | Undirected Weighted Graph
rma10 |  46K $\times$ 46K | $2^{-9.8}$ | 2.37M | Computational Fluid Dynamics
cant |  62K $\times$ 62K | $2^{-9.9}$ | 2.03M | Finite Element Method
consph |  83K $\times$ 83K | $2^{-10.1}$ | 3.05M | Finite Element Method
cop20k_A |  121K $\times$ 121K | $2^{-12.4}$ | 1.36M | Finite Element Method

We consider the following sparse matrix kernels:

  * •

BATAX. The BATAX kernel computed using the following formula (Nelson et al.,
2015): $f(j)=\Sigma_{i,k}\beta\cdot A(i,j)\cdot A(i,k)\cdot X(k)$. We consider
its gradient with respect to the vector: $\frac{\partial f}{\partial X}$.

  * •

SMMM. The summation of elements of matrix-matrix-multiplication (Schleich et
al., 2022): $f=\Sigma_{i,j,k}A(i,k)\cdot B(k,j)$. We consider its gradient
with respect to the second matrix: $\frac{\partial f}{\partial B}$.

  * •

SMVM. The summation of the elements of a matrix-vector-multiplication (Smeding
and Vákár, 2023): $f=\Sigma_{i,j}A(i,j)\cdot X(j)$. We consider its gradient
with respect to the vector: $\frac{\partial f}{\partial X}$.

The first two kernels produce a matrix result, while the last one produces a
vector.

We also consider the following vector kernels from
$\text{d}\widetilde{\textsc{f}}$ ((Shaikhha et al., 2019)):

  * •

VVA. The addition of two vectors: $f(i)=V_{1}(i)+V_{2}(i)$. We consider its
gradient with respect to the first vector: $\frac{\partial f}{\partial
V_{1}}$.

  * •

VVD. The dot product of two vectors: $f=\Sigma_{i}V_{1}(i)\cdot V_{2}(i)$. We
consider its gradient with respect to the first vector: $\frac{\partial
f}{\partial V_{1}}$.

  * •

VSM. The vector-scalar multiplication: $f(i)=V(i)\cdot s\cdot s$. The scalar
value was multiplied by itself to make the differentiation more complicated.
We consider its gradient with respect to the scalar value: $\frac{\partial
f}{\partial s}$.

The first kernel produces a matrix result, while the last two produce a
vector.

As competitors for the matrix kernels, we consider the following systems:

  * •

$\nabla$SD (dict): Generated C++ code by $\nabla$SD from tensorized forward-
mode AD that uses a nested dictionary.

  * •

$\nabla$SD (arr): Generated C++ code by $\nabla$SD from tensorized forward-
mode AD that also benefits from the physical storage specified by the sparse
representations.

  * •

TensorFlow: The reverse-mode AD of the TensorFlow framework, using its
gradient and jacobian APIs.

  * •

PyTorch: The reverse-mode AD of the PyTorch framework, using its jacobian API.

For the vector kernels, in addition to the above competitors, we consider the
following systems:111The generated code were made available in
$\text{d}\widetilde{\textsc{f}}$ here:
https://github.com/amirsh/autodiff2/tree/master/icfp19

  * •

Tapenade (R): Generated C code by Tapenade using reverse-mode AD.

  * •

Tapenade (F): Generated C code by Tapenade using forward-mode AD.

  * •

$\text{d}\widetilde{\textsc{f}}$: Generated C code by
$\text{d}\widetilde{\textsc{f}}$ using its forward-mode AD.

  * •

$\text{d}\widetilde{\textsc{f}}$ \+ DPS: Generated C code by
$\text{d}\widetilde{\textsc{f}}$ using its forward-mode AD with the DPS
optimization.

Both real-world and synthetic datasets were utilized in our study. To obtain
the former, we gathered seven sparse matrices from the SuiteSparse Matrix
Collection (Davis and Hu, 2011). Table 1 summarizes these datasets. For
synthetic data, we created random matrices and vectors with different density
and dimension configurations.

In all these benchmarks, we consider the gradient with respect to a dense
variable. All experiments use a single-core. We take the average time of
running five runs.

|  |
---|---|---
Figure 11. Performance results for differentiation of kernels over real-world
sparse matrices. $\nabla$SD uses tensorized forward mode over sparse matrices,
while TensorFlow and PyTorch use reverse-mode over dense matrices.

###  7.2. Benchmarks over Real-world Datasets

In this section, we consider real-world sparse matrices. We compare $\nabla$SD
with TensorFlow and PyTorch. Although both systems support sparse operations,
none support AD over sparse matrices.

Figure 11 shows the results for the real-world sparse matrices. We make the
following observations. First, for all kernels $\nabla$SD outperforms
TensorFlow and PyTorch thanks to leveraging sparse representations. Second,
the generated code by $\nabla$SD that leverages the array-based physical
representation can run faster than the version that uses nested dictionaries
in SMVM and SMMM. Finally, $\nabla$SD scales to large sparse matrices, as
opposed to TensorFlow and PyTorch, which do not manage to process such
matrices due to storing the entire matrix including the zero elements.

###  7.3. Benchmarks over Synthetic Datasets

In this section, we use synthetic datasets to better analyse the impact of the
sparsity and dimensions. We consider two scenarios: (1) fixing the dimension
and varying the sparsity, and (2) fixing the sparsity and varying the
dimension.

|  |
---|---|---
|  |
Figure 12. Performance results for differentiation of different sparse matrix kernels by varying sparsity. |  |
---|---|---
|  |
Figure 13. Performance results for differentiation of different sparse matrix
kernels by varying dimensions.

Matrix Kernels. Figure 12 shows the results for matrix kernels by varying the
sparsity for small and large matrices. We make the following observations.
First, for low densities, which are in the same range as the real-world
datasets, we see a clear advantage for $\nabla$SD over TensorFlow and PyTorch.
Second, the performance of TensorFlow and PyTorch is not dependent on the
sparsity, as expected. As the matrices get denser, the gap between these two
frameworks and $\nabla$SD becomes smaller. Finally, for the BATAX kernel,
there is an advantage for the dictionary-based representation over the array-
based one. However, for the other two kernels, especially for larger matrices,
the array-based representation performs better.

Figure 13 shows the results by varying the dimension for two different
sparsities. Matrices with high densities show a better performance for
TensorFlow and PyTorch over $\nabla$SD. The gap even widens for larger
dimensions. However, we observe the opposite impact for a lower density.

|  |
---|---|---
|  |
Figure 14. Performance results for differentiation of different sparse vector kernels by varying sparsity. |  |
---|---|---
|  |
Figure 15. Performance results for differentiation of different sparse vector
kernels by varying dimensions.

Vector Kernels. Figure 14 and Figure 15 show the results for vector kernels by
varying the sparsity and density. For smaller dimensions, there is no clear
advantage for $\nabla$SD; for VVA $\text{d}\widetilde{\textsc{f}}$ performs
better than $\nabla$SD. However, for larger dimensions, we observe a clear
advantage for $\nabla$SD, especially for matrices with lower densities. For
VSM, we observe a clear advantage for forward-mode-based systems over reverse-
mode-based ones; for lower densities $\nabla$SD outperforms the rest, whereas
for higher densities $\text{d}\widetilde{\textsc{f}}$ is the most performant
system.

|
---|---
|
Figure 16. Impact of the normalization and DPS over the performance of
$\nabla$SD on the BATAX kernel.

###  7.4. Impact of Optimizations

In this section, we investigate the impact of optimizations on the performance
of $\nabla$SD. For each data representation, we consider the following
alternatives for the generated C++ code: (1) without low-level
transformations, (2) with multiplication normalization, and (3) with
multiplication normalization and removal of intermediate tensors in the nested
loops using DPS.

Figure 16 shows the results for the BATAX kernel. Similarly to the previous
sets of experiments, we see a slight advantage for the dictionary-based
representation over the array-based one. This is due to the intermediate
dictionary created in the array-based version (cf. Section 5.5). Furthermore,
each of the optimizations has a positive impact on performance. Especially for
lower densities, the impact of optimizations becomes more pronounced.

##  8\. Related Work

Automatic Differentiation. There are several existing automatic
differentiation (AD) frameworks and libraries for imperative and functional
programming languages. ADIFOR (Bischof et al., 1992) and Tapenade (Hascoët and
Pascual, 2013) perform AD for Fortran and C programs respectively, while Adept
(Hogan, 2014) and ADIC (Narayanan et al., 2010) perform AD for C++ using
expression templates. ForwardDiff (Revels et al., 2016) uses vector forward-
mode AD for differentiating Julia programs, while DiffSharp (Baydin et al.,
2015) is an AD library implemented in F# that provides both forward-mode and
reverse-mode. Stalingrad (Pearlmutter and Siskind, 2008) is an optimizing
compiler for a dialect of Scheme with a first-class AD operator and supports
both forward mode and reverse mode of AD. Similarly, Karczmarczuk
(Karczmarczuk, 1998) presents a Haskell implementation for both forward and
reverse mode AD, and Elliott (Elliott, 2018) provides a generalization of AD
based on category theory for implementing both forward and reverse-mode AD.
There has been recent efforts on providing correct and asymptotically
efficient reverse-mode AD for functional languages (Krawiec et al., 2022;
Smeding and Vákár, 2023; Radul et al., 2023), the ideas of which are
implemented in JAX (Bradbury et al., 2018; Frostig et al., 2018) and Dex
(Paszke et al., 2021).

Machine learning libraries like Tensorflow and PyTorch are implemented based
on tensor abstractions. These systems come with a predefined set of efficient
kernels for manipulating tensors and can use compilation backends for further
optimization. Lantern (Wang et al., 2019) uses multi-stage programming to
perform reverse-mode AD. However, none of the mentioned frameworks supports AD
for sparse data structures with irregular storage formats (TensorFlow, 2023;
JAX, 2023; PyTorch, 2023). There have been efforts on statically incorporating
sparsities, however, this requires manually specifying the sparsity patterns
by the programmers (Peng and Dubach, 2023) and do not work for large sparse
matrices with arbitrary patterns (Ghorbani et al., 2022).

Sparse Tensor Algebra. Sparse tensor algebra has been the focus of much
research and development in recent years, leading to the emergence of several
frameworks and systems designed to support it. TACO (Kjolstad et al., 2017;
Chou et al., 2018) is a system capable of handling both sparse and dense
computations over tensor algebra. Another noteworthy framework is the sparse
polyhedral framework (Strout et al., 2018), which extends the capabilities of
polyhedral compilation to support sparse tensor algebra. In addition, packages
such as SciPy (Virtanen et al., 2019), TensorFlow, PyTorch, and the MATLAB
Tensor Toolbox (Bader and Kolda, 2007) offer support for various sparse matrix
representations, enabling efficient computation on sparse tensor data
structures commonly found in scientific and engineering applications.

Despite the progress made in this area, automatic differentiation for sparse
tensors is still not widely supported. One of the few recent efforts (Nytko et
al., 2022) provides manual gradients for a limited set of kernels. The primary
challenge in differentiating sparse tensors is the irregular data
representation, making differentiation a complex process. To address this
issue, we propose separating the logical sparse tensor representation from its
physical storage format, allowing for more efficient and effective
differentiation.

##  9\. Conclusion & Outlook

In this paper, we present $\nabla$SD, the first framework that supports
automatic differentiation for sparse tensors. Our main insight is to separate
the logical concerns from the physical data storage representations. We
provide a tensorized forward-mode transformation over the logical fragment of
SDQLite, a recently introduced language that fuses the physical storage of
sparse tensors with the logical specification of kernels. By benefiting from
the algebraic optimizations globally applied using equality saturation we
further improve the performance of differentiated programs and fuse their
physical representations. We show experimentally that our framework
outperforms the state-of-the-art differentiable tensor frameworks over both
real-world and synthetic datasets.

For the future, we plan to add support for reverse-mode AD (by borrowing ideas
from (Radul et al., 2023; Berg et al., 2022)). In addition, we plan to add the
support for scheduling transforms to add the GPU backend (Senanayake et al.,
2020). This way we can use our framework for training deep, yet sparse
learning models such as Graph Neural Networks (GNNs) (Nytko et al., 2022).
Finally, we aim to add the support for the entire SDQL (Shaikhha et al., 2022)
(including boolean and integer semi-rings required for set and bag
construction as well as tupling constructs) in order to provide automatic
differentiation for relational databases and hybrid relational-linear algebra
workloads.

## Acknowledgments

The first author thanks Huawei for their support of the distributed data
management and processing laboratory at the University of Edinburgh. The
second author is supported by a Royal Society University Research Fellowship.

## References

  * (1)
  * rbd (2023) 2023\.  Fast & memory efficient hashtable based on robin hood hashing for C++.  https://github.com/martinus/robin-hood-hashing.
  * Abadi et al. (2015) Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. 2015\.  TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems.  https://www.tensorflow.org/ Software available from tensorflow.org.
  * Bader and Kolda (2007) Brett W. Bader and Tamara G. Kolda. 2007.  Efficient MATLAB Computations with Sparse and Factored Tensors.  _SIAM J. Sci. Comput._ 30, 1 (2007), 205–231.  https://doi.org/10.1137/060676489
  * Baydin et al. (2015) Atilim Gunes Baydin, Barak A. Pearlmutter, and Jeffrey Mark Siskind. 2015.  DiffSharp: Automatic Differentiation Library.  _CoRR_ abs/1511.07727 (2015).  arXiv:1511.07727 http://arxiv.org/abs/1511.07727
  * Berg et al. (2022) Birthe van den Berg, Tom Schrijvers, James McKinna, and Alexander Vandenbroucke. 2022.  Forward-or Reverse-Mode Automatic Differentiation: What’s the Difference?  _arXiv preprint arXiv:2212.11088_ (2022).
  * Betancourt (2018) Michael Betancourt. 2018\.  A geometric theory of higher-order automatic differentiation.  _arXiv preprint arXiv:1812.11592_ (2018).
  * Bischof et al. (1992) Christian H. Bischof, Alan Carle, George F. Corliss, Andreas Griewank, and Paul D. Hovland. 1992\.  ADIFOR - Generating Derivative Codes form Fortran Programs.  _Sci. Program._ 1, 1 (1992), 11–29.  https://doi.org/10.1155/1992/717832
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018\.  _JAX: composable transformations of Python+NumPy programs_.  http://github.com/google/jax
  * Chou et al. (2018) Stephen Chou, Fredrik Kjolstad, and Saman P. Amarasinghe. 2018\.  Format abstraction for sparse tensor algebra compilers.  _Proc. ACM Program. Lang._ 2, OOPSLA (2018), 123:1–123:30.  https://doi.org/10.1145/3276493
  * Davis and Hu (2011) Timothy A. Davis and Yifan Hu. 2011.  The university of Florida sparse matrix collection.  _ACM Trans. Math. Softw._ 38, 1 (2011), 1:1–1:25.  https://doi.org/10.1145/2049662.2049663
  * Elliott (2018) Conal Elliott. 2018\.  The simple essence of automatic differentiation.  _Proc. ACM Program. Lang._ 2, ICFP (2018), 70:1–70:29.  https://doi.org/10.1145/3236765
  * Flanagan et al. (1993) Cormac Flanagan, Amr Sabry, Bruce F. Duba, and Matthias Felleisen. 1993.  The Essence of Compiling with Continuations. In _Proceedings of the ACM SIGPLAN’93 Conference on Programming Language Design and Implementation (PLDI), Albuquerque, New Mexico, USA, June 23-25, 1993_ , Robert Cartwright (Ed.). ACM, 237–247.  https://doi.org/10.1145/155090.155113
  * Frostig et al. (2018) Roy Frostig, Matthew James Johnson, and Chris Leary. 2018\.  Compiling machine learning programs via high-level tracing.  _Systems for Machine Learning_ 4, 9 (2018).
  * Ghorbani et al. (2022) Mahdi Ghorbani, Mathieu Huot, Shideh Hashemian, and Amir Shaikhha. 2022.  Compiling Structured Tensor Algebra.  _arXiv preprint arXiv:2211.10482_ (2022).
  * Hascoët and Pascual (2013) Laurent Hascoët and Valérie Pascual. 2013.  The Tapenade automatic differentiation tool: Principles, model, and specification.  _ACM Trans. Math. Softw._ 39, 3 (2013), 20:1–20:43.  https://doi.org/10.1145/2450153.2450158
  * Hogan (2014) Robin J. Hogan. 2014\.  Fast Reverse-Mode Automatic Differentiation using Expression Templates in C++.  _ACM Trans. Math. Softw._ 40, 4 (2014), 26:1–26:16.  https://doi.org/10.1145/2560359
  * Huot and Shaikhha (2022) Mathieu Huot and Amir Shaikhha. 2022.  Denotationally Correct, Purely Functional, Efficient Reverse-mode Automatic Differentiation.  _CoRR_ abs/2212.09801 (2022).  https://doi.org/10.48550/arXiv.2212.09801 arXiv:2212.09801
  * Huot et al. (2020) Mathieu Huot, Sam Staton, and Matthijs Vákár. 2020\.  Correctness of Automatic Differentiation via Diffeologies and Categorical Gluing.. In _FoSSaCS_. 319–338.
  * Iglesias-Zemmour (2013) Patrick Iglesias-Zemmour. 2013\.  _Diffeology_. Vol. 185.  American Mathematical Soc.
  * JAX (2023) JAX. 2023.  GitHub Issue #13118.  https://github.com/google/jax/issues/13118.  Accessed: 2023-01-30.
  * Karczmarczuk (1998) Jerzy Karczmarczuk. 1998\.  Functional Differentiation of Computer Programs. In _Proceedings of the third ACM SIGPLAN International Conference on Functional Programming (ICFP ’98), Baltimore, Maryland, USA, September 27-29, 1998_ , Matthias Felleisen, Paul Hudak, and Christian Queinnec (Eds.). ACM, 195–203.  https://doi.org/10.1145/289423.289442
  * Khan and Barton (2015) Kamil A. Khan and Paul I. Barton. 2015.  A vector forward mode of automatic differentiation for generalized derivative evaluation.  _Optim. Methods Softw._ 30, 6 (2015), 1185–1212.  https://doi.org/10.1080/10556788.2015.1025400
  * Kjolstad et al. (2017) Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. 2017.  The Tensor Algebra Compiler.  _Proc. ACM Program. Lang._ 1, OOPSLA, Article 77 (Oct. 2017), 29 pages.  https://doi.org/10.1145/3133901
  * 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.  https://doi.org/10.1145/3498710
  * Minamide (1998) Yasuhiko Minamide. 1998\.  A functional representation of data structures with a hole. In _Proceedings of the 25th ACM SIGPLAN-SIGACT symposium on Principles of programming languages_. 75–84.
  * Moré (2006) Jorge J Moré. 2006\.  The Levenberg-Marquardt algorithm: implementation and theory. In _Numerical Analysis: Proceedings of the Biennial Conference Held at Dundee, June 28–July 1, 1977_. Springer, 105–116.
  * Narayanan et al. (2010) Sri Hari Krishna Narayanan, Boyana Norris, and Beata Winnicka. 2010\.  ADIC2: Development of a component source transformation system for differentiating C and C++. In _Proceedings of the International Conference on Computational Science, ICCS 2010, University of Amsterdam, The Netherlands, May 31 - June 2, 2010_ _(Procedia Computer Science, Vol. 1)_ , Peter M. A. Sloot, G. Dick van Albada, and Jack J. Dongarra (Eds.). Elsevier, 1845–1853.  https://doi.org/10.1016/j.procs.2010.04.206
  * Nelson et al. (2015) Thomas Nelson, Geoffrey Belter, Jeremy G. Siek, Elizabeth R. Jessup, and Boyana Norris. 2015\.  Reliable Generation of High-Performance Matrix Algebra.  _ACM Trans. Math. Softw._ 41, 3 (2015), 18:1–18:27.  https://doi.org/10.1145/2629698
  * Nytko et al. (2022) Nicolas Nytko, Ali Taghibakhshi, Tareq Uz Zaman, Scott MacLachlan, Luke N Olson, and Matt West. 2022.  Optimized Sparse Matrix Operations for Reverse Mode Automatic Differentiation.  _arXiv preprint arXiv:2212.05159_ (2022).
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019.  PyTorch: An Imperative Style, High-Performance Deep Learning Library. In _Advances in Neural Information Processing Systems 32_. Curran Associates, Inc., 8024–8035.
  * Paszke et al. (2021) Adam Paszke, Daniel D Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Matthew J Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin. 2021.  Getting to the point: index sets and parallelism-preserving autodiff for pointful array programming.  _Proceedings of the ACM on Programming Languages_ 5, ICFP (2021), 1–29.
  * Pearlmutter and Siskind (2008) Barak A. Pearlmutter and Jeffrey Mark Siskind. 2008.  Reverse-mode AD in a functional framework: Lambda the ultimate backpropagator.  _ACM Trans. Program. Lang. Syst._ 30, 2 (2008), 7:1–7:36.  https://doi.org/10.1145/1330017.1330018
  * Peng and Dubach (2023) Mai Jacob Peng and Christophe Dubach. 2023.  LAGrad: Statically Optimized Differentiable Programming in MLIR. In _Proceedings of the 32nd ACM SIGPLAN International Conference on Compiler Construction_ (Montréal, QC, Canada) _(CC 2023)_. Association for Computing Machinery, New York, NY, USA, 228–238.  https://doi.org/10.1145/3578360.3580259
  * PyTorch (2023) PyTorch. 2023.  GitHub Issue #12498.  https://github.com/pytorch/pytorch/issues/12498.  Accessed: 2023-01-30.
  * Radul et al. (2023) Alexey Radul, Adam Paszke, Roy Frostig, Matthew J. Johnson, and Dougal Maclaurin. 2023.  You Only Linearize Once: Tangents Transpose to Gradients.  _Proc. ACM Program. Lang._ 7, POPL (2023), 1246–1274.  https://doi.org/10.1145/3571236
  * Revels et al. (2016) Jarrett Revels, Miles Lubin, and Theodore Papamarkou. 2016\.  Forward-Mode Automatic Differentiation in Julia.  _CoRR_ abs/1607.07892 (2016).  arXiv:1607.07892 http://arxiv.org/abs/1607.07892
  * Schleich et al. (2022) Maximilian Schleich, Amir Shaikhha, and Dan Suciu. 2022\.  Optimizing Tensor Programs on Flexible Storage.  _arXiv preprint arXiv:2210.06267_ (2022).
  * Senanayake et al. (2020) Ryan Senanayake, Changwan Hong, Ziheng Wang, Amalee Wilson, Stephen Chou, Shoaib Kamil, Saman Amarasinghe, and Fredrik Kjolstad. 2020.  A Sparse Iteration Space Transformation Framework for Sparse Tensor Algebra.  _Proc. ACM Program. Lang._ 4, OOPSLA, Article 158 (Nov. 2020), 30 pages.  https://doi.org/10.1145/3428226
  * Shaikhha (2022) Amir Shaikhha. 2022\.  Deep Fusion for Efficient Nested Recursive Computations. In _Proceedings of the 21st ACM SIGPLAN International Conference on Generative Programming: Concepts and Experiences_. 33–44.
  * Shaikhha et al. (2017) Amir Shaikhha, Andrew Fitzgibbon, Simon Peyton Jones, and Dimitrios Vytiniotis. 2017.  Destination-passing style for efficient memory management. In _Proceedings of the 6th ACM SIGPLAN International Workshop on Functional High-Performance Computing_. 12–23.
  * Shaikhha et al. (2019) Amir Shaikhha, Andrew W. Fitzgibbon, Dimitrios Vytiniotis, and Simon Peyton Jones. 2019.  Efficient differentiable programming in a functional array-processing language.  _Proc. ACM Program. Lang._ 3, ICFP (2019), 97:1–97:30.  https://doi.org/10.1145/3341701
  * Shaikhha et al. (2022) Amir Shaikhha, Mathieu Huot, Jaclyn Smith, and Dan Olteanu. 2022\.  Functional collection programming with semi-ring dictionaries.  _Proc. ACM Program. Lang._ 6, OOPSLA1 (2022), 1–33.  https://doi.org/10.1145/3527333
  * Smeding and Vákár (2023) Tom Smeding and Matthijs Vákár. 2023.  Efficient Dual-Numbers Reverse AD via Well-Known Program Transformations.  _Proc. ACM Program. Lang._ 7, POPL (2023), 1573–1600.  https://doi.org/10.1145/3571247
  * Strout et al. (2018) Michelle Mills Strout, Mary W. Hall, and Catherine Olschanowsky. 2018.  The Sparse Polyhedral Framework: Composing Compiler-Generated Inspector-Executor Code.  _Proc. IEEE_ 106, 11 (2018), 1921–1934.  https://doi.org/10.1109/JPROC.2018.2857721
  * Tang et al. (2020) Xuan Tang, Teseo Schneider, Shoaib Kamil, Aurojit Panda, Jinyang Li, and Daniele Panozzo. 2020\.  EGGS: Sparsity-Specific Code Generation.  _Comput. Graph. Forum_ 39, 5 (2020), 209–219.  https://doi.org/10.1111/cgf.14080
  * Tate et al. (2009) Ross Tate, Michael Stepp, Zachary Tatlock, and Sorin Lerner. 2009\.  Equality saturation: a new approach to optimization. In _Proceedings of the 36th annual ACM SIGPLAN-SIGACT symposium on Principles of programming languages_. 264–276.
  * TensorFlow (2023) TensorFlow. 2023\.  GitHub Issue #43497.  https://github.com/tensorflow/tensorflow/issues/43497.  Accessed: 2023-01-30.
  * Virtanen et al. (2019) Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C. J. Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy. 2019.  SciPy 1.0-Fundamental Algorithms for Scientific Computing in Python.  _CoRR_ abs/1907.10121 (2019).  arXiv:1907.10121 http://arxiv.org/abs/1907.10121
  * Wang et al. (2019) Fei Wang, Daniel Zheng, James M. 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 (2019), 96:1–96:31.  https://doi.org/10.1145/3341700
  * Willsey et al. (2021) Max Willsey, Chandrakana Nandi, Yisu Remy Wang, Oliver Flatt, Zachary Tatlock, and Pavel Panchekha. 2021\.  egg: Fast and extensible equality saturation.  _Proc. ACM Program. Lang._ 5, POPL (2021), 1–29.  https://doi.org/10.1145/3434304