0.0.0.1 Keywords.

reverse-mode automatic differentiation, probabilistic programming, algebraic effects, categorical semantics.

1 Introduction↩︎

Probabilistic programming languages provide constructs for specifying probabilistic models directly and for manipulating their statistical properties within the language. Their promise is to make common tasks in probabilistic data analysis, such as sampling, inference, and the optimisation of expected objectives, more compositional and closer to the structure of the models themselves.

Reverse-mode automatic differentiation (AD), often called backpropagation, underlies much of modern optimisation and machine learning. Its practical importance comes from the need to compute gradients of scalar or low-dimensional objectives with respect to large numbers of parameters [1]. Compared to the simpler forward-mode of AD, reverse-mode AD inverts data flow: forward-mode propagates sensitivities along the computation, whereas reverse-mode propagates and accumulates cotangents against the direction of computation. This makes the design and semantics of reverse-mode automatic differentiation substantially more delicate.

Combinatory Homomorphic Automatic Differentiation (CHAD) gives a particularly clean categorical account of reverse-mode AD. In the CHAD framework, automatic differentiation is not specified as an ad hoc syntactic traversal, but as a semantics-driven program transformation: a structure-preserving map out of the syntactic categorical model of the source language [2], [3]. This account is not merely explanatory, but it guides us in how to differentiate more complex forms of computation. Once the relevant semantic structure has been identified, the corresponding program transformation is obtained by following the universal structure of the language. In the deterministic setting, this has led both to correctness proofs by categorical logical relations and to efficient implementations via familiar functional-programming techniques, such as sparse representations and defunctionalisation or closure conversion [4].

The goal of the present paper is to explain how the reverse-mode CHAD transformation applies to programs with finite-support probabilistic effects and differentiable weights. In this impure setting, backpropagation is not only a matter of propagating cotangents through the deterministic input-output behaviour of a program. The weights of probabilistic traces, and more generally the effectful structure of computations, become part of the reverse derivative communication channel. We formulate this using monadic types and handlers: monadic operations generate probabilistic computations, while handlers consume them, for instance by taking expected values.

There is already foundational work on forward-mode automatic differentiation for probabilistic programming, notably ADEV [5]. The present paper addresses the corresponding reverse-mode problem. The semantic challenges of the two modes are different: forward-mode propagates sensitivities with the computation, whereas reverse mode must accumulate cotangent contributions backwards through both ordinary program structure and probabilistic weights. Our contribution is to give a compositional reverse-mode transformation for this finite probabilistic setting, formulated using monadic types and handlers, and to prove its correctness by a categorical logical-relations argument. The CHAD framework is essential here: it derives the transformation from the semantic structure rather than retrofitting a correctness proof after the fact, while keeping the construction close to implementation.

1.1 Motivating example: expected-profit optimisation under finite probabilistic choice↩︎

To make the problem concrete, we begin with a small finite probabilistic optimisation example. The model is deliberately simple, but it contains the features that drive the paper: parameter-dependent probabilistic weights, deterministic arithmetic, and a handler that consumes a probabilistic computation by taking an expected value. Reverse-mode AD for such a program must propagate cotangents through all of these components, not only through the deterministic arithmetic. We first describe the model mathematically, and later return to its encoding in the source language and to the corresponding code transformation (Section 5).

We consider the problem of introducing a product to a population. To obtain a finite probabilistic model, we make the simplifying assumption that the population is highly trend-following. There are then two possible outcomes: either the trend does not catch on and no one buys the product, or it catches on and everyone who can afford the product buys one.

The parameters of the model are the per-product profit \(p\) and the marketing budget \(m\). Marketing influences the probability of the successful-trend scenario, while the per-product profit determines how many people can purchase the product if that scenario occurs.

In symbols, the probability of zero total sales is \[\fn[n]{\opn{sig}}{-s_m\cdot(m-m_0)},\] where \(\fn[n]{\opn{sig}}{x}=(1+e^{-x})^{-1}\) is the sigmoid function and \(s_m,m_0\in\R\) are constants: \(s_m\) represents the model’s sensitivity to marketing, while \(m_0\) is the marketing budget at which the success and failure scenarios are equally likely. If, instead, the campaign is successful, the number of purchases is \[P\cdot \fn[n]{\opn{sig}}{-s_p\cdot(p-p_0)},\] where \(P\) is the total population, \(s_p\) is a known constant controlling the sensitivity to changes in \(p\), and \(p_0\) is the per-product profit at which the product is affordable to half the population. Thus, by reducing the per-product profit, a larger portion of the population is able to afford the product.

The expected total profit, which we intend to maximise, is therefore \[\fn[n]{\opn{ETP}}{p,m} = p \cdot \underbrace{ P \cdot \fn[n]{\opn{sig}}{-s_p\cdot(p-p_0)} \cdot \fn[n]{\opn{sig}}{s_m\cdot(m-m_0)} }_{\text{expected number of sales}} - m .\] The stationarity conditions obtained by differentiating this expected objective are a pair of coupled transcendental equations. Thus, even in this simple example, optimisation is naturally performed numerically, for instance by gradient descent.

1.2 Contributions and Outline↩︎

The example above illustrates the central semantic issue addressed in this paper: probabilistic programs carry differentiable information not only through their deterministic computations, but also through the effectful structure by which probabilistic computations are generated and handled. The body of the paper addresses this issue as follows.

First, we give a factorisation-based account of the finite atomic distribution (FAD) monad. In \(\Set\), a FAD is a finitely supported nonnegative weight function; we reconstruct it by factorising the weighted-sum strong monad morphism induced by the barycentre algebra on lists of point/log-weight pairs (Section 2.2). This is the bridge to the differentiated semantics: in \(\mathsf{Fam}(\Vect[o])\), a distribution \(w\) on a space \(X\) with cotangent spaces \(A(x)\) has fibre \[\bigoplus_{x\in\operatorname{supp}(w)}(A(x)\oplus\R),\] whose summands record cotangents for the atom \(x\) and for its log-weight. Thus the lifted FAD monad captures movement of atoms and variation of weights (Sections 2.3 and 2.4).

Second, we use this lift to extend reverse-mode CHAD to a first-order probabilistic language with products, coproducts, Euclidean base types \(\R[n]\), differentiable primitives, the monadic type constructor \(\opn[b]{M}\), categorical distributions, and expectation-like handler algebras. Its ordinary semantics lives in \(\Set\), its differentiated semantics in \(\mathsf{Fam}(\Vect[o])\), and the transformation is the induced structure-preserving translation into a target language for derivatives.

Third, we prove correctness by logical relations. We construct a category \(\cat[]{LR}\) relating the ordinary \(\Set\)-semantics of source programs to the \(\mathsf{Fam}(\Vect[o])\)-semantics of their differentiated programs. Its projections recover, respectively, the ordinary source denotation and the CHAD denotation. The resulting correctness theorem (Theorem 2) says that, for real-input/real-output terms, the transformed program computes the original value and its reverse derivative.

Fourth, we encode the motivating example in the source language using categorical distributions, monadic composition, and an expectation handler. The transformed program exhibits the intended cotangent flow through deterministic arithmetic, atom locations, and log-weights.

Finally, although the paper focuses on finite discrete probability, the construction also identifies a reusable pattern for discrete-output monadic effects. Section 6 explains this general pattern, and Appendix 15 spells it out for finite multiset non-determinism, exceptions, and \(\R[n]\)-writer accumulation. As such, this paper makes important steps towards both richer probabilistic-differentiable languages and towards general differentiable programming with algebraic effects and handlers.

2 The Finite Atomic Distribution Monad↩︎

The probabilistic effect studied in the main body of the paper is represented by the finite atomic distribution (FAD) monad. For a set \(X\), its elements are finitely supported, nonnegative weight distributions on \(X\). Operationally, this is the monad of finite probabilistic computations: monadic composition sums the weights of finite execution traces, while handlers, such as expectation, consume the resulting weighted distribution. For reverse-mode CHAD, however, the ordinary set-level monad is not enough. We need a lifted version that also records how cotangents propagate through the probabilistic structure itself, including both the locations of atoms and their weights.

This section develops the FAD monad in the form needed later in the paper. We begin in \(\Set\), recalling its familiar description as the monad of finitely supported nonnegative weight functions. We then give an equivalent factorisation presentation in \(\Set\): lists of point/log-weight pairs serve as raw weighted syntax, their weighted-sum semantics gives a strong monad morphism, and factorising this morphism identifies precisely the presentations that determine the same finite atomic distribution. The following subsection packages these data as a FAD factorisation structure. This structure is then instantiated in \(\mathsf{Fam}(\Vect[o])\), where the associated FAD monad carries cotangent information for atoms and log-weights, and in \(\cat[]{LR}\), where the same construction is used in the logical-relations correctness argument.

2.1 FAD Monad on \(\Set\)↩︎

We first recall the direct set-level description of the FAD monad. This is the familiar finite-support distribution monad on \(\Set\), used here in its unnormalised weighted form; it appears both in the probabilistic-programming literature [6] and in implementations [7].

We use standard categorical notation: products and coproducts are written \(1\), \(X_1\times X_2\), \(\prod_{i\in I}X_i\), and \(0\), \(X_1\amalg X_2\), \(\coprod_{i\in I}X_i\), with projections \(\pi_i\), inclusions \(\iota_i\), tuplings \(\tupling{f_i}_{i\in I}\), and cotuplings \(\cotupling{f_i}_{i\in I}\). Exponentials are denoted by \(\inthom{X}{Y}\) or by \(Y^{X}\).

A finite atomic distribution (FAD) on a set \(X\) is a finitely supported nonnegative function \(w:X\to\R\). Explicitly, we define a FAD as the following notion of unnormalised weight distribution: \[\fn[n]{w}{x}\geq 0 \;\text{for all }x\in X, \qquad \abs{\operatorname{supp}w}<\infty, \qquad \operatorname{supp}w:=\{x\in X\mid \fn[n]{w}{x}\neq 0\}.\] The set \(\operatorname{supp}w\) is the support of \(w\); its elements are the atoms of the distribution, and \(\fn[n]{w}{x}\) is the weight of the atom \(x\). We write \(\Fun[nn]{M}{X}\) for the set of finite atomic distributions on \(X\).

The functorial action pushes weights forward. For \(f:X\to Y\), define \(\Fun[nn]{M}{f}:\Fun[nn]{M}{X}\to\Fun[nn]{M}{Y}\) by \[\fn[n]{\Fun[nn]{M}{f}}{w}(y) := \sum_{\substack{x\in\operatorname{supp}w\\ \fn[n]{f}{x}=y}} \fn[n]{w}{x}.\] The unit \(\Nat[nn]{\eta}{X}:X\to\Fun[nn]{M}{X}\) sends each point to its Dirac distribution, that is, \(\fn[n]{\Nat[nn]{\eta}{X}}{x}(x')=1\) if \(x=x'\), and \(\fn[n]{\Nat[nn]{\eta}{X}}{x}(x')=0\) otherwise. The multiplication \(\Nat[nn]{\mu}{X}:\Fun[nn]{M^2}{X}\to\Fun[nn]{M}{X}\), where \(M^2:=M\circ M\), flattens a FAD of FADs by summing all contributions to each point: \[\fn[n]{\Nat[nn]{\mu}{X}}{\omega}(x) := \sum_{\substack{w\in\operatorname{supp}\omega\\ x\in\operatorname{supp}w}} \fn[n]{\omega}{w}\cdot\fn[n]{w}{x}.\] The maps \(\eta:\mathrm{Id}_{\Set}\to M\) and \(\mu:M^2\to M\) are natural and equip \(M\) with a monad structure.

This completes the direct description of the FAD monad on \(\Set\). We next give a second presentation of the same monad, not by changing the object being defined, but by changing how it is obtained: finite atomic distributions are recovered as the image of a weighted-sum map from ordered point/log-weight presentations. This factorisation presentation is the one that admits the axiomatic formulation used in the rest of the section.

2.2 Recovering FAD by Weighted-Sum Factorisation↩︎

The preceding subsection described \(\Fun[nn]{M}{X}\) directly as the set of finitely supported nonnegative weight functions on \(X\). We now recover the same object as the image of a weighted-sum transformation. The domain of this transformation is the ordered-list representation of weighted finite support: a list of point/log-weight pairs over \(X\) acts on a test function \(f:X\to\R\) by evaluating \(f\) at the listed points and summing with weights \(\exp\gamma_j\). Its image factorisation has \(\Fun[nn]{M}{X}\) as the middle object: the left factor forms the corresponding finite atomic distribution by adding the weights of coincident points, while the right factor embeds such a distribution as its integration functional. This is the calculation in \(\Set\) whose categorical content is isolated in the next subsection.

We begin with the real-valued case, which provides the barycentre algebra used below. For a list of point/log-weight pairs \(\tuple{ \tuple{t_{1}, \gamma_{1}}, \cdots, \tuple{t_{n}, \gamma_{n}} } \in \coprod_{n \in \mathbb{N}} \left( \R \times \R \right)^{n}\), where \(t_{j},\gamma_{j}\in\R\), define its unnormalised barycentre by¹ \[\label{eq:barycentre-in-set} \fn[n]{ \operatorname{bary}}{ \tuple{ \tuple{t_{j}, \gamma_{j}} }_{j = 1}^{n} } := \sum_{j = 1}^{n} t_{j} \cdot \exp \gamma_{j}.\tag{1}\] Here \(\gamma_{j}\) is a log-weight, and \(\exp\gamma_{j}\) is the corresponding positive weight assigned to \(t_j\). If the list contains collisions, meaning that \(t_{i}=t_{j}=:t\) for distinct \(i,j\in\{1,\cdots,n\}\), then the total weight assigned to \(t\) is the sum of \(\exp\gamma_{k}\) over all \(k\) such that \(t_k=t\).

For an arbitrary set \(X\), we write \[\Fun[nn]{ \mathsf{AST}}{ X } := \coprod_{n \in \mathbb{N}} \left( X \times \R \right)^{n}\] for the set of lists of point/log-weight pairs over \(X\). Given a function \(f:X\to\R\), define \[\label{eq:definition-ast-in-set} \begin{align} \Fun[nn]{ \mathsf{AST}}{ f }: \Fun[nn]{ \mathsf{AST}}{ X } &\to \Fun[nn]{ \mathsf{AST}}{ \R };\qquad \tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} \mapsto \tuple{\tuple{\fn[n]{ f }{ x_{j} }, \gamma_{j}}}_{j = 1}^{n}. \end{align}\tag{2}\] Here \(\mathsf{AST}\) stands for abstract syntax tree, for reasons explained in Section 2.3. Thus \(\Fun[nn]{ \mathsf{AST}}{ f }\) acts on a raw weighted presentation by applying \(f\) to its points and leaving the log-weights unchanged. Composing with the barycentre gives the weighted-sum map \(\Nat[nn]{ \operatorname{ws}}{ X }\): \[\label{eq:weighted-sum-definition-in-set} \begin{align} \Nat[nn]{ \operatorname{ws}}{ X }: \Fun[nn]{ \mathsf{AST}}{ X } &\to \inthom{\inthom{X}{\R}}{\R} \\ \tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} &\mapsto \left( f \mapsto \fn[p]{ \operatorname{bary}\circ \Fun[nn]{ \mathsf{AST}}{ f } }{ \tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} } = f \mapsto \sum_{j = 1}^{n} \fn[n]{ f }{ x_{j} } \cdot \exp \gamma_{j} \right). \end{align}\tag{3}\]

We factorise \(\Nat[nn]{ \operatorname{ws}}{ X }\) through \(\Fun[nn]{ M }{ X }\). Define maps \(\Nat[nn]{ \operatorname{categorical}}{ X }\) and \(\Nat[nn]{ \int}{ X }\) such that \[\Nat[nn]{ \operatorname{ws}}{ X } = \Nat[nn]{ \int}{ X } \circ \Nat[nn]{ \operatorname{categorical}}{ X }: \Fun[nn]{ \mathsf{AST}}{ X } \to \Fun[nn]{ M }{ X } \to \inthom{\inthom{X}{\R}}{\R}.\] The left factor \(\Nat[nn]{ \operatorname{categorical}}{ X }\) sends a raw presentation to the corresponding finite atomic distribution, accumulating the weights of all occurrences of each point. For any \(\tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} \in \Fun[nn]{ \mathsf{AST}}{ X }\), define \[\label{eq:categorical-definition-in-set} \fn[n]{ \Nat[nn]{ \operatorname{categorical}}{ X } }{ \tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} }: X \to \R \qq{s.t.} x \mapsto \sum_{\substack{1 \le j \le n \\ x_{j} = x}} \exp \gamma_{j}.\tag{4}\] Its support is exactly \(\{x_{1},\cdots,x_{n}\}\), since the displayed sum is zero precisely when it is empty. The name \(\operatorname{categorical}\) reflects the standard probability theory nomenclature of finitely supported weight distributions. \(\Nat[nn]{ \operatorname{categorical}}{ X }\) is surjective, but not injective; it forgets about ordering and collisions.

The right factor \(\Nat[nn]{ \int}{ X }\) sends a finite atomic distribution to the corresponding integration functional on test functions. For \(w \in \Fun[nn]{ M }{ X }\), define \[\label{eq:iwrt-definition-in-set} \fn[n]{ \Nat[nn]{ \int}{ X } }{ w }: \inthom{X}{\R} \to \R \qq{s.t.} f \mapsto \sum_{x \in \operatorname{supp}w} \fn[n]{ f }{ x } \cdot \fn[n]{ w }{ x }.\tag{5}\] This map is injective because indicator functions separate finite atomic distributions: for each \(x\in X\), \(\fn[p]{ \fn[n]{ \Nat[nn]{ \int}{ X } }{ w } }{ \mathbb{1}_{\left\{ x \right\}} } = \fn[n]{w}{x}.\) Hence, for \(w,w'\in\Fun[nn]{M}{X}\), \(\fn[n]{ \Nat[nn]{ \int}{ X } }{ w } = \fn[n]{ \Nat[nn]{ \int}{ X } }{ w' }\) if and only if \(w=w'\).

The composite has the required weighted-sum semantics. Indeed, for any list \(\tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} \in \Fun[nn]{ \mathsf{AST}}{ X }\) and any function \(f:X\to\R\), \[\begin{align} \fn[p]{ \fn[p]{ \Nat[nn]{ \int}{ X } \circ \Nat[nn]{ \operatorname{categorical}}{ X } }{ \tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} } }{ f } &= \sum_{x \in \left\{ x_{1}, \cdots, x_{n} \right\} } \left( \fn[n]{ f }{ x } \cdot \sum_{\substack{1 \le j \le n \\ x_{j} = x}} \exp \gamma_{j} \right) \\ &= \sum_{j = 1}^{n} \fn[n]{ f }{ x_{j} } \cdot \exp \gamma_{j}= \fn[p]{ \fn[n]{ \Nat[nn]{ \operatorname{ws}}{ X } }{ \tuple{\tuple{x_{j}, \gamma_{j}}}_{j = 1}^{n} } }{ f }. \end{align}\] Thus \(\Nat[nn]{ \operatorname{ws}}{ X } = \Nat[nn]{ \int}{ X } \circ \Nat[nn]{ \operatorname{categorical}}{ X }\). Since \(\Nat[nn]{ \operatorname{categorical}}{ X }\) is surjective and \(\Nat[nn]{ \int}{ X }\) is injective, this is the image factorisation of \(\Nat[nn]{ \operatorname{ws}}{ X }\) in \(\Set\).

The calculation has isolated the three ingredients that we will abstract next. The ordered lists form the raw weighted-syntax monad \(\mathsf{AST}\); the barycentre algebra induces the weighted-sum morphism into a continuation monad; and the FAD monad on \(\Set\) is recovered as the middle object of the resulting image factorisation. We next package these ingredients as a FAD factorisation structure on a category \(\cat[]{X}\) and define the associated FAD monad on \(\cat[]{X}\).

2.3 FAD Monads by Axiomatic Factorisation↩︎

We now abstract the factorisation just described in \(\Set\). The aim is to specify the structure on a category \(\cat[]{X}\) that is sufficient to form weighted raw syntax, interpret it by a barycentre algebra, and factorise the resulting strong monad morphism. Once these data are fixed, the FAD monad on \(\cat[]{X}\) is not additional primitive structure: it is the monad obtained as the middle term of the factorisation \[\mathsf{AST}^{\cat[]{X}} \xrightarrow{\;\operatorname{categorical}^{\cat[]{X}}\;} M^{\cat[]{X}} \xrightarrow{\;\int^{\cat[]{X}}\;} \mathsf{K}{V}^{\cat[]{X}} .\] For \(\cat[]{X}=\Set\), with the choices specified below, this construction recovers the FAD monad on \(\Set\).

We assume that \(\cat[]{X}\) is bicartesian closed with countable coproducts. Thus we fix a terminal object \(1\), an initial object \(0\), binary products and coproducts, countable coproducts, and, for each object \(A\), a right adjoint \(\inthom{A}{\mathord{-}}:\cat[]{X}\to\cat[]{X}\) to \((\mathord{-})\times A:\cat[]{X}\to\cat[]{X}\).

We also fix a monoid object \(R\) in \(\cat[]{X}\), with unit \(1_R:1\to R\) and multiplication \(\cdot_R:R\times R\to R\). This is the monoid of weights; in the motivating instance, it is the additive monoid of log-weights. Finally, we fix an \(R\)-semimodule \(V\), called the barycentre semimodule, with structure maps \[+_V:V\times V\to V, \qquad 0_V:1\to V, \qquad \cdot_R^V:V\times R\to V .\] Here an \(R\)-semimodule means a commutative monoid object \(V\) equipped with a right action \(V\times R\to V\) of the monoid object \(R\) by commutative-monoid endomorphisms.

In \(\Set\), the relevant choices are \(R=\R\) (thought of as log-weights), with \[\label{eq:monoid-of-weights-operations-in-set} \begin{align} 1_R: 1 &\to \R \\ \ast &\mapsto 0 \end{align} \qq{and} \begin{align} \cdot_R: \R \times \R &\to \R \\ \tuple{a, b} &\mapsto a + b, \end{align}\tag{6}\] and \(V=\R\) (thought of as real values), with \[\label{eq:barycentre-semimodule-operations-in-set} \begin{align} 0_V: 1 &\to \R \\ \ast &\mapsto 0 \end{align} \qq{and} \begin{aligned} +_V: \R \times \R &\to \R \\ \tuple{a, b} &\mapsto a + b \end{aligned} \qq{and} \begin{align} \cdot_R^V: \R \times \R &\to \R \\ \tuple{v, r} &\mapsto v \cdot \exp r . \end{align}\tag{7}\] The case \(V=\R[p]\), for \(p\in\mathbb{N}\), is obtained in the same way.

From \(R\) and \(V\) we form two monads. The raw weighted-syntax monad and the continuation monad on \(V\) have underlying endofunctors \[\label{eq:definition-ast-continuation-general} \mathsf{AST}^{\cat[]{X}} := \coprod_{n \in \mathbb{N}} \left( \mathord{-}\times R \right)^{n} \qq{and} \mathsf{K}{V}^{\cat[]{X}} := \inthom{\inthom{\mathord{-}}{V}}{V}.\tag{8}\] The monad structure on \(\mathsf{AST}^{\cat[]{X}}\) is obtained by combining the ordered-list monad \(A\mapsto\coprod_{n\in\mathbb{N}}A^n\) [8] with the \(R\)-writer monad \(A\mapsto A\times R\), induced by the monoid \(\tuple{R,\cdot_R,1_R}\); see, for instance, Chapter 7 of [9]. These monads are combined using the canonical distributive law recalled in Section 6.3 of [10]. Thus \(\mathsf{AST}^{\cat[]{X}}\) is an ordered-list presentation of weighted finite support. The identifications imposed by permutations and collisions are not built into \(\mathsf{AST}^{\cat[]{X}}\); they are produced by factorisation.

Both \(\mathsf{AST}^{\cat[]{X}}\) and \(\mathsf{K}{V}^{\cat[]{X}}\) are strong monads. For \(\mathsf{AST}^{\cat[]{X}}\), the strength is induced by distributivity of products over countable coproducts and by the canonical maps \(A \times B^{n} \to A^{n} \times B^{n} \xrightarrow{\sim} \left( A \times B \right)^{n}\) obtained from the diagonal morphism and tuplings of projections. Then, by using \(B\times R\) in place of \(B\), this gives a natural transformation \[A \times \Fun[nn]{ \mathsf{AST}^{\cat[]{X}} }{ B } \to \Fun[np]{ \mathsf{AST}^{\cat[]{X}} }{ A \times B }.\] The continuation monad carries its standard strength in every cartesian closed category.

The semimodule structure on \(V\) defines an \(\mathsf{AST}^{\cat[]{X}}\)-algebra structure on \(V\), the barycentre algebra. Explicitly, it is the morphism \[\label{eq:definition-barycentre-algebra-general} \operatorname{bary}^{\cat[]{X}}: \Fun[nn]{ \mathsf{AST}^{\cat[]{X}} }{ V } \to V; \quad \textrm{defined by} \quad \operatorname{bary}^{\cat[]{X}} = \cotupling{\underset{n}{\Sigma} \circ \left( \cdot_R^V\right)^{n}}_{n \in \mathbb{N}},\tag{9}\] where \(\underset{n}{\Sigma}:V^n\to V\) is the \(n\)-ary sum determined by \(+_V\) and \(0_V\), and \(\left(\cdot_R^V\right)^n:(V\times R)^n\to V^n\) is the product of \(n\) copies of \(\cdot_R^V\).

Under continuation-passing style, the barycentre algebra corresponds to a strong monad morphism \(\operatorname{ws}^{\cat[]{X}}\) to the continuation monad \(\inthom{\inthom{-}{V}}{V}\) whose \(A\)-component sends \(a \in \mathsf{AST}^{\cat[]{X}}A\) to \(\lambda k.\,\operatorname{bary}^{\cat[]{X}}\!\left(\mathsf{AST}^{\cat[]{X}}k(a)\right)\).

Applying Lemma 1 to the barycentre algebra gives a strong monad morphism \[\operatorname{ws}^{\cat[]{X}}: \mathsf{AST}^{\cat[]{X}} \to \mathsf{K}{V}^{\cat[]{X}} \qquad\textrm{with components}\qquad \Nat[nn]{ \operatorname{ws}^{\cat[]{X}} }{ A }: \Fun[nn]{ \mathsf{AST}^{\cat[]{X}} }{ A } \to \inthom{\inthom{A}{V}}{V}.\]

It remains to specify the factorisation. In \(\Set\), we used the image factorisation. In \(\cat[]{X}\), we fix an orthogonal factorisation system \(\tuple{\mathcal{E},\mathcal{M}}\). Thus every \(f:A\to B\) factors as \(f=m\circ e\), with \(e\in\mathcal{E}\) and \(m\in\mathcal{M}\), and \(\mathcal{E}\) is orthogonal to \(\mathcal{M}\): for every equation \(v\circ e=m\circ u\), with \(e\in\mathcal{E}\) and \(m\in\mathcal{M}\), there is a unique \(w\) such that \(u=w\circ e\) and \(v=m\circ w\). In \(\Set\), taking \(\mathcal{E}\) to be the epis (i.e., surjections) and \(\mathcal{M}\) to be the (regular) monos (i.e., injections) recovers the usual image factorisation. We assume that \(\mathcal{E}\) is closed under binary products. Since left classes of orthogonal factorisation systems are closed under coproducts, this implies closure of \(\mathcal{E}\) under \(\mathsf{AST}^{\cat[]{X}}\).

For each object \(A\), factorise the component \(\Nat[nn]{\operatorname{ws}^{\cat[]{X}}}{A}\) as \[\begin{tikzcd}[column sep = 2 cm] \Fun[nn]{ \mathsf{AST}^{\cat[]{X}} }{ A } && \Fun[nn]{ \mathsf{K}{V}^{\cat[]{X}} }{ A } \\ & \Fun[nn]{ M^{\cat[]{X}} }{ A } & \arrow[from = 1-1, to = 1-3, "{\Nat[nn]{ \operatorname{ws}^{\cat[]{X}} }{ A }}"] \arrow[from = 1-1, to = 2-2, swap, "{\Nat[nn]{ \operatorname{categorical}^{\cat[]{X}} }{ A }}"] \arrow[from = 2-2, to = 1-3, swap, "{\Nat[nn]{ \int^{\cat[]{X}} }{ A }}"] \end{tikzcd}\] with \(\Nat[nn]{ \operatorname{categorical}^{\cat[]{X}} }{ A }\in\mathcal{E}\) and \(\Nat[nn]{ \int^{\cat[]{X}} }{ A }\in\mathcal{M}\). At this stage \(\Fun[nn]{ M^{\cat[]{X}} }{ A }\) denotes only the intermediate objects of the factorisation. The next theorem ensures that these objects assemble into a strong monad.

We can now name the structure and the monad it determines.

With \(\cat[]{X}=\Set\), \(R=\R\), \(V=\R\), and the (epi, regular mono) factorisation, the associated FAD monad \(M^{\Set}\) is canonically isomorphic to the FAD monad on \(\Set\). The point of the preceding abstraction is that the same construction can now be instantiated in categories with additional semantic structure. The first such instance is \(\mathsf{Fam}(\Vect[o])\), where the associated FAD monad refines the FAD monad on the underlying set by adding cotangent fibres for atom locations and log-weights.

2.4 FAD Monad on \(\mathsf{Fam}(\Vect[o])\)↩︎

We now exhibit a FAD factorisation structure on \(\cat[]{F}:=\mathsf{Fam}(\Vect[o])\), and write \(M^{\cat[]{F}}\) for its associated FAD monad. This is the semantic category for differentiated values in reverse-mode CHAD: an object carries a set of possible values together with a specified cotangent space over each value, and a morphism carries both the underlying function on values and the corresponding contravariant action on cotangents. We think of this action as a transposed derivative.

The category \(\mathsf{Fam}(\Vect[o])\) is the free coproduct completion of \(\Vect[o]\) [12]. We use it as the category of sets equipped with cotangent fibres. An object is a pair \(\tuple{X,A}\), where \(X\) is a set and \(A:X\to\Vect[o]\) is a functor, viewing \(X\) as a discrete category. Equivalently, \(A\) is a family \(\{\Fun[np]{A}{x}\}_{x\in X}\) of vector spaces indexed by \(X\).

A morphism \(\tuple{X,A}\to\tuple{Y,B}\) is a pair \(\tuple{f,\alpha}\), where \(f:X\to Y\) is a function and \(\alpha:A\to B\circ f\) is a natural transformation in \(\Vect[o]\). Equivalently, for each \(x\in X\), it determines a linear map \[\Nat[nn]{\alpha}{x}: \Fun[np]{B}{\fn[n]{f}{x}} \to \Fun[np]{A}{x}.\] Thus \(f\) acts on base points, while \(\Nat[nn]{\alpha}{x}\) acts contravariantly on cotangents. In the intended interpretation, \(\Fun[np]{A}{x}\) is the cotangent space at \(x\), and \(\Nat[nn]{\alpha}{x}\) plays the role of the transposed derivative of \(f\) at \(x\). We write \(\left(\tuple{X,A}\right)_{\mathrm{set}}=X, \quad \left(\tuple{X,A}\right)_{\mathrm{fbr}}=A, \quad \left(\tuple{f,\alpha}\right)_{\mathrm{set}}=f, \quad \left(\tuple{f,\alpha}\right)_{\mathrm{fbr}}=\alpha\) for the base and fibre parts. We use this notation throughout the construction below.

The category \(\mathsf{Fam}(\Vect[o])\) is complete and cocomplete. In particular, for any family \(\{\tuple{X_i,A_i}\}_{i\in I}\) of objects, products and coproducts are given by \[\label{eq:definition-product-coproduct-fam-vectop} \prod_{i \in I} \tuple{ X_{i}, A_{i} } = \tuple{ \left( \prod_{i \in I} X_{i} \right), \left( \left\{ x_{i} \right\}_{i \in I} \mapsto \bigoplus_{i \in I} \Fun[nn]{ A_{i} }{ x_{i} } \right) } \qq{and} \coprod_{i \in I} \tuple{ X_{i}, A_{i} } = \tuple{ \left( \coprod_{i \in I} X_{i} \right), \cotupling{A_{i}}_{i \in I} }.\tag{10}\] Moreover, \(\mathsf{Fam}(\Vect[o])\) has exponentials [2]: \([\tuple{X,A},\tuple{Y,B}] = \tuple{\prod_{x\in X}\sum_{y \in Y}\Vect(\Fun[nn]{B}{y}, \Fun[nn]{A}{x}),f\mapsto \bigoplus_{x\in X}\Fun[np]{B}{ \fn[n]{\pi_1\!}{\fn[n]{f\!}{x}}}}\) . Hence \(\cat[]{F}=\mathsf{Fam}(\Vect[o])\) is bicartesian closed with countable coproducts, as required by the axiomatic construction.

We use the (epi, regular mono) factorisation system on \(\mathsf{Fam}(\Vect[o])\). It admits the following explicit description.

Therefore \(\cat[]{F}\) has the categorical structure and factorisation system required for a FAD factorisation structure. It remains to choose the monoid of weights and the barycentre semimodule. Both are obtained by lifting the corresponding \(\Set\)-structures to the real cotangent bundle \[\TR := \tuple{\R,\mathrm{const}\R},\qquad\text{where}\qquad \Fun[pn]{\mathrm{const}Y}{\;x}\;=\;Y.\]

The monoid of weights is \(\TR\). Its unit and multiplication are \[\label{eq:monoid-of-weights-operations-in-fam-vectop} \begin{array}{lrcl} 1_R&: &\tuple{ 1, \mathrm{const}0 } &\to \TR\\ \left(1_R\right)_{\mathrm{set}} &:& 1 &\to \R \\ &&\ast &\mapsto 0\\ \Nat[nn]{ \left(1_R\right)_{\mathrm{fbr}} }{\ast} &: &\R &\to 0 \\ && v &\mapsto 0 \end{array} \qquad \begin{array}{lccl} \cdot_R&: &\TR \times \TR &\to \TR\\ \left(\cdot_R\right)_{\mathrm{set}} &:& \R \times \R &\to \R \\ &&\tuple{ a, b } &\mapsto a + b\\ \Nat[nn]{ \left(\cdot_R\right)_{\mathrm{fbr}} }{ \tuple{ a, b } } &: &\R &\to \R \oplus \R \\ &&v &\mapsto \tuple{ v, v } . \end{array}\tag{11}\] The base parts agree with Equation 6 ; the fibre parts are the corresponding transposed derivatives.

We also take the barycentre semimodule to be \(\TR\).² Its zero element and addition are \[\label{eq:barycentre-semimodule-operations-in-fam-vectop} \begin{array}{lrcl} 0_V&: &\tuple{ 1, \mathrm{const}0 } &\to \TR\\ \left(0_V\right)_{\mathrm{set}} &:& 1 &\to \R \\ &&\ast &\mapsto 0\\ \Nat[nn]{ \left(0_V\right)_{\mathrm{fbr}} }{\ast} &: &\R &\to 0 \\ &&v &\mapsto 0 \end{array} \qquad \begin{array}{lrcl} +_V&: &\TR \times \TR &\to \TR\\ \left(+_V\right)_{\mathrm{set}} &:& \R \times \R &\to \R \\ &&\tuple{ a, b } &\mapsto a + b\\ \Nat[nn]{ \left(+_V\right)_{\mathrm{fbr}} }{ \tuple{ a, b } } &: &\R &\to \R \oplus \R \\ &&v &\mapsto \tuple{ v, v } . \end{array}\tag{12}\] Again, the base parts are those of Equation 7 , and the fibre parts are their transposed derivatives. The action of the monoid of weights on the barycentre semimodule is \[\label{eq:barycentre-semimodule-monoid-action-in-fam-vectop} \begin{array}{lrcl} \cdot_R^V&: &\TR \times \TR &\to \TR\\ \left(\cdot_R^V\right)_{\mathrm{set}} &:& \R \times \R &\to \R \\ &&\tuple{ v, r } &\mapsto v \cdot \exp r\\ \Nat[nn]{ \left(\cdot_R^V\right)_{\mathrm{fbr}} }{ \tuple{ v, r } } &: &\R &\to \R \oplus \R \\ &&t &\mapsto \tuple{ t \cdot \exp r, t \cdot v \cdot \exp r } . \end{array}\tag{13}\] The base part is the action used in \(\Set\), while the fibre part is its transposed derivative.

These choices make \(\cat[]{F}=\mathsf{Fam}(\Vect[o])\) into a FAD factorisation structure. Let \(M^{\cat[]{F}}\) denote the associated FAD monad. For an object \(\tuple{X,A}\), the factorisation construction gives \[\label{eq:fad-monad-on-objects-in-fam-vectop} \Fun[nn]{ M^{\cat[]{F}} }{ \tuple{ X, A } } := \tuple{ \Fun[nn]{ M }{ X }, \left( w \mapsto \bigoplus_{x \in \operatorname{supp}w} \left( \Fun[np]{ A }{ x } \oplus \R \right) \right) }.\tag{14}\] Thus the base of \(\Fun[nn]{ M^{\cat[]{F}} }{ \tuple{X,A} }\) is the set of finite atomic distributions on \(X\). The fibre over \(w\in\Fun[nn]{M}{X}\) has one summand for each atom \(x\in\operatorname{supp}w\): the component \(\Fun[np]{A}{x}\) carries cotangent information for the position of the atom, while the real component carries cotangent information for \(\log \fn[n]{w}{x}\).

For a morphism \(\tuple{f,\alpha}:\tuple{X,A}\to\tuple{Y,B}\), the base part is \[\left( \Fun[nn]{ M^{\cat[]{F}} }{ \tuple{ f, \alpha } } \right)_{\mathrm{set}} = \Fun[nn]{ M }{ f } .\] For \(w\in\Fun[nn]{M}{X}\), the fibre part is \[\label{eq:fad-monad-on-morphisms-in-fam-vectop} \begin{align} \Nat[nn]{ \left(\Fun[nn]{ M^{\cat[]{F}} }{ \tuple{ f, \alpha } }\right)_{\mathrm{fbr}} }{ w }: \bigoplus_{y \in \fn[n]{ f }{ \operatorname{supp}w }} \left( \Fun[np]{ B }{ y } \oplus \R \right) &\to \bigoplus_{x \in \operatorname{supp}w} \left( \Fun[np]{ A }{ x } \oplus \R \right) \\ \sum_{y \in \operatorname{supp}\fn[n]{ \Fun{ M }{ f } }{ w }} \left( b_{y} \cdot \mathbf{ e }_{y} + s_{y} \cdot \mathbf{ \varepsilon }_{y} \right) &\mapsto \sum_{x \in \operatorname{supp}w} \left( \frac{ \fn[n]{ w }{ x } \cdot \left( \fn[n]{ \Nat{ \alpha }{ x } }{ b_{\fn[n]{ f }{ x }} } \cdot \mathbf{ e }_{x} + s_{\fn[n]{ f }{ x }} \cdot \mathbf{ \varepsilon }_{x} \right) }{ \fn[p]{ \fn[n]{ \Fun{ M }{ f } }{ w } }{ \fn[n]{ f }{ x } } } \right) . \end{align}\tag{15}\] Here \(\mathbf{e}_{\bullet}\) denotes the atom-location component and \(\mathbf{\varepsilon}_{\bullet}\) the log-weight component in the relevant direct sum. The formula says that cotangents are pulled back along \(\tuple{f,\alpha}\) by distributing the cotangent at an image atom among all atoms mapping to it, with the contribution of \(x\) scaled by the relative weight \[\frac{ \fn[n]{ w }{ x } }{ \fn[p]{ \fn[n]{ \Fun{ M }{ f } }{ w } }{ \fn[n]{ f }{ x } } } .\] The atom component is then transported through \(\Nat[nn]{\alpha}{x}\), while the log-weight component is copied to the corresponding \(\mathbf{\varepsilon}_{x}\)-summand.

Under the base projection, the monoid of weights, the barycentre semimodule, and the monad \(M^{\cat[]{F}}\) recover the corresponding structures in \(\Set\). Thus \(M^{\cat[]{F}}\) is a lift of the FAD monad on \(\Set\) along the base projection, with the additional fibre information required for reverse-mode differentiation. Over a distribution \(w\in\Fun[nn]{M}{X}\), this fibre consists of one cotangent component for each atom \(x\in\operatorname{supp}w\), and one real cotangent component for the log-weight \(\log\fn[n]{w}{x}\).

2.4.1 Expected-value gradients and score contributions↩︎

We now spell out how the lifted FAD monad separates the usual contributions to the derivative of an expected value. The point is not to introduce an external analytic argument, but to show that the decomposition is already encoded by the fibre maps in \(\mathsf{Fam}(\Vect[o])\). In particular, for a function \(\fn[n]{f}{\theta,z}\), with \(z\) sampled from a finite distribution \(\fn[n]{q}{\theta}\) depending on the parameter \(\theta\), the lifted semantics separates three sources of variation: the explicit dependence of \(f\) on \(\theta\), the dependence of the atoms \(z\) on \(\theta\), and the dependence of the log-weights on \(\theta\).

To make this comparison readable, we use differential notation for objects and morphisms in \(\mathsf{Fam}(\Vect[o])\). Thus we write \[T^{\ast}U = \tuple{ U, \left\{T^{\ast}_{u}U\right\}_{u\in U} }\] and write morphisms as \[\hat{h} = \tuple{ h, \grad h }: T^{\ast} X \to T^{\ast} Y, \qquad \left( \grad h \right)_{x} : T^{\ast}_{\fn[n]{h}{x}}Y \to T^{\ast}_{x}X,\] where \(\Nat[pn]{ \grad h }{x}\) denotes the \(x\)-fibre of the fibre part.

Consider the lifted real-valued expectation algebra \[\hat{\opn{E}}_{\R}^{\cat[]{F}}: \Fun[np]{M^{\cat[]{F}}}{\TR} \to \TR .\] For \(u\in\Fun[nn]{M}{\R}\), its base and fibre parts are \[\label{eq:lifted-real-expectation-fibre} \fn[n]{ \opn{E}_{\R}^{\cat[]{F}} }{u} = \sum_{r\in\operatorname{supp}u} \fn[n]{ u }{ r } \cdot r, \qquad \fn[n]{ \Nat[pn]{ \grad \opn{E}_{\R}^{\cat[]{F}} }{ u } }{ t } = \sum_{r \in \operatorname{supp}u} \left( \fn[n]{ u }{ r } \cdot t \cdot\mathbf{e}_{r} + \fn[n]{ u }{ r } \cdot r \cdot t \cdot \mathbf{\varepsilon}_{r} \right).\tag{16}\] The two summands in the fibre component have different origins. The \(\mathbf{e}_{r}\)-term is the cotangent contribution of the atom value \(r\); the \(\mathbf{\varepsilon}_{r}\)-term is the cotangent contribution of the log-weight \(\log \fn[n]{u}{r}\).

Let \[\hat{g}=\tuple{g,\grad g}:T^{\ast}X\to\TR, \qquad w\in\Fun[nn]{M}{X}, \qquad u:=\fn[p]{\Fun[nn]{M}{g}}{w}.\] The Frobenius identity gives \[\sum_{r\in\operatorname{supp}u} \fn[n]{ u }{ r } \cdot r = \sum_{x\in\operatorname{supp}w} \fn[n]{ w }{ x } \cdot \fn[n]{ g }{ x }.\] On fibres, composing the map induced by \(\Fun[nn]{M^{\cat[]{F}}}{\hat{g}}\) with Equation 16 gives \[\label{eq:frobenius-cotangent-rule} \fn[n]{ \Nat[pn]{ \grad \! \left( \hat{ \opn{E} }_{\R}^{\cat[]{F}} \circ \Fun[nn]{ M^{\cat[]{F}} }{ \hat{g} } \right) }{ w } }{t} = \sum_{x \in \operatorname{supp}w} \left( \fn[n]{ w }{ x } \cdot \fn[n]{ \Nat[pn]{ \grad g }{ x } }{ t } \cdot \mathbf{e}_x + \fn[n]{ w }{ x } \cdot \fn[n]{ g }{ x } \cdot t \cdot \mathbf{ \varepsilon }_x \right).\tag{17}\] Indeed, the fibre map of \(\Fun[nn]{M^{\cat[]{F}}}{\hat{g}}\) contributes, in the \(x\)-summand, the relative factor \[\frac{\fn[n]{w}{x}}{\fn[p]{\fn[n]{\Fun{M}{g}}{w}}{\fn[n]{g}{x}}},\] while Equation 16 contributes the corresponding mass \[\fn[p]{\fn[n]{\Fun{M}{g}}{w}}{\fn[n]{g}{x}}.\] These factors cancel, leaving precisely the expression in Equation 17 .

Now let \[\hat{q} = \tuple{ q, \grad q }: T^{\ast} \Theta \to \Fun[np]{ M^{\cat[]{F}} }{ T^{\ast} Z }\] be a parameterised FAD. For each \(\theta\in\Theta\), choose a local collision-free presentation \[\fn[n]{ q }{ \theta } = \fn[n]{ \operatorname{categorical}}{ \tuple{ \tuple{ \fn[n]{ z_{i} }{ \theta }, \fn[n]{ \ell_{i} }{ \theta } } }_{i \in I} } \qq{and} \fn[n]{ w_{i} }{ \theta } := \exp \fn[n]{ \ell_{i} }{ \theta } .\] Let \[\hat{\underline q} := \Nat[nn]{ \operatorname{st}^{\cat[]{F}} }{ T^{\ast} \Theta, T^{\ast} Z } \circ \tupling{ \mathrm{id}_{T^{\ast} \Theta}, \hat{q} } .\] Then \[\fn[n]{ \underline{q} }{ \theta } = \fn[n]{ \operatorname{categorical}}{ \tuple{ \tuple{ \theta, \fn[n]{ z_{i} }{ \theta } }, \fn[n]{ \ell_{i} }{ \theta } }_{i \in I} }.\] With respect to this local labelling, the fibre part of \(\hat{\underline{q}}\) is \[\label{eq:suggestive-fibre-parametrised-fad} \fn[n]{ \Nat[pn]{ \grad \underline{q} }{ \theta } }{ \sum_{i \in I} \left( \tuple{d_i,a_i} \cdot \mathbf{e}_{\tuple{\theta,z_i(\theta)}} + s_i\cdot\mathbf{ \varepsilon }_{\tuple{\theta,z_i(\theta)}} \right) } = \sum_{i\in I} d_i + \sum_{i \in I} \fn[n]{ \Nat[pn]{ \grad z_{i} }{ \theta } }{ a_{i} } + \sum_{i \in I} \fn[n]{ \Nat[pn]{ \grad \ell_{i} }{ \theta } }{ s_{i} },\tag{18}\] where \[d_i\in T^{\ast}_{\theta}\Theta, \qquad a_i\in T^{\ast}_{\fn[n]{z_i}{\theta}}Z, \qquad s_i\in\R .\]

Finally, let \[\widehat{f} = \tuple{f,\grad f}: T^{\ast}\Theta\times T^{\ast}Z \to \TR .\] We write the two components of its fibre part as \[\fn[n]{ \Nat[pn]{ \grad f }{\tuple{ \theta, z }} }{ t } = \tuple{ \fn[n]{ \Nat[nn]{ \left( \grad f \right)^{\Theta} }{ \tuple{ \theta, z } } }{ t }, \fn[n]{ \Nat[nn]{ \left( \grad f \right)^{Z} }{ \tuple{ \theta, z } } }{ t } }.\] Apply Equation 17 to the composite \[T^{\ast} \Theta \xrightarrow{\hat{\underline q}} \Fun[np]{ M^{\cat[]{F}} }{ T^{\ast} \Theta \times T^{\ast} Z } \xrightarrow{\Fun[nn]{ M^{\cat[]{F}} }{ \widehat{f} }} \Fun[nn]{ M^{\cat[]{F}} }{ \TR } \xrightarrow{\opn{E}^{\cat[]{F}}_{\R}} \TR\] and then use Equation 18 . This gives \[\label{eq:three-term-weighted-particle-gradient} \begin{align} \fn[n]{ \Nat[pn]{ \grad\! \left( \hat{ \opn{E} }_{\R}^{\cat[]{F}} \circ \Fun[nn]{ M^{\cat[]{F}} }{ \widehat{f} } \circ \hat{\underline{q}} \right) }{ \theta } }{ t } &= \sum_{i \in I} \fn[n]{ w_{i} }{ \theta } \cdot \left[ \fn[n]{ \Nat[nn]{ \left( \grad f \right)^{\Theta} }{ \tuple{ \theta, \fn[n]{ z_{i} }{ \theta } } } }{ t } + \fn[n]{ \Nat[pn]{ \grad z_{i} }{ \theta } }{ \fn[n]{ \Nat[nn]{ \left( \grad f \right)^{Z} }{ \tuple{ \theta, \fn[n]{ z_{i} }{ \theta } } } }{ t } } \right. \\ &\qquad \qquad \qquad + \left. \fn[n]{ f }{ \tuple{ \theta, \fn[n]{ z_{i} }{ \theta } } } \cdot \fn[n]{ \Nat[pn]{ \grad \ell_{i} }{ \theta } }{ t } \right]. \end{align}\tag{19}\] The three terms inside the brackets have distinct origins. The first is the explicit parameter contribution, coming from the dependence of \(\fn[n]{f}{\theta,z}\) on \(\theta\). The second is the moving-atom contribution, obtained by transporting cotangents through the dependence of \(z_i(\theta)\) on \(\theta\). The third is the score contribution, coming from the dependence of the log-weight \(\ell_i(\theta)=\log \fn[n]{w_i}{\theta}\) on \(\theta\).

This calculation explains the operational role of the lifted FAD monad in \(\mathsf{Fam}(\Vect[o])\). It separates the reverse-mode contributions carried by a finite probabilistic computation into atom-location terms and log-weight terms, and it does so internally, through the fibre maps of the associated FAD monad. The next subsection lifts the construction once more, from \(\mathsf{Fam}(\Vect[o])\) to the logical-relations category used to prove that the syntactic CHAD transformation computes these cotangent contributions correctly.

2.5 FAD Monad on the Logical-Relations Category↩︎

The final instance of the FAD factorisation construction is the logical-relations category used in the correctness proof. Whereas \(\mathsf{Fam}(\Vect[o])\) records the differentiated semantics of programs, \(\cat[]{LR}\) relates this differentiated semantics to the ordinary semantics in \(\Set\). The purpose of this subsection is to equip \(\cat[]{LR}\) with a FAD factorisation structure, and hence with an associated FAD monad \(M^{\cat[]{LR}}\). This is the semantic ingredient needed later to interpret the source language in the logical-relations model and to prove correctness of the reverse-mode CHAD transformation.

We use the comma-category presentation of logical relations from [2]. Let \(G:\cat[]{C}\to\cat[]{D}\) be a functor. We write \[\cat[]{LR}:=\commacat{\cat[]{D}}{G}.\] An object of \(\cat[]{LR}\) is a triple \(\tuple{D,C,f}\), where \(D\) is an object of \(\cat[]{D}\), \(C\) is an object of \(\cat[]{C}\), and \(f:D\to\Fun[nn]{G}{C}\) is a morphism in \(\cat[]{D}\). A morphism \[\tuple{u,v}:\tuple{D,C,f}\to\tuple{D',C',f'}\] consists of morphisms \(u:D\to D'\) in \(\cat[]{D}\) and \(v:C\to C'\) in \(\cat[]{C}\) such that \(\Fun[nn]{G}{v}\circ f = f'\circ u.\) There is therefore a forgetful functor \[U:\cat[]{LR}\to\cat[]{D}\times\cat[]{C}, \qquad \textrm{defined by} \qquad \Fun[nn]{U}{\tuple{D,C,f}}:=\tuple{D,C}, \qquad \Fun[nn]{U}{\tuple{u,v}}:=\tuple{u,v}.\]

For the correctness argument in this paper, we take \[\cat[]{D}:=\Set, \qquad \cat[]{C}:=\Set\times\mathsf{Fam}(\Vect[o]),\] and define \(G\) to be the hom-functor \[\label{eq:definition-g-functor-logical-relations} G := \fn[n]{ \cat[]{C} }{ \tuple{ \mathbb{E}, T^{\ast}\mathbb{E}}, \mathord{-}},\tag{20}\] where \(\mathbb{E}\) is a fixed finite-dimensional real vector space, regarded as a set, and \(T^{\ast}\mathbb{E}:=\tuple{\mathbb{E},\mathrm{const}\mathbb{E}^{\ast}}\) is its cotangent bundle.

The space \(\mathbb{E}\) is used as a probe object. For example, when \(\mathbb{E}=\R\), morphisms from \(\tuple{\mathbb{E},T^{\ast}\mathbb{E}}\) into an object of \(\cat[]{C}\) can be read as parametrised curves equipped with cotangent data. In this way, differentiability of a map can be tested by its action on such probes: postcomposition should send admissible probes in the domain to admissible probes in the codomain.

Thus, in the present instance, an object of \(\cat[]{LR}\) is a triple \(\tuple{D,C,p}\) where \(p:D\to\fn[n]{\cat[]{C}}{\tuple{\mathbb{E},T^{\ast}\mathbb{E}},C}\) is a function. When \(p\) is injective, \(D\) may be regarded as a specified class of probes \(\tuple{\mathbb{E},T^{\ast}\mathbb{E}}\to C\), namely those satisfying the predicate represented by \(p\). A morphism \(\tuple{u,v}:\tuple{D,C,p}\to\tuple{D',C',p'}\) asserts that postcomposition with \(v:C\to C'\) sends probes satisfying \(p\) to probes satisfying \(p'\).

We write \[\left(\tuple{D,C,p}\right)_{\mathrm{dom}}:=D, \qquad \left(\tuple{D,C,p}\right)_{\mathrm{cod}}:=C, \qquad \left(\tuple{u,v}\right)_{\mathrm{dom}}:=u, \qquad \left(\tuple{u,v}\right)_{\mathrm{cod}}:=v\] for the domain and codomain parts of objects and morphisms of \(\cat[]{LR}\).

In the correctness proof, the predicates represented in \(\cat[]{LR}\) will express differentiability and correctness of the syntactic reverse derivative. Accordingly, the source language must admit an interpretation in \(\cat[]{LR}\). This requires, in particular, that \(\cat[]{LR}\) carry the categorical structure needed for the language: finite products and coproducts, the FAD monad, and the distinguished objects and operations used to interpret base types and primitives.³ The remainder of this subsection verifies the FAD part of this structure by applying the FAD factorisation construction of Section 2.3.

The main structural input is the monadic-comonadic analysis of these comma categories from the CHAD logical-relations construction.

In the present case, \(\cat[]{C}=\Set\times\mathsf{Fam}(\Vect[o])\) is complete and cocomplete; hence so is \(\Set\times\cat[]{C}\). Since \(\cat[]{C}\) has copowers, the hom-functor \(G\) has a left adjoint, and Lemma 4 applies. Thus \[U:\cat[]{LR}\to\Set\times\cat[]{C}\] is both monadic and comonadic. The monadic part implies that \(U\) creates limits, while the comonadic part implies that \(U\) creates colimits. Hence \(\cat[]{LR}\) is complete and cocomplete, and its limits and colimits are computed by taking the corresponding limits and colimits in the underlying \(\Set\)- and \(\cat[]{C}\)-components, with the unique induced logical-relation structure.

The same componentwise description gives the factorisation system required by Section 2.3. A morphism \(\tuple{u,v}\) lies in the left class precisely when \(u\) and \(v\) lie in the corresponding left classes in \(\Set\) and \(\Set\times\mathsf{Fam}(\Vect[o])\); similarly for the right class. Writing \(v=\tuple{v_1,v_2}\), the left-class condition says that \(u\) and \(v_1\) are surjective and that \(v_2\) lies in the left class described in Section 2.4. The right-class condition says that \(u\) and \(v_1\) are injective and that \(v_2\) lies in the corresponding right class. This is the usual image factorisation, equivalently the (epi, regular mono) factorisation system in \(\cat[]{LR}\). Since the left classes in \(\Set\) and \(\mathsf{Fam}(\Vect[o])\) are closed under binary products, and \(U\) creates products, the left class in \(\cat[]{LR}\) is closed under binary products as well.

It remains to specify the monoid of weights and the barycentre semimodule in \(\cat[]{LR}\). Both are obtained by imposing the differentiability predicate on the corresponding structures in \(\cat[]{C}=\Set\times\mathsf{Fam}(\Vect[o])\). The monoid of weights is \(\tuple{ \opn{Diff} \left(\mathbb{E},\R\right), \tuple{\R,\TR}, \operatorname{dv}_{\R} }\) where \(\opn{Diff} \left(\mathbb{E},\R\right)\) denotes the set of differentiable functions \(\mathbb{E}\to\R\), and \(\operatorname{dv}_{\R}\) is the usual differentiability predicate. More generally, for each \(k\in\mathbb{N}\), \[\label{eq:definition-predicate-monoid-of-weights-in-lr} \operatorname{dv}_{\R[k]}: \opn{Diff} \left( \mathbb{E}, \R[k] \right) \to \fn[n]{ \Set }{ \mathbb{E}, \R[k] } \times \fn[n]{ \mathsf{Fam}(\Vect[o]) }{ T^{\ast}\mathbb{E}, \TR[k] } \qq{s.t.} f \mapsto \tuple{ f, \tuple{ f, \grad f } }.\tag{21}\] The barycentre semimodule is given by \(\tuple{ \opn{Diff} \left(\mathbb{E},\R\right), \tuple{\R,\TR}, \operatorname{dv}_{\R} }\).⁴

The structure maps for the monoid and the semimodule are induced componentwise from their counterparts in \(\Set\) and \(\mathsf{Fam}(\Vect[o])\), with the domain part given by postcomposition on differentiable functions. For example, the operation \(+_V\) is defined by \[\label{eq:definition-vplus-in-lr} \begin{array}{lllcl} \left(+_V\right)_{\mathrm{dom}} &:= +_V^{\Set} \circ \mathord{-} &:& \opn{Diff} \left( \mathbb{E}, \R[2] \right) &\to \opn{Diff} \left( \mathbb{E}, \R \right)\\ \left(+_V\right)_{\mathrm{cod}} &:= \tuple{ +_V^{\Set}, +_V^{\cat[]{F}} } &:& \tuple{ \R[2], \TR[2] } &\to \tuple{ \R, \TR } . \end{array}\tag{22}\] Here we use the canonical identifications \[\opn{Diff} \left(\mathbb{E},\R\right)\times \opn{Diff} \left(\mathbb{E},\R\right) \simeq \opn{Diff} \left(\mathbb{E},\R[2]\right) \quad \text{and} \quad \tuple{\R,\TR}\times\tuple{\R,\TR} \simeq \tuple{\R[2],\TR[2]}.\] The remaining structural morphisms are defined analogously. In each case, the codomain part is the corresponding morphism in \(\cat[]{C}=\Set\times\mathsf{Fam}(\Vect[o])\), while the domain part restricts it to the differentiable probes.

This componentwise construction is the sense in which the logical-relations structure is fibred over \(\cat[]{C}\): objects, structural maps, limits, colimits, and factorisations in \(\cat[]{LR}\) project to their corresponding constructions in \(\cat[]{C}\).

We may therefore apply the FAD factorisation construction of Section 2.3 to \(\cat[]{LR}\). The resulting monad \(M^{\cat[]{LR}}\) is fibred over \(\cat[]{C}=\Set\times\mathsf{Fam}(\Vect[o])\). Concretely, for an object \(\tuple{D,C,p}\) and a morphism \(\tuple{u,v}:\tuple{D,C,p}\to\tuple{D',C',p'}\) in \(\cat[]{LR}\), \[\label{eq:fad-monad-on-lr-fibred-over-c} \left(\Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ D, C, p } }\right)_{\mathrm{cod}} = \Fun[pn]{ M \times M^{\cat[]{F}} }{ C } \qq{and} \left(\Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ u, v } }\right)_{\mathrm{cod}} = \Fun[pn]{ M \times M^{\cat[]{F}} }{ v }.\tag{23}\]

The domain part records the admissible differentiable probes into the FAD object. Explicitly, \[\label{eq:fad-monad-on-lr-object-map-domain-part} \left(\Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ D, C, p } }\right)_{\mathrm{dom}} = \left\{ \omega: D \to \opn{Diff} \left( \mathbb{E}, \R \right) \;\middle|\; \begin{align} &\abs{ \operatorname{supp}\omega } <\infty \text{ and } \\ & \forall \left( x \in \operatorname{supp}\omega \right), \forall \left( v \in \mathbb{E}\right), \fn[n]{ \fn[n]{ \omega }{ x } }{ v } >0 \end{align} \right\}.\tag{24}\] Thus the domain part consists of finitely supported \(\opn{Diff} \left(\mathbb{E},\R\right)\)-valued weight distributions on \(D\), with strictly positive non-zero weights at every probe point. On morphisms, it is given by the usual pushforward of finite support: \[\label{eq:fad-monad-on-lr-morphism-map-domain-part} \left(\Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ u, v } }\right)_{\mathrm{dom}}: \left(\Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ D, C, p } }\right)_{\mathrm{dom}} \to \left(\Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ D', C', p' } }\right)_{\mathrm{dom}} \qq{s.t.} \omega \mapsto \left( y \mapsto \sum_{\substack{x \in \operatorname{supp}\omega \\ \fn[n]{ u }{ x } = y}} \fn[n]{ \omega }{ x } \right).\tag{25}\]

Let \(X\) denote the \(\Set\)-component of \(C\). When \(\mathbb{E}=\R\), the domain set \(\left( \Fun[nn]{ M^{\cat[]{LR}} }{ \tuple{ D, C, p } } \right)_{\mathrm{dom}}\) can be read as the class of differentiable curves \(h:\mathbb{E}\to\Fun[nn]{M}{X}\) selected by the logical relation. Locally, such a curve is represented by a finite weighted presentation \[\fn[n]{ h }{ v } = \fn[n]{ \operatorname{categorical}}{ \tuple{ \tuple{ \fn[n]{ f_{j} }{ v }, \fn[n]{ g_{j} }{ v } } }_{j = 1}^{n} },\] where the curves \(f_j:\mathbb{E}\to X\) are represented by elements of \(D\) through \(p\), and the functions \(g_j:\mathbb{E}\to\R\) are differentiable. The logical relation simultaneously records the corresponding transposed derivatives.

We have therefore constructed the associated FAD monad \(M^{\cat[]{LR}}\) on the logical-relations category. It is compatible with the codomain projection to \(\Set\times\mathsf{Fam}(\Vect[o])\), while its domain part selects the differentiable probes into finite atomic distributions, including differentiable atom data and differentiable log-weights. This completes the semantic preparation for the syntactic categories: the source language will be interpreted in \(\Set\), the differentiated target semantics in \(\mathsf{Fam}(\Vect[o])\), and the correctness argument in \(\cat[]{LR}\). The next section introduces the syntactic category of the source language and the structure that these semantic interpretations must preserve.

3 Syntactic categories↩︎

The CHAD perspective is to formulate automatic differentiation as a semantic construction. Rather than defining a derivative transformation merely as an ad hoc transformation, one first identifies the categorical structure freely generated by the language and then defines a functor preserving that structure. In this way, a code transformation is controlled by a categorical transformation. The syntactic category is the object that makes this passage precise.

Given a type theory, its syntactic category has as objects the types generated by the theory⁵ and as morphisms the \(\alpha\)-equivalence classes of programs. Composition is given by capture-avoiding substitution. The categorical structure of the syntactic category records the corresponding type-forming operations: for instance, products, coproducts, and exponentials interpret product, sum, and function types, respectively. An interpretation of the theory in a category \(\cat[]{X}\) is then a functor from the syntactic category to \(\cat[]{X}\) preserving the chosen structure.

In this paper we use two type theories, the source language and the target language. The source language is the language of programs to be differentiated. The target language extends it with the additional structure needed to express reverse-mode derivatives. The CHAD transformation sends source programs to target programs by the corresponding structure-preserving functor between their syntactic categories.

3.1 Source language↩︎

In this section, we describe the source language. Its syntactic category is the free distributive category generated by basic types \(\srcreal[n]\), \(n \in \mathbb{N}\), and basic built-in operations \(\mathsf{op}\), equipped with a chosen strong monad, a categorical operation, and a handler algebra.

Types in this language are given according to the following context-free grammar, \[\label{eq:source-language-type-construction-grammar} \tau, \sigma, \rho ::= 0 \mid 1 \mid \tau \times \sigma \mid \tau \sqcup \sigma \mid \srcreal[k] \left( k \in \mathbb{N}\right) \mid \opn[b]{M} \tau\tag{26}\] with \(\srcreal\) being used as abbreviation for \(\srcreal[1]\).

For terms, we denote variables by \(x, y, z, \dots\) and contexts as \(\Gamma = x_{1} : \tau_{1}, \cdots, x_{n} : \tau_{n}\). Term construction is defined by the following grammar \[\label{eq:source-language-term-construction-grammar} \begin{align} t, s, r ::= & \left. x \mid \tuple{\;} \mid \tuple{t, s} \mid \opn[b]{pr}_{1} t \mid \opn[b]{pr}_{2} t \mid \opn[b]{inl} t \mid \opn[b]{inr} t \right. \\ & \left. \mid \mathsf{case}\;t\;\mathsf{of}\;\mathsf{inl}(x)\Rightarrow s\;|\;\mathsf{inr}(y)\Rightarrow r \right. \left. \mid \fn[n]{ \opn[b]{abort} }{ t } \right. \\ & \left. \mid \mathsf{let}\;x=t\;\mathsf{in}\;s \right. \left. \mid \fn[n]{ \mathsf{op}}{ t_{1}, \dots, t_{k} } \right. \left. \mid \fn[n]{ \opn[b]{return} }{ t } \right. \\ & \left. \mid \fn[n]{ \opn[b]{bind} }{ t; x . s } \right. \left. \mid \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } \right. \left. \mid \fn[n]{ \opn[b]{E}_{\nu} }{ t } \right. \end{align}\tag{27}\] where \(\nu = \srcreal[p]\), for some \(p \in \mathbb{N}\), is a fixed type on which the handler algebra \(\opn[b]{E}_{\nu}\) is defined, and \(\mathsf{op}\) ranges over a family of chosen built-in operations, each with a signature \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\), representing the basic operations available in the source language. This family is required to include all affine maps and pointwise linear algebra on products and sums, but extra operations may be added depending on the use case: a differential equation solver might require Bessel functions and spherical harmonics, while a signal analyser might require trigonometric functions, for example.

Typing judgements for the \(0\), \(1\), product, and coproduct types are standard, and they may be found in Appendix 8. Any operation \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\) defines a typing rule \[\prftree[r]{\mathrm{\small Operations}}{ \Gamma \vdash t_{j}: \srcreal[n_{j}] \quad \left( j \in \left\{ 1, \dots, k \right\} \right) }{ \Gamma \vdash \fn[n]{ \mathsf{op}}{ t_{1}, \dots, t_{k} } : \srcreal[m] }\]

For types of the form \(\opn[b]{M} \tau\), we add the following extra rules \[\label{eq:source-language-monad-typing-rules} \begin{gather} \prftree[r]{\( \opn[b]{M} \text{E} \)}{ \Gamma \vdash t : \opn[b]{M} \nu }{ \Gamma \vdash \fn[n]{ \opn[b]{E}_{\nu} }{ t } : \nu } \qquad \prftree[r]{\( \opn[b]{M} \text{I1} \)}{ \Gamma \vdash t : \opn[b]{M} \sigma }{ \Gamma, x : \sigma \vdash s : \opn[b]{M} \tau }{ \Gamma \vdash \fn[n]{ \opn[b]{bind} }{ t; x . s } : \opn[b]{M} \tau } \\ \prftree[r]{\( \opn[b]{M} \text{I2} \)}{ \Gamma \vdash t : \tau }{ \Gamma \vdash \fn[n]{ \opn[b]{return} }{ t }: \opn[b]{M} \tau } \qquad \prftree[r]{\( \opn[b]{M} \text{I3} \)}{ \Gamma \vdash t_{i}: \tau }{ \Gamma \vdash w_{i}: \srcreal }{ \left( i \in \left\{ 1, \cdots, n \right\} \right) } { \Gamma \vdash \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{t_{i}, w_{i}} }_{i = 1}^{n} } : \opn[b]{M} \tau } \end{gather}\tag{28}\]

Finally, we add equations. In addition to the standard equational theory for the \(0\), \(1\), product, and coproduct types, we add equations for the bind, return, and algebra constructions. First, the Kleisli laws for the bind and return constructions are \[\label{eq:source-language-kleisli-laws} \begin{gather} \fn[n]{ \opn[b]{bind} }{ \fn[n]{ \opn[b]{return} }{ t }; x . s } \equiv s[t/x] \qquad \fn[n]{ \opn[b]{bind} }{ t; x . \fn[n]{ \opn[b]{return} }{ x } } \equiv t \\ \fn[n]{ \opn[b]{bind} }{ \fn[n]{ \opn[b]{bind} }{ t; x . s }; y . r } \equiv \fn[n]{ \opn[b]{bind} }{ t; x . \fn[n]{ \opn[b]{bind} }{ s ; y . r } } \end{gather}\tag{29}\] where we use \(s[t/x]\) for the capture-avoiding substitution of the term \(t\) for the variable \(x\) in the expression \(s\).

There are also equational rules for the handler algebra \(\opn[b]{E}_{\nu}\), which ensure that it behaves as an Eilenberg-Moore algebra for \(\opn[b]{M}\) over the type \(\nu = \srcreal[p]\), \[\label{eq:source-language-eilenberg-moore-algebra-laws} \fn[n]{ \opn[b]{E}_{\nu} }{ \fn[n]{ \opn[b]{return} }{ v } } \equiv v \qquad \fn[n]{ \opn[b]{E}_{\nu} }{ \fn[n]{ \opn[b]{bind} }{ t; x . u } } \equiv \fn[n]{ \opn[b]{E}_{\nu} }{ \fn[n]{ \opn[b]{bind} }{ t; x . \fn[n]{ \opn[b]{return} }{ \fn[n]{ \opn[b]{E}_{\nu} }{ u } } } }\tag{30}\]

We define the category \(\mathsf{Syn}\) as the syntactic category of this source language. In categorical terms, \(\mathsf{Syn}\) can be seen as the freely generated distributive category on the objects \(\srcreal[n]\), \(n \in \mathbb{N}\), and operations \(\mathsf{op}\), equipped with a fixed strong monad, a chosen \(\opn[b]{categorical}\) operation, and a chosen Eilenberg–Moore algebra \(\opn[b]{E}_{\nu}\) over some object \(\nu = \srcreal[p]\) for some \(p \in \mathbb{N}\).

Given any other distributive category \(\cat[]{X}\), with a strong monad \(M^{\cat[]{X}}: \cat[]{X} \to \cat[]{X}\), chosen categorical constructors, and an Eilenberg–Moore algebra \(\opn{alg}: \Fun[nn]{ M^{\cat[]{X}} }{ V } \to V\) for some object \(V\), we may define an interpretation functor \(\llbracket \cdot\rrbracket: \mathsf{Syn}\to \cat[]{X}\) by choosing, for each \(n \in \mathbb{N}\), an object \(X_{n}\) in \(\cat[]{X}\) such that \(X_{p} = V\). The object map of this functor is given by the type interpretations, which are inductively defined on the structure of the types as \[\label{eq:source-language-interpretation-inductive-definition} \begin{align} \llbracket 0\rrbracket &= 0_{\cat[]{X}} \\ \llbracket \tau \times \sigma\rrbracket &= \llbracket \tau\rrbracket \times \llbracket \sigma\rrbracket \end{align} \quad \begin{aligned} \llbracket 1\rrbracket &= 1_{\cat[]{X}} \\ \llbracket \tau \sqcup \sigma\rrbracket &= \llbracket \tau\rrbracket \amalg \llbracket \sigma\rrbracket \end{aligned} \quad \begin{align} \llbracket \srcreal[n]\rrbracket &= X_{n} \\ \llbracket \opn[b]{M} \tau\rrbracket &= \Fun[nn]{ M^{\cat[]{X}} }{ \llbracket \tau\rrbracket } \end{align}\tag{31}\] where \(0_{\cat[]{X}}\) and \(1_{\cat[]{X}}\) denote chosen initial and terminal objects in \(\cat[]{X}\), and \(\times\) and \(\amalg\) denote chosen product and coproduct structures, respectively.

Contexts are interpreted as products of the underlying types, that is, \(\llbracket x_{1}: \tau_{1}, \cdots, x_{n} : \tau_{n}\rrbracket = \prod_{i = 1}^{n} \llbracket \tau_{i}\rrbracket\), and a typing judgement \(\Gamma \vdash t : \tau\) is then interpreted as a morphism \(\llbracket \Gamma \vdash t : \tau\rrbracket: \llbracket \Gamma\rrbracket \to \llbracket \tau\rrbracket\).

Variables, constants, and constructors for the distributive structure are interpreted in the standard way (see Appendix 8). A built-in operator \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\) must be associated with a morphism \(\llbracket \mathsf{op}\rrbracket: \prod_{j = 1}^{k} X_{n_{j}} \to X_{m}\), so that \[\label{eq:source-language-built-in-operator-interpretation} \llbracket \Gamma \vdash \fn[n]{ \mathsf{op}}{ t_{1}, \cdots, t_{k} } : \srcreal[m]\rrbracket = \llbracket \mathsf{op}\rrbracket \circ \tupling{\llbracket \Gamma \vdash t_{j} : \srcreal[n_{j}]\rrbracket}_{j = 1}^{k}: \llbracket \Gamma\rrbracket \to X_{m}\tag{32}\] and the monadic operations are interpreted in terms of the unit \(\eta^{\cat[]{X}}\), multiplication \(\mu^{\cat[]{X}}\), and strength \(\operatorname{st}^{\cat[]{X}}\) of the monad \(M^{\cat[]{X}}\), \[\begin{align} \llbracket \Gamma \vdash \fn[n]{ \opn[b]{return} }{ t } : \opn[b]{M} \tau\rrbracket &= \Nat[nn]{ \eta^{\cat[]{X}} }{ \llbracket \tau\rrbracket } \circ \llbracket \Gamma \vdash t : \tau\rrbracket \\ \llbracket \Gamma \vdash \fn[n]{ \opn[b]{bind} }{ t; x . s } : \opn[b]{M} \tau\rrbracket &= \mu^{\cat[]{X}}_{\llbracket \tau\rrbracket} \circ \Fun[nn]{ M^{\cat[]{X}} }{ \llbracket \Gamma, x : \sigma \vdash s : \opn[b]{M} \tau\rrbracket } \circ \operatorname{st}^{\cat[]{X}}_{ \llbracket \Gamma\rrbracket, \llbracket \sigma\rrbracket } \circ \tupling{ \mathrm{id}_{\llbracket \Gamma\rrbracket}, \llbracket \Gamma \vdash t : \opn[b]{M} \sigma\rrbracket } \end{align}\]

Terms constructed with \(\opn[b]{categorical}\) are interpreted using the chosen morphisms \(\operatorname{categorical}_{\tau,n}^{\cat[]{X}}\). \[\llbracket \Gamma \vdash \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{t_{i}, w_{i}} }_{i = 1}^{n} } : \opn[b]{M} \tau \rrbracket = \operatorname{categorical}_{\tau,n}^{\cat[]{X}} \circ \tupling{ \tupling{ \llbracket \Gamma \vdash t_{i} : \tau\rrbracket, \llbracket \Gamma \vdash w_{i} : \srcreal\rrbracket } }_{i = 1}^{n}.\] For the FAD language treated here, no additional equations for these constructors are imposed.

Finally, the handler algebra \(\opn[b]{E}_{\nu}\) is interpreted as the chosen algebra on \(X_{p}\), \[\label{eq:source-language-handler-algebra-interpretation} \llbracket \Gamma \vdash \fn[n]{ \opn[b]{E}_{\nu} }{ t } : \nu\rrbracket = \opn{alg} \circ \llbracket \Gamma \vdash t : \opn[b]{M} \nu\rrbracket\tag{33}\]

Following this prescription, one can prove the following soundness theorem,

We may now state the universal property of the syntactic category of the source language.

The proof of existence and uniqueness follows from the prescribed construction above.

3.2 Interpretation of the source language in \(\Set\)↩︎

Applying the universal property of the source language (Lemma 6) to the data below gives a unique interpretation functor \(\llbracket \cdot\rrbracket_{\Set}: \mathsf{Syn}\to \Set\) interpreting the source language in the category of sets. For every \(n \in \mathbb{N}\), this interpretation is characterised by \[\label{eq:interpretation-reals-in-set} \llbracket \srcreal[n]\rrbracket_{\Set} := \R[n]\tag{34}\] and, for any built-in operator \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\) in the source language, we interpret \(\mathsf{op}\) as a differentiable function \[\label{eq:interpretation-operations-in-set} \llbracket \mathsf{op}\rrbracket_{\Set} : \R[n_{1}] \times \cdots \times \R[n_{k}] \to \R[m].\tag{35}\] The strong monad \(M^{\Set}\) is the FAD monad \(M\), with its unit, multiplication, and strength. The handler algebra for \(\nu = \srcreal[p]\) is interpreted as the expectation algebra \[\opn{alg}^{\Set}_{\nu}: \Fun[nn]{ M }{ \R[p] } \to \R[p], \qquad w \mapsto \sum_{x \in \operatorname{supp}w} x \cdot \fn[n]{ w }{ x }.\] Finally, \(\opn[b]{categorical}\) is interpreted using the construction from Equation 4 : for each type \(\tau\), we use \[\operatorname{categorical}^{\Set}_{\tau}:\coprod_{n \in \mathbb{N}} \left( \llbracket \tau\rrbracket_{\Set} \times \R \right)^{n} \to \Fun[nn]{ M }{ \llbracket \tau\rrbracket_{\Set} }.\]

3.3 Target language↩︎

Code written in the source language (Section 3.1) can be automatically differentiated using the CHAD algorithm. The transformed code is written in a separate language, called the target language.

The target language is dependently typed, split into a cartesian type theory and, for every context of cartesian types, a linear type theory. Semantically, the cartesian fragment of the target language allows one to reason about base spaces, while the linear type theories describe cotangents.

We denote by \(\CSyn\) the term model of the cartesian fragment of the target language. Crucially, \(\mathsf{Syn}\) can be seen as a full subcategory of \(\CSyn\). For each context \(\Gamma\) of cartesian types, we denote by \(\LSyn[\Gamma]\) the syntactic category of the associated linear type theory. Since any substitution \(\Gamma \to \Gamma'\) between cartesian contexts \(\Gamma\) induces a functor \(\LSyn[\Gamma'] \to \LSyn[\Gamma]\), we have a strict indexed category \(\LSyn: \CSyn[o] \to \Cat\). The syntactic category of the target language is then given by the Grothendieck construction \(\mathsf{Tgt}\) (see Section 2.3 of [2]). An explicit description of objects and morphisms in this category can be seen in Section 3.4.

As in the source language, we write \(\tau, \sigma, \rho\) for the cartesian types, and reserve underlined Greek letters \(\underline{\tau}, \underline{\sigma}, \underline{\rho}\) for linear types. We then have cartesian typing judgements of the form \(\Gamma \vdash t : \tau\) and linear typing judgements of the form \(\Gamma ; v : \underline{\sigma} \vdash s : \underline{\rho}\).

We can divide the description of the target language into three layers:

• Non-monadic core: deterministic operations. For brevity, we omit the description of this layer here (see Appendix  9).

• Oplax strong-monad core: rules for a strong monad once a compatible cotangent lift has been supplied.

• Distribution and handler: language features that are specific to the FAD monad.

In doing so, we isolate the data needed to apply CHAD to effects.

Cartesian types are described by the following grammar: \[\label{eq:target-language-cartesian-type-construction-grammar} \tau, \sigma, \rho ::= \left( \text{As in Equation~\eqref{eq:source-language-type-construction-grammar}} \right) \mid \underline{\sigma} \multimap\underline{\rho} \mid \Sigma \left( x : \sigma \right) . \tau'\tag{36}\] meaning that, in the cartesian fragment, we also have function types between linear types and dependent sums. For linear types, we use the grammar \[\label{eq:target-language-linear-type-construction-grammar} \underline{\tau}, \underline{\sigma}, \underline{\rho} ::= \underline{\srcreal[k]} \; \left( k \in \mathbb{N}\right) \mid \underline{1} \mid \underline{\tau} \times \underline{\sigma} \mid \opn[b]{case} \; t \; \opn[b]{of} \; \left\{ \opn[b]{in}_{i} x_{i} \mapsto \underline{\sigma}_{i} \right\}_{i = 1}^{n} \mid \opn[b]{M}^{w}_{p : \opn[b]{M} \tau} \left[ x \mapsto \underline{ \sigma } \right]\tag{37}\] meaning that the base Euclidean-space types have associated linear types (in the sense that the cotangent space of a Euclidean space at any point can be identified with the base space itself), as well as the fibre-wise zero object \(1\). Moreover, we allow fibre-wise binary biproducts \(\times\) and families of linear types indexed by finite disjunctions, and we introduce the cotangent linear type associated with monadic types.

We also distinguish cartesian and linear terms. The cartesian terms are described by the grammar \[\label{eq:target-language-cartesian-term-construction-grammar} t, s, r ::= \left( \text{As in Equation~\eqref{eq:source-language-term-construction-grammar}} \right) \mid \tuple{ a, t } \mid \opn[b]{case} \; p \; \opn[b]{of} \; \tuple{ x, y } \mapsto s \mid \lambda v \; . \; s\tag{38}\] where the constructor \(\tuple{ a, t }\) and the eliminator \(\opn[b]{case} \; p \; \opn[b]{of} \; \tuple{ x, y } \mapsto s\) are used in the typing rules for the cartesian \(\Sigma\)-type, and \(\lambda v \; . \; s\) is used for abstraction over the unique linear variable \(v\).

For linear terms, we use the following grammar: \[\label{eq:target-language-linear-term-construction-grammar} \begin{align} t, s, r ::= v &\mid \mathsf{let}\;v=t\;\mathsf{in}\;s \mid \tuple{ t, s } \mid \opn[b]{pr}_{i} t \mid \opn[b]{0} \mid t + s \mid \linapp{r}{t} \\ & \mid \fn[n]{ \fn[n]{ \mathsf{op}}{ t_{1}, \cdots, t_{k} } }{v} \mid \fn[n]{ \opn[b]{return}^{w}_{t} }{s} \mid \fn[n]{ \opn[b]{bind}^{w}_{s; x \; . \; t} }{r; u} \\ & \mid \fn[n]{ \opn[b]{categorical}^{w}_{\tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n}} }{ \tuple{ \tuple{ t_{i}', w_{i}' } }_{i = 1}^{n} } \mid \fn[n]{ \opn[b]{E}^{w}_{\nu, t} }{ s } \end{align}\tag{39}\] The three lines correspond to the three layers: the first line consists of the non-monadic core operations (linear variables, let bindings, pairings and projections, the zero vector, sums, and \(\linapp{r}{t}\) for linear application with \(r\) a cartesian term of type \(\underline{\sigma} \multimap\underline{\rho}\)). The second line has the derivatives of the built-in operations \(\mathsf{op}\) of the source language, as well as the linear fibre components of the monad operations \(\opn[b]{return}\) and \(\opn[b]{bind}\). Finally, the third line has the linear fibre components of \(\opn[b]{categorical}\) and of the handler algebra \(\opn[b]{E}_{\nu}\).

The typing rules for the non-monadic core are standard, so we omit them here (see Appendix 9).

There are two typing rules for the monadic core, corresponding to the \(\opn[b]{return}\) and \(\opn[b]{bind}\) operations. For \(\opn[b]{return}\), the rule is \[\label{eq:target-language-return-typing-law} \prftree{ \Gamma \vdash t : \tau }{ \Gamma, x : \tau \vdash \underline{\sigma} \text{ type} }{ \Gamma; v : \underline{\sigma}[t/x] \vdash s : \underline{\rho} }{ \Gamma; v : \opn[b]{M}^{w}_{\fn[n]{ \opn[b]{return} }{ t } : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ x \mapsto \underline{\sigma} \right] \vdash \fn[n]{ \opn[b]{return}^{w}_{t} }{ s }: \underline{\rho} }\tag{40}\] and for \(\opn[b]{bind}\), the rule is \[\label{eq:target-language-bind-typing-law} \begin{gather} \Gamma, x : \sigma \vdash \underline{\sigma} \text{ type} \qquad \Gamma, y : \tau \vdash \underline{\tau} \text{ type} \qquad \Gamma \vdash \underline{\rho} \text{ type} \qquad \Gamma \vdash s : \Fun[nn]{ \opn[b]{M} }{ \sigma } \\ \prftree{ \Gamma, x : \sigma \vdash t : \Fun[nn]{ \opn[b]{M} }{ \tau } }{ \Gamma; v : \opn[b]{M}^{w}_{s : \Fun[nn]{ \opn[b]{M} }{ \sigma }} \left[ x \mapsto \underline{\sigma} \right] \vdash r : \underline{\rho} }{ \Gamma, x : \sigma; v : \opn[b]{M}^{w}_{t : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ y \mapsto \underline{\tau} \right] \vdash u : \underline{\rho} \times \underline{\sigma} }{ \Gamma; v : \opn[b]{M}^{w}_{\fn[n]{ \opn[b]{bind} }{ s; x \; . \; t } : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ y \mapsto \underline{\tau} \right] \vdash \fn[n]{ \opn[b]{bind}^{w}_{s; x \; . \; t} }{ r; u } : \underline{\rho} } \end{gather}\tag{41}\]

These typing rules are the target-language interface for the chosen lifted strong monad. For the FAD monad, the lift \(M^{\cat[]{F}}\) also supplies transpose rules for \(\opn[b]{categorical}\) and for the handler algebra \(\opn[b]{E}_{\nu}\). For \(\opn[b]{categorical}\), we have, for all \(n \in \mathbb{N}\), \[\label{eq:target-language-typing-rule-categorical} \begin{gather} \Gamma, x : \tau \vdash \underline{\tau} \text{ type} \qquad \Gamma \vdash \underline{\rho} \text{ type} \\ \prftree{ \Gamma \vdash w_{i} : \srcreal }{ \Gamma \vdash t_{i}: \tau }{ \Gamma; v : \underline{\srcreal} \vdash w_{i}' : \underline{\rho} }{ \Gamma; v : \underline{\tau}[t_{i}/x] \vdash t_{i}' : \underline{\rho} }{ \left( 1 \le i \le n \right) }{ \Gamma; v : \opn[b]{M}^{w}_{ \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } : \Fun[nn]{ \opn[b]{M} }{ \tau } } \left[ x \mapsto \underline{\tau} \right] \vdash \fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } }{ \tuple{ \tuple{ t_{i}', w_{i}' } }_{i = 1}^{n} } : \underline{\rho} } \end{gather}\tag{42}\] and, for the handler algebra, recalling that \(\nu = \srcreal[p]\), we define \(\underline{\nu} := \underline{\srcreal[p]}\). Then, \[\label{eq:target-language-typing-rule-handler-algebra} \prftree{ \Gamma \vdash t : \Fun[nn]{ \opn[b]{M} }{ \nu } }{ \Gamma; v : \opn[b]{M}^{w}_{t: \Fun[nn]{ \opn[b]{M} }{ \nu }} \left[ \underline{\hphantom{x}} \mapsto \underline{\nu} \right] \vdash s : \underline{\rho} }{ \Gamma; v : \underline{\nu} \vdash \fn[n]{ \opn[b]{E}^{w}_{\nu, t} }{ s } : \underline{\rho} }\tag{43}\]

The equations for the cartesian fragment are standard, whereas those for the linear fragment can be seen as originating from a formal derivative procedure over the cartesian base counterparts (see Appendix 9). For instance, the right unit law for \(\opn[b]{return}\) and \(\opn[b]{bind}\) (see Equation 29 ) \[\fn[n]{\opn[b]{bind}}{ s; x.\fn[n]{\opn[b]{return}}{x} } \equiv s.\] When imposed on the cartesian fragment, this law gives rise to the equality \[\label{eq:target-language-monad-w-right-unit} \fn[n]{ \opn[b]{bind}^{w}_{ s; x.\fn[n]{\opn[b]{return}}{x} } }{ r; \tuple{ \opn[b]{0}, \fn[n]{\opn[b]{return}^{w}_{x}}{v} } } \equiv r .\tag{44}\]

3.4 Interpretation of the target language in \(\mathsf{Fam}(\Vect[o])\)↩︎

In Section 3.2, we provided set-based semantics for the source language via the interpretation functor \(\llbracket \cdot\rrbracket_{\Set} : \mathsf{Syn}\to \Set\). Here we provide concrete semantics for the target language syntactic category \(\mathsf{Tgt}\) by interpreting it in \(\mathsf{Fam}(\Vect[o])\) (see Section 2.4).

An object of \(\mathsf{Tgt}\) may be seen as a pair \(\tuple{ p : \tau, \underline{\sigma} }\) (up to \(\alpha\)-equivalence) such that \(\underline{\sigma}\) is a linear type in \(\LSyn[p : \tau]\). A morphism \(\tuple{ t, s }: \tuple{ p : \tau, \underline{\rho} } \to \tuple{ q : \sigma, \underline{\chi} }\) is a pair of a cartesian arrow \(t: p \to q\) in \(\CSyn\), and of a linear arrow \(s: \underline{\chi}[t/q] \to \underline{\rho}\) in \(\LSyn[p : \tau]\).

We start by providing an interpretation of the cartesian fragment \(\CSyn\) of the target language. This fragment is interpreted in \(\Set\) by extending the interpretation of \(\mathsf{Syn}\), as described in Section 3.2. For that reason, we keep the same notation, writing \(\llbracket 0\rrbracket_{\Set} := \emptyset\), \(\llbracket \srcreal[k]\rrbracket_{\Set} := \R[k]\), \(\llbracket \tau \times \sigma\rrbracket_{\Set} := \llbracket \tau\rrbracket_{\Set} \times \llbracket \sigma\rrbracket_{\Set}\), and so on. The linear interpretation of a linear type under a cartesian context \(p : \tau\) is also described inductively over the structure of the linear types, and is described in Appendix 10. In particular, we give \(\fn[n]{ \llbracket \underline{\srcreal[k]} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } := \R[k]\) and, for any cartesian judgement \(p: \tau \vdash q:\Fun[nn]{\opn[b]{M}}{\rho}\), \(\fn[n]{ \llbracket \opn[b]{M}^{w}_{q : \Fun[nn]{ \opn[b]{M} }{ \rho }} \left[ x \mapsto \underline{\sigma} \right] \rrbracket_{p : \tau}^{\mathrm{lin}} }{\gamma} = \bigoplus_{x \in \operatorname{supp}\fn[n]{ \llbracket q\rrbracket_{\Set} }{\gamma} } \left( \fn[n]{ \llbracket \underline{\sigma}\rrbracket_{p : \tau, x : \rho}^{\mathrm{lin}} }{ \gamma, x } \oplus \R \right)\) for any cartesian variable \(p : \tau\) and \(\gamma \in \llbracket \tau\rrbracket_{\Set}\).

So, for any object \(\tuple{ p : \tau, \underline{\sigma} }\) in \(\mathsf{Tgt}\), we define its interpretation in \(\mathsf{Fam}(\Vect[o])\) as \[\label{eq:target-language-fam-vectop-interpretation-objects} \llbracket \tuple{ p : \tau, \underline{\sigma} }\rrbracket_{\cat[]{F}} := \tuple{ \llbracket \tau\rrbracket_{\Set}, \left( \gamma \mapsto \fn[n]{ \llbracket \underline{\sigma}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } \right) }\tag{45}\] and morphisms are interpreted, displaying the corresponding ordinary \(\Vect\)-map in the fibre, as \[\label{eq:target-language-fam-vectop-interpretation-morphisms} \llbracket \tuple{ t, s }\rrbracket_{\cat[]{F}} := \tuple{ \llbracket t\rrbracket_{\Set}, \left( \gamma \mapsto \left( \opmor[p]{ \fn[n]{ \llbracket s\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \in \fn[n]{ \Vect }{ \fn[n]{ \llbracket \underline{\chi}\rrbracket_{q : \sigma}^{\mathrm{lin}} }{ \fn[n]{ \llbracket t\rrbracket_{\Set} }{ \gamma } }, \fn[n]{ \llbracket \underline{\rho}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \right) \right) }\tag{46}\]

3.5 CHAD code transformation↩︎

The core idea of CHAD is the translation of the automatic differentiation code transformation into a structure-preserving categorical transformation, i.e. a functor. This is a functor \(\chad: \mathsf{Syn}\to \mathsf{Tgt}\). For each term in context \(\Gamma \vdash t : \tau\) in the source language, we get a term under context \[\label{eq:chad-code-transformation-terms-under-context} \Fun[np]{ \chad }{ \Gamma }_{1} \vdash \Fun[np]{ \chad[\Gamma] }{ t } : \Sigma \left( p : \Fun[np]{ \chad }{ \tau }_{1} \right) \; . \; \left( \Fun[np]{ \chad }{ \tau }_{2} \multimap \Fun[np]{ \chad }{ \Gamma }_{2} \right)\tag{47}\] In other words, the term under context \(\Gamma \vdash t : \tau\) is translated as a term of this dependent-sum type: its first projection has type \(\Fun[np]{ \chad }{ \tau }_{1}\), and the corresponding fibres are given by the function types \(\Fun[np]{ \chad }{ \tau }_{2} \multimap\Fun[np]{ \chad }{ \Gamma }_{2}\), corresponding to transposed derivatives (backpropagators).

Similar to denotational semantics, type and term transformations are defined inductively over the structure of the types and terms. The definitions for base types, as well as for products and coproducts is straightforward, and can be seen in Appendix 11. For the monadic types, we define \[\label{eq:chad-code-transformation-measure-type-derivative} \Fun[np]{ \chad }{ \Fun[nn]{ \opn[b]{M} }{ \tau } }_{1} := \Fun[nn]{ \opn[b]{M} }{ \Fun[nn]{ \chad }{ \tau }_{1} } \qq{and} \Fun[np]{ \chad }{ \Fun[nn]{ \opn[b]{M} }{ \tau } }_{2} := \opn[b]{M}^{w}_{p : \Fun[nn]{ \opn[b]{M} }{ \Fun[np]{ \chad }{ \tau }_{1} }} \left[ x \mapsto \Fun[np]{ \chad }{ \tau }_{2} [x/p] \right]\tag{48}\] and the terms related to the monadic structure have their derivatives defined as \[\begin{align} \Fun[np]{ \chad[\Gamma] }{ \fn[n]{ \opn[b]{return} }{ t } } &:= \opn[b]{case} \; \Fun[np]{ \chad[\Gamma] }{ t } \; \opn[b]{of} \; \tuple{ p, b } \to \tuple{ \fn[n]{ \opn[b]{return} }{ p }, \lambda v \; . \; \fn[n]{ \opn[b]{return}^{w}_{p} }{ \fn[n]{ b }{ v } } } \tag{49} \\ \Fun[np]{ \chad[\Gamma] }{ \fn[n]{ \opn[b]{bind} }{ t; x \; . \; s } } &:= \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma] }{ t } \; \opn[b]{of} \tuple{ p, b } \to \opn[b]{let} \; r = \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma, x : \sigma] }{ s } \; \opn[b]{of} \; \tuple{ q, c } \to q \; \opn[b]{in} \tag{50} \\ &\qquad \tuple{ \fn[n]{ \opn[b]{bind} }{ p; x \; . \; r }, \lambda v \; . \; \fn[n]{ \opn[b]{bind}^{w}_{p; x \; . \; r} }{ \fn[n]{ b }{ v }; \fn[p]{ \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma, x : \sigma] }{ s } \; \opn[b]{of} \; \tuple{ q, c } \to c }{ v } } } \notag \end{align}\] while the term constructors specific to the FAD monad have their code transformations defined as \[\label{eq:chad-code-transformation-categorical} \begin{align} &\Fun[np]{ \chad[\Gamma] }{ \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } } := \\ & \qquad \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma] }{ t_{1} } \; \opn[b]{of} \tuple{ u_{1}, u_{1}' } \to \cdots \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma] }{ t_{n} } \; \opn[b]{of} \tuple{ u_{n}, u_{n}' } \to \\ & \qquad \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma] }{ w_{1} } \; \opn[b]{of} \tuple{ a_{1}, a_{1}' } \to \cdots \opn[b]{case} \; \Fun[nn]{ \chad[\Gamma] }{ w_{n} } \; \opn[b]{of} \tuple{ a_{n}, a_{n}' } \to \\ & \qquad \qquad \tuple{ \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ u_{i}, a_{i} } }_{i = 1}^{n} }, \lambda v \; . \; \fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{ \tuple{ u_{i}, a_{i} } }_{i = 1}^{n} } }{ \tuple{ \tuple{ \fn[n]{ u_{i}' }{ v }, \fn[n]{ a_{i}' }{ v } } }_{i = 1}^{n} } } \end{align}\tag{51}\] \[\label{eq:chad-code-transformation-handler-algebra}\qquad\qquad\qquad \Fun[np]{ \chad[\Gamma] }{ \fn[n]{ \opn[b]{E}_{\nu} }{ t } } := \opn[b]{case} \; \Fun[np]{ \chad[\Gamma] }{ t } \; \opn[b]{of} \tuple{ p, b } \to \tuple{ \fn[n]{ \opn[b]{E}_{\nu} }{ p }, \lambda v \; . \; \fn[n]{ \opn[b]{E}^{w}_{\nu, p} }{ \fn[n]{ b }{ v } } }\tag{52}\]

4 Logical relations and correctness↩︎

4.1 Interpretation of source language in the category of logical relations↩︎

As in \(\Set\), we can use the universal property of the syntactic category of the source language to construct an interpretation of the source language in the category \(\cat[]{LR}\) of logical relations. It is enough to specify the interpretations of the elementary real types, the built-in operations, the monadic type constructor, \(\opn[b]{categorical}\), and the handler algebra.

For the elementary real number types, we define, for all \(n \in \mathbb{N}\), \[\label{eq:interpretation-reals-in-lr} \llbracket \srcreal[n]\rrbracket_{\cat[]{LR}} := \tuple{ \opn{Diff} \left( \mathbb{E}, \R[n] \right), \tuple{ \R[n], \TR[n] }, \operatorname{dv}_{\R[n]} }\tag{53}\] and, for any built-in operation \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\) in the source language, its domain part is given by postcomposition with \(\llbracket \mathsf{op}\rrbracket_{\Set}\), while its codomain part is the base function together with its transposed derivative, \[\label{eq:interpretation-operations-in-lr} \begin{align} \left(\llbracket \mathsf{op}\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} &:= \left( \llbracket \mathsf{op}\rrbracket_{\Set} \circ \mathord{-}\right): \underbrace{ \opn{Diff} \left( \mathbb{E}, \R[n_{1}] \right) \times \cdots \times \opn{Diff} \left( \mathbb{E}, \R[n_{k}] \right) }_{ \text{identified with } \opn{Diff} \left( \mathbb{E}, \R[n_{1}] \times \cdots \times \R[n_{k}] \right) } \to \opn{Diff} \left( \mathbb{E}, \R[m] \right) \\ \left(\llbracket \mathsf{op}\rrbracket_{\cat[]{LR}}\right)_{\mathrm{cod}} &:= \tuple{ \llbracket \mathsf{op}\rrbracket_{\Set}, \tuple{ \llbracket \mathsf{op}\rrbracket_{\Set}, \grad \llbracket \mathsf{op}\rrbracket_{\Set} } }: \tuple{ \prod_{j = 1}^{k} \R[n_{j}], \prod_{j = 1}^{k} \TR[n_{j}] } \to \tuple{\R[m], \TR[m]} \end{align}\tag{54}\]

The strong monad is taken to be the FAD monad \(M^{\cat[]{LR}}\) on \(\cat[]{LR}\) (see Section 2.5), with the corresponding interpretations of unit, multiplication, and strength. The interpretation of \(\opn[b]{categorical}\) is given by the factorisation procedure on \(\cat[]{LR}\) described in Section 2.5. For \(\nu = \srcreal[p]\), the handler algebra \(\opn[b]{E}_{\nu}\) is interpreted by the corresponding \(M^{\cat[]{LR}}\)-algebra \(\opn{alg}^{\cat[]{LR}}_{\nu}: \Fun[nn]{ M^{\cat[]{LR}} }{ \llbracket \nu\rrbracket_{\cat[]{LR}} } \to \llbracket \nu\rrbracket_{\cat[]{LR}},\) whose codomain part is the pair consisting of the expectation algebra in \(\Set\) and the lifted expectation algebra in \(\mathsf{Fam}(\Vect[o])\), and whose domain part sends a differentiable function of FADs to the differentiable expectation function obtained by applying that algebra pointwise.

It then follows from the universal property of the syntactic category of the source language (Lemma 6) that these data define a unique interpretation functor \(\llbracket \cdot\rrbracket_{\cat[]{LR}}: \mathsf{Syn}\to \cat[]{LR}\).

Crucially, by composing this interpretation functor with \(U_{\Set} := \Pi_{1} \circ \Pi_{2} \circ U\), where \(U: \cat[]{LR} \to \Set \times \cat[]{C}\) is the comma-category forgetful functor, \(\Pi_{2}: \Set \times \cat[]{C} \to \cat[]{C}\) is the second projection into \(\cat[]{C} = \Set \times \mathsf{Fam}(\Vect[o])\), and \(\Pi_{1}: \cat[]{C} \to \Set\) is the first projection, we get another distributive functor. Indeed, \(U\) creates limits and colimits, \(\Pi_{1}\) and \(\Pi_{2}\) preserve them, and the codomain part of \(M^{\cat[]{LR}}\) is fibred over the FAD monads in \(\Set\) and \(\mathsf{Fam}(\Vect[o])\). This functor \(U_{\Set} \circ \llbracket \cdot\rrbracket_{\cat[]{LR}}\) satisfies the same data as the \(\Set\)-interpretation, in particular Equations 34 and 35 . Hence, by uniqueness, \[\label{eq:lr-interpretation-fibred-over-set} U_{\Set} \circ \llbracket \cdot\rrbracket_{\cat[]{LR}} = \llbracket \cdot\rrbracket_{\Set}\tag{55}\]

Analogously, composing \(\llbracket \cdot\rrbracket_{\cat[]{LR}}: \mathsf{Syn}\to \cat[]{LR}\) with \(U_{\cat[]{F}} := \Pi_{2} \circ \Pi_{2} \circ U: \cat[]{LR} \to \mathsf{Fam}(\Vect[o])\), we get a distributive functor \(\mathsf{Syn}\to \mathsf{Fam}(\Vect[o])\) that interprets the monad type constructor (and its associated unit and multiplication) as the FAD monad in \(\mathsf{Fam}(\Vect[o])\). It follows from the universal property of \(\mathsf{Syn}\) that \[\label{eq:lr-interpretation-fibred-over-fam-vectop} U_{\cat[]{F}} \circ \llbracket \cdot\rrbracket_{\cat[]{LR}} = \llbracket \cdot\rrbracket_{\cat[]{F}} \circ \chad\tag{56}\]

4.2 Correctness Theorem↩︎

In this section, we use the previously defined machinery to prove the following correctness theorem.

To prove this, fix a well-typed source term \(x: \srcreal[n] \vdash t : \srcreal[m]\). Interpreting this judgement in the logical-relations category via the functor \(\llbracket \cdot\rrbracket_{\cat[]{LR}}: \mathsf{Syn}\to \cat[]{LR}\), we get a morphism \[\llbracket t\rrbracket_{\cat[]{LR}}: \tuple{ \opn{Diff} \left( \mathbb{E}, \R[n] \right), \tuple{ \R[n], \TR[n] }, \operatorname{dv}_{\R[n]} } \to \tuple{ \opn{Diff} \left( \mathbb{E}, \R[m] \right), \tuple{ \R[m], \TR[m] }, \operatorname{dv}_{\R[m]} }\] in \(\cat[]{LR}\). By the defining commutative square for morphisms in a comma category, \[\Fun[nn]{ G }{ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{cod}} } \circ \operatorname{dv}_{\R[n]} = \operatorname{dv}_{\R[m]} \circ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}}: \opn{Diff} \left( \mathbb{E}, \R[n] \right) \to \fn[n]{ \cat[]{C} }{ \tuple{ \mathbb{E}, T^{\ast}\mathbb{E}}, \tuple{ \R[m], \TR[m] } }.\]

Since \(\mathbb{E}\) was an arbitrary finite-dimensional real vector space, we may take \(\mathbb{E}= \R[n]\). In this case, we apply both sides of the equality to \(\mathrm{id}_{\R[n]} \in \opn{Diff} \left( \R[n], \R[n] \right)\). Using \(\operatorname{dv}_{\R[n]} \left( \mathrm{id}_{\R[n]} \right) = \mathrm{id}_{\tuple{ \R[n], \TR[n] }}\), the left-hand side is \[\fn[p]{ \Fun[nn]{ G }{ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{cod}} } \circ \operatorname{dv}_{\R[n]} }{ \mathrm{id}_{\R[n]} } = \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{cod}} \circ \mathrm{id}_{\tuple{ \R[n], \TR[n] }} = \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{cod}} = \Fun[pp]{ \Pi_{2} \circ U }{ \llbracket t\rrbracket_{\cat[]{LR}} },\] and the right-hand side is \[\fn[p]{ \operatorname{dv}_{\R[m]} \circ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ \mathrm{id}_{\R[n]} } = \tuple{ \fn[n]{ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ \mathrm{id}_{\R[n]} }, \tuple{ \fn[n]{ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ \mathrm{id}_{\R[n]} }, \grad \left( \fn[n]{ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ \mathrm{id}_{\R[n]} } \right) } }.\] Applying the first projection \(\Pi_{1}: \cat[]{C} \to \Set\) to both expressions, it follows from Equation 55 that \[\fn[n]{ \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ \mathrm{id}_{\R[n]} } = \Fun[pp]{ \Pi_{1} \circ \Pi_{2} \circ U }{ \llbracket t\rrbracket_{\cat[]{LR}} } = \Fun[nn]{ U_{\Set} }{ \llbracket t\rrbracket_{\cat[]{LR}} } = \llbracket t\rrbracket_{\Set}\] and thus, we can prove Theorem 2, by using Equation 56 : \[\begin{align} \llbracket \Fun[np]{ \chad }{ t }\rrbracket_{\cat[]{F}} &= \Fun[nn]{ U_{\cat[]{F}} }{ \llbracket t\rrbracket_{\cat[]{LR}} } = \Fun[np]{ \Pi_{2} }{ \Fun[pp]{ \Pi_{2} \circ U }{ \llbracket t\rrbracket_{\cat[]{LR}} } } = \tuple{ { \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ (\mathrm{id}_{\R[n]}}), \grad \left( { \left(\llbracket t\rrbracket_{\cat[]{LR}}\right)_{\mathrm{dom}} }{ (\mathrm{id}_{\R[n]}) } \right) }= \tuple{ \llbracket t\rrbracket_{\Set}, \grad{\llbracket t\rrbracket_{\Set}} }. \end{align}\]

5 Returning to the example↩︎

In this section, we return to the stochastic optimisation problem discussed in Section 1.1. The goal is to encode it in the source language and then apply the code transformation. We then study how the transformed code relates to the mathematical description of the derivative semantics.

5.1 Encoding the example in the source language↩︎

As a first step in encoding the example in the source language, we must set up the family of built-in operations. In addition to the basic operations for representing affine maps and linear algebra, we also add five nullary operations (i.e. constants) \(s_{m}, s_{p}, m_{0}, p_{0}, P : \srcreal\) and two unary operations \(\opn[b]{sig}, \opn[b]{lsig}: \srcreal \to \srcreal\), representing the sigmoid and the logarithm of the sigmoid function, respectively. Then, in the context \(\Gamma = p : \srcreal, m : \srcreal\), the program is given by the term \[\label{eq:example-source-code-encoding} \begin{align} &\opn[b]{let} \; s = \opn[b]{categorical} \left( \left \langle \tuple{ 0, \fn[n]{ \opn[b]{lsig} }{ -s_{m} \cdot \left( m - m_{0} \right) } }, \right. \right. \\ &\qquad \qquad \qquad \qquad \hphantom{a} \left. \left. \tuple{ P \cdot \fn[n]{ \opn[b]{sig} }{ -s_{p} \cdot \left( p - p_{0} \right) }, \fn[n]{ \opn[b]{lsig} }{ s_{m} \cdot \left( m - m_{0} \right) } } \right\rangle \right) \\ &\opn[b]{in} \; \opn[b]{let} \; t = \fn[n]{ \opn[b]{bind} }{s; x . \fn[n]{ \opn[b]{return} }{ p \cdot x - m } } \; \opn[b]{in} \; \fn[n]{ \opn[b]{E}_{\srcreal} }{ t } \end{align}\tag{58}\] where we have chosen \(\nu = \srcreal\). First, we introduce a distribution \(s\) for the two possible total-sales scenarios. Then, we map the function \(x \mapsto p \cdot x - m\) over those scenarios using \(\opn[b]{bind}\) and \(\opn[b]{return}\), giving the profit distribution \(t\). Finally, we compute the expected total profit using the \(\opn[b]{E}_{\srcreal}\) handler.

In this simple example, the final result is an expected value. However, the source language allows different expected values to be mixed within the same expression. For instance, suppose that one wants to maximise a risk-adjusted profit given by the ratio of the square of the expected profit to the profit variance. The source-language version of this calculation is \[\begin{align} &\opn[b]{let} \; s = \opn[b]{categorical} \left( \left \langle \tuple{ 0, \fn[n]{ \opn[b]{lsig} }{ -s_{m} \cdot \left( m - m_{0} \right) } }, \right. \right. &&\text{\textcolor{lightgray}{// Sales distribution}} \\ &\qquad \qquad \qquad \qquad \hphantom{a} \left. \left. \tuple{ P \cdot \fn[n]{ \opn[b]{sig} }{ -s_{p} \cdot \left( p - p_{0} \right) }, \fn[n]{ \opn[b]{lsig} }{ s_{m} \cdot \left( m - m_{0} \right) } } \right\rangle \right) \\ & \opn[b]{in} \; \opn[b]{let} \; t = \fn[n]{ \opn[b]{bind} }{s; x . \fn[n]{ \opn[b]{return} }{ p \cdot x - m } } &&\text{\textcolor{lightgray}{// Profit distribution}} \\ & \opn[b]{in} \; \opn[b]{let} \; a = \fn[n]{ \opn[b]{E}_{\srcreal} }{ t } &&\text{\textcolor{lightgray}{// Expected profit}} \\ & \opn[b]{in} \; \opn[b]{let} \; b = \fn[n]{ \opn[b]{bind} }{ t; x . \fn[n]{ \opn[b]{return} }{ x^{2} } } &&\text{\textcolor{lightgray}{// Squared profit}} \\ & \opn[b]{in} \; \opn[b]{let} \; v = \fn[n]{ \opn[b]{E}_{\srcreal} }{ b } - a^{2} &&\text{\textcolor{lightgray}{// Variance}} \\ & \opn[b]{in} \; a^{2} / v \end{align}\] The source language also supports more complex combinations of expected values, including expected values of expressions that themselves contain expected values.

5.2 Code transformation↩︎

We now consider the differentiated code from the example in Equation 58 . The full code can be seen in Appendix 12. Here we highlight some of the main features.

Recall that the differentiation is done with respect to the context \(\Gamma = p : \srcreal, m : \srcreal\), and we adopt the following abbreviations: \[\begin{align} a_{0} &:= 0 \\ a_{1} &:= P \cdot \fn[n]{ \opn[b]{sig} }{ -s_{p} \cdot \left( p - p_{0} \right) } \end{align} \qquad \begin{align} \gamma_{0} &:= \fn[n]{ \opn[b]{lsig} }{ -s_{m} \cdot \left( m - m_{0} \right) } \\ \gamma_{1} &:= \fn[n]{ \opn[b]{lsig} }{ s_{m} \cdot \left( m - m_{0} \right) } \end{align}\]

The first lines of the differentiated code are case statements that pattern match on \[\begin{align} &\opn[b]{case} \; \tuple{ a_{0}, \grad a_{0} } \; \opn[b]{of} \tuple{ \theta_{0}, \theta_{0}' } \to \opn[b]{case} \; \tuple{ a_{1}, \grad a_{1} } \; \opn[b]{of} \tuple{ \theta_{1}, \theta_{1}' } \to \\ &\opn[b]{case} \; \tuple{ \gamma_{0}, \grad \gamma_{0} } \; \opn[b]{of} \tuple{ \rho_{0}, \rho_{0}' } \to \opn[b]{case} \; \tuple{ \gamma_{1}, \grad \gamma_{1} } \; \opn[b]{of} \tuple{ \rho_{1}, \rho_{1}' } \to \\ &\qquad \tuple{ \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ \theta_{j}, \rho_{j} } }_{j = 0}^{1} }, \lambda v \; . \; \fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{ \tuple{ \theta_{j}, \rho_{j} } }_{j = 0}^{1} } }{ \tuple{ \tuple{ \fn[n]{ \theta_{j}' }{ v }, \fn[n]{ \rho_{j}' }{ v } } }_{j = 0}^{1} } } \end{align}\] These lines correspond to the derivative of \(\opn[b]{categorical}\). Since \(a_{i} = \theta_{i}\) and \(\gamma_{i} = \rho_{i}\), the first entry of the tuple is simply a recalculation of the same categorical distribution of sales. The second entry uses the \(\opn[b]{categorical}^{w}\) linear term construction to construct a cotangent value combining the transposed derivative of the atoms (\(\theta_{i}'\)) and of the log-weights (\(\rho_{i}'\)), as is the case with the denotational semantics in \(\mathsf{Fam}(\Vect[o])\) (Section 2.4).

Another point to highlight is the block concerning the transposed derivative of the handler: \[\begin{array}{lll} & \opn[b]{case} & \opn[b]{case} \tuple{ t, \lambda v \; . \; \opn[b]{coproj}_{\fn[n]{ \opn{idx} }{ t; \Gamma, s, t }} v} \\ & & \opn[b]{of} \tuple{ \zeta, \zeta' } \to \tuple{ \fn[n]{ \opn[b]{E}_{\srcreal} }{ \zeta }, \lambda v \; . \; \fn[n]{ \opn[b]{E}_{\srcreal, \zeta}^{w} }{ \fn[n]{ \zeta' }{ v } } } \\ & \opn[b]{of} &\tuple{ \epsilon, \epsilon' } \to \tuple{ \epsilon[\delta/t], \lambda v \; . \; \mathsf{let}\;w= \fn[n]{ \epsilon'[\delta/t] }{ v } \;\mathsf{in}\; \opn[b]{pr}_{1} w + \fn[n]{ \delta' }{ \opn[b]{pr}_{2} w } } \end{array}\] Here, the transposed derivative appears as the transposed derivative of the variable \(t\) with respect to the extended context \(\Gamma, s : \Fun[nn]{ \opn[b]{M} }{ \srcreal }, t : \Fun[nn]{ \opn[b]{M} }{ \srcreal }\). We then compute its expected value, and use the \(\opn[b]{E}^{w}_{\srcreal}\) term constructor to compute the derivative of the handler algebra. In the last line, \(t\) is substituted by the \(\delta\) term, whose definition is given in a surrounding case construction.

6 Beyond probability: discrete-output algebraic effects↩︎

The FAD monad is the main case study, but our construction is not intrinsically probabilistic. The reusable data are a strong monad presented by discrete-output operations, a chosen Eilenberg–Moore handler algebra, and compatible lifts of the monad, the operations, and the handler to \(\mathsf{Fam}(\Vect[o])\) and to the logical-relations category. With these data, the cotangent rules for algebraic operations, generic \(\opn[b]{return}\) and \(\opn[b]{bind}\), and handlers from Sections 3.1–3.4 have the same form.

For a typed algebraic signature \(\Sigma\) with generic effects \(\omega:I_{\omega}\to T O_{\omega}\), the parameter types \(I_\omega\) may be differentiable: \(\llbracket I_{\omega}\rrbracket_{\cat[]{F}}\) may have non-trivial cotangent fibres. The direct fibred-polynomial construction requires the operation output type to be finite and discrete \(O_\omega= 1\sqcup \cdots\sqcup 1\) (so \(\llbracket O_{\omega}\rrbracket_{\cat[]{F}}\) has trivial fibre). Before quotienting by equations, the free AST monad has the form \[\mathsf{AST}_{\Sigma}X \simeq \mu Y.\; X \sqcup \coprod_{\omega\in\Sigma} I_{\omega}\times \prod_{o\in O_{\omega}}Y:\] we take a finite product of continuation subtrees. The lift to \(\mathsf{Fam}(\Vect[o])\) adds cotangents for the parameter and branches, while the branch index \(o\) has zero fibre. (Outputs with non-trivial fibres can destroy fibredness.) Appendix 15 gives the corresponding generic CHAD theorem.

6.0.0.1 Finite non-determinism

For finite non-determinism, we replace the FAD monad with the finite multiset monad \(\mathsf{MFin}X = \coprod_{k \in \mathbb{N}} \left(X^{k}/\mathfrak{S}_{k}\right),\) where multiplicity is retained but order is forgotten. From the signature with no equations, we can take the corresponding \(\opn{AST}\) monad as the list monad (where order is relevant). This is analogous to what happens to FADs: where FADs can be seen as lists of point/log-weight pairs modulo collisions and reordering, multisets can be seen as lists modulo reordering. These additional equations are then imposed by an \(\opn{AST}\)-algebra.

For multisets, a \(\opn{AST}\)-algebra (i.e. a \(\opn{List}\)-algebra) is precisely a monoid structure on the underlying object. By making this monoid commutative, one then imposes order-irrelevance. Thus \(\tuple{\R[n],0,+}\), for example, induces a \(\opn{List}\)-algebra \(\opn{sum}: \Fun[np]{ \opn{List} }{ \R[n] } \to \R[n]\) for multiplicity-sensitive finite non-deterministic computations valued in \(\R[n]\). The ordinary lifted monad also quotients the fibres: over a multiset represented by \(\tuple{x_{1},\dots,x_{k}}\), the fibre is the stabiliser-invariant part of \(\bigoplus_i A_{x_i}\), equivalently one cotangent slot per support point after collisions. Hence repeated indistinguishable outcomes do not keep independent occurrence cotangents. The finite-choice reverse rule is the FAD categorical rule with the log-weight components removed, followed by this quotienting of cotangents.

6.0.0.2 Exceptions

For exceptions, \(\Fun[nn]{ M_{E} }{ X } = X \amalg E\) may be taken as its own \(\opn{AST}\). In general, a handler on a carrier \(A\) is determined by a recovery map \(h:E\to A\), giving the algebra \([\mathrm{id}_A,h]:A\amalg E\to A\). To recover the monad from the CPS factorisation, it is enough to use the free algebra on carrier \(A=\Fun[nn]{ M_E }{ 2 }=2\amalg E\) as a canonical algebra: monad multiplication \(\Nat[nn]{ \mu }{ 2 }: \left( 2 \amalg E \right) \amalg E \to 2 \amalg E\). If \(E\) has continuous structure, a raised exception carries an exception-label cotangent until it is handled, and the handler uses the transposed derivative of the recovery map. Since here the monad and the \(\opn{AST}\) coincide, exceptions fit the same pattern without requiring a quotient construction.

6.0.0.3 Accumulation

For writer accumulation, let \(\Fun[nn]{ W_{M} }{ X } = X \times M\) be the writer monad associated to a monoid \(M\). For the basic case \(M = \R[n]\) we have a cancellable monoid, and again it may be taken as its own \(\opn{AST}\). Its algebras are right \(M\)-actions \(A\times M\to A\), and so for the \(\opn{AST}\)-algebra we use the carrier \(\R[n]\) with action \(\R[n] \times \R[n] \to \R[n]\), \(\tuple{ u, v } \mapsto u + v\). The lifted writer monad attaches to \(\tuple{x,r}\) the fibre \(A_x\oplus\R[n]\), and the handler rule is the transposed derivative of addition.

6.0.0.4 Remark (source syntax for the additional first-order effects).

For finite multiset non-determinism we add operations \(\opn[b]{choose}_k\) with permutation and commutative-monoid laws, with a fold handler for the chosen commutative monoid; operationally, a choice node can be implemented as a parallel fork and the handler as the corresponding join/reduction. For exceptions one may add \(\fn[n]{\opn[b]{raise}}{e}:\opn[b]{M}\tau\), with handlers determined by maps \(E\to\tau\). For the \(\R[n]\)-writer monad one may add \(\fn[n]{\opn[b]{tell}}{r}:\opn[b]{M}1\), together with the basic handler \(\opn[b]{M}\srcreal[n]\to\srcreal[n]\) induced by addition, which consumes the accumulated output. The target language then adds the corresponding linear terms for the transposed derivatives of operations and handlers; the generic \(\opn[b]{return}^{w}\) and \(\opn[b]{bind}^{w}\) rules are unchanged. These fork–join and accumulation patterns are common in high-performance code that we want to differentiate.

7 Related work↩︎

This work connects probabilistic programming semantics, stochastic gradient estimation, and semantic accounts of AD. The continuation view of distributions has classical roots in Schwartz’s theory of distributions [13]. In probabilistic programming, Vákár, Kammar, and Staton construct a probabilistic powerdomain by factoring integration into a continuation monad [14], while Kammar and McDermott develop the factorisation-systems technology for monadic lifting and logical relations [11].

Stochastic gradient estimation is central in machine learning; [15] surveys score, pathwise, and measure-valued estimators. Stochastic computation graphs showed how these compose by AD-like backpropagation [16], and Storchastic allows different estimators at different stochastic nodes [17]. ADEV is especially close, extending forward-mode AD to expectations of probabilistic programs with discrete and continuous distributions [5]. We contribute the corresponding reverse-mode account: reverse mode exposes cotangent-flow and fibredness issues that do not arise in forward mode. In particular, continuous probability is not merely the discrete construction with sums replaced by integrals; the relevant continuation/image construction need not be fibred over primal outcomes, so we leave it to future work. Compared to ADEV, we also use image factorisation to carve out distributions and their reverse-derivative structure from the continuation monad, and formulate expectation as an algebraic-effect handler [18], [19].

There is also growing work on deterministic reverse AD with algebraic effects and handlers. De Vilhena and Pottier verify a define-by-run reverse-mode AD library based on effect handlers [20], while Sigal studies effect-handler implementations of AD in Frank, OCaml, and related settings [21], [22]. Our focus is different: leaving implementation strategy open, with Efficient CHAD as a natural starting point [4], we give a denotationally correct source transformation for languages whose source terms themselves have algebraic effects. This builds on the CHAD paradigm for pure source languages [2], [23], [24] and its extension to iteration and partiality [3], showing that the technique also applies to impure, discrete-output algebraic effects and handlers.

Funding↩︎

This research was supported by the ERC project FoRECAST. The second author was also supported by the Centre for Mathematics of the University of Coimbra (CMUC, doi:10.54499/UID/00324/2025), funded by the Fundação para a Ciência e a Tecnologia (FCT) through grants UID/00324/2025 and UID/PRR/00324/2025.

8 Source language for CHAD↩︎

In this section, we present the complete type theory of the source language for CHAD, whose abbreviated introduction can be found in Section 3.1.

This is a first-order language whose term model can be understood as a free distributive category generated by the basic types \(\srcreal[k]\), \(k \in \mathbb{N}\), and by a family of chosen operations \(\mathsf{op}\), and equipped with a strong monad, a categorical constructor, and a specific Eilenberg–Moore algebra.

8.1 Type and term grammar↩︎

The type grammar for the source language is the one displayed in Equation 59 \[\label{eq:appendix-source-language-type-construction-grammar} \tau, \sigma, \rho ::= 0 \mid 1 \mid \tau \times \sigma \mid \tau \sqcup \sigma \mid \srcreal[k] \left( k \in \mathbb{N}\right) \mid \opn[b]{M} \tau\tag{59}\] with \(\srcreal\) being used as an abbreviation for \(\srcreal[1]\).

Similarly, the term construction grammar is the one in Equation 60 , denoting variables by \(x, y, z, \dots\) and contexts as \(\Gamma = x_{1} : \tau_{1}, \cdots, x_{n} : \tau_{n}\) \[\label{eq:appendix-source-language-term-construction-grammar} \begin{align} t, s, r ::= & \left. x \mid \tuple{\;} \mid \tuple{t, s} \mid \opn[b]{pr}_{1} t \mid \opn[b]{pr}_{2} t \mid \opn[b]{inl} t \mid \opn[b]{inr} t \right. \\ & \left. \mid \mathsf{case}\;t\;\mathsf{of}\;\mathsf{inl}(x)\Rightarrow s\;|\;\mathsf{inr}(y)\Rightarrow r \right. \left. \mid \fn[n]{ \opn[b]{abort} }{ t } \right. \\ & \left. \mid \mathsf{let}\;x=t\;\mathsf{in}\;s \right. \left. \mid \fn[n]{ \mathsf{op}}{ t_{1}, \dots, t_{k} } \right. \left. \mid \fn[n]{ \opn[b]{return} }{ t } \right. \\ & \left. \mid \fn[n]{ \opn[b]{bind} }{ t; x . s } \right. \left. \mid \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } \right. \left. \mid \fn[n]{ \opn[b]{E}_{\nu} }{ t } \right. \end{align}\tag{60}\] where \(\nu = \srcreal[p]\), for some \(p \in \mathbb{N}\), is the fixed type on which the handler algebra \(\opn[b]{E}_{\nu}\) is defined, and \(\mathsf{op}\) ranges over a family of chosen built-in operations, each with a signature \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\), representing the basic operations available in the source language. These operations should be thought of as deterministic operations, and are chosen according to the intended usage of the language. In the \(\opn[b]{categorical}\) constructor, the second components \(w_i\) are real-valued log-weights; the corresponding nonnegative weights are produced by the interpretation of \(\opn[b]{categorical}\).

8.2 Typing rules↩︎

The typing rules for the language are as follows.

8.2.1 Nullary product and coproduct↩︎

The elimination rule for nullary sum can be seen as a type-theoretic instance of the explosion principle, whereas the introduction rule for the nullary product states the existence of an element of \(1\) in any context, \[\prftree[r]{\( 0 \text{E} \)}{\Gamma \vdash t : 0}{\Gamma \vdash \fn[n]{ \opn[b]{abort} }{ t } : \rho} \qquad \prftree[r]{\( 1 \text{I} \)}{\;}{\Gamma \vdash \tuple{\;} : 1}\]

8.2.2 Binary products↩︎

The introduction and elimination rules of binary products behave similarly to their set-based semantics, cartesian products. \[\prftree[r]{\( \times \text{E1} \)}{ \Gamma \vdash u: \tau \times \sigma }{ \Gamma \vdash \opn[b]{pr}_{1} u : \tau } \qquad \prftree[r]{\( \times \text{E2} \)}{ \Gamma \vdash u: \tau \times \sigma }{ \Gamma \vdash \opn[b]{pr}_{2} u : \sigma } \qquad \prftree[r]{\( \times \text{I} \)}{\Gamma \vdash t: \tau}{\Gamma \vdash s: \sigma}{ \Gamma \vdash \tuple{t, s} : \tau \times \sigma }\]

8.2.3 Binary sums↩︎

As with binary products, binary sums typing rules are based on their intended set-based semantics, the disjoint union. The elimination rule is \[\prftree[r]{\( \sqcup \text{E} \)}{ \Gamma \vdash t: \tau \sqcup \sigma }{ \Gamma, x: \tau \vdash u : \rho }{ \Gamma, y: \sigma \vdash v : \rho }{ \Gamma \vdash \mathsf{case}\;t\;\mathsf{of}\;\mathsf{inl}(x)\Rightarrow u\;|\;\mathsf{inr}(y)\Rightarrow v: \rho }\] and the introduction rules are \[\prftree[r]{\( \sqcup \text{I1} \)}{ \Gamma \vdash t: \tau }{ \Gamma \vdash \opn[b]{inl} t : \tau \sqcup \sigma } \qquad \prftree[r]{\( \sqcup \text{I2} \)}{ \Gamma \vdash s: \sigma }{ \Gamma \vdash \opn[b]{inr} s : \tau \sqcup \sigma }\]

8.2.4 Operations↩︎

Any built-in operation \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\) defines a typing rule \[\prftree[r]{\mathrm{\small Operations}}{ \Gamma \vdash t_{j}: \srcreal[n_{j}] \quad \left( j \in \left\{ 1, \dots, k \right\} \right) }{ \Gamma \vdash \fn[n]{ \mathsf{op}}{ t_{1}, \dots, t_{k} } : \srcreal[m] }\]

8.2.5 Structural rules↩︎

We add a structural rule for the let-binding, \[\prftree[r]{\mathrm{\small Let introduction}}{ \Gamma \vdash t: \sigma }{ \Gamma, x : \sigma \vdash s : \rho }{ \Gamma \vdash \mathsf{let}\;x=t\;\mathsf{in}\;s : \rho }\] and another one for assumptions, meaning that variables already in the context may be taken as terms, \[\prftree[r]{\mathrm{\small Assumption }}{ \left( x : \tau \right) \in \Gamma }{ \Gamma \vdash x : \tau }\]

8.2.6 Syntactic sugar and pattern-matching↩︎

The \(n\)-ary product and coproduct notation used below is syntactic sugar for the binary grammar. We take both products and coproducts to be right-associated: \[\tau_1 \times \cdots \times \tau_n := \begin{cases} 1, & n=0,\\ \tau_1, & n=1,\\ \tau_1 \times (\tau_2 \times \cdots \times \tau_n), & n\geq 2, \end{cases}\] and similarly \[\tau_1 \sqcup \cdots \sqcup \tau_n := \begin{cases} 0, & n=0,\\ \tau_1, & n=1,\\ \tau_1 \sqcup (\tau_2 \sqcup \cdots \sqcup \tau_n), & n\geq 2. \end{cases}\]

The corresponding term notation is defined recursively by \[\tuple{t_1,\dots,t_n} := \begin{cases} \tuple{\;}, & n=0,\\ t_1, & n=1,\\ \tuple{t_1,\tuple{t_2,\dots,t_n}}, & n\geq 2, \end{cases}\] and by \[\opn[b]{in}^{n}_{1}t:=\opn[b]{inl}t \quad (n\geq 2), \qquad \opn[b]{in}^{n}_{i}t := \opn[b]{inr}(\opn[b]{in}^{n-1}_{i-1}t) \quad (2\leq i\leq n),\] with \(\opn[b]{in}^{1}_{1}t:=t\). We write \(\opn[b]{in}_{i}t\) when \(n\) is clear.

The \(n\)-ary coproduct case expression is defined by nested binary cases: \[\casenst{t}{i}{1}{x_i}{s_i} := \mathsf{let}\;x_1=t\;\mathsf{in}\;s_1,\] and, for \(n\geq 2\), \[\casenst{t}{i}{n}{x_i}{s_i} := \mathsf{case}\;t\;\mathsf{of}\;\mathsf{inl}(x_1)\Rightarrow s_1\;|\;\mathsf{inr}(y)\Rightarrow \casenst{y}{j}{n-1}{x_{j+1}}{s_{j+1}} ,\] where \(y\) is fresh.

Similarly, \(n\)-ary product pattern matching is defined by nested projections: \[\caseprod{t}{x_1}{s} := \mathsf{let}\;x_1=t\;\mathsf{in}\;s,\] and, for \(n\geq 2\), \[\caseprod{t}{x_1,\dots,x_n}{s} := \mathsf{let}\;z=t\;\mathsf{in}\; \mathsf{let}\;x_1=\opn[b]{pr}_{1}z\;\mathsf{in}\; \caseprod{\opn[b]{pr}_{2}z}{x_2,\dots,x_n}{s} ,\] where \(z\) is fresh. With these definitions, the displayed \(n\)-ary product and coproduct typing rules are admissible.

8.2.7 Monadic types↩︎

For monadic types, we add typing rules for the \(\opn[b]{return}\) and \(\opn[b]{bind}\) term constructions. Together with their equational rules, this ensures that they behave as the unit and Kleisli extension of a strong monad. \[\prftree[r]{\( \opn[b]{M} \text{E} \)}{ \Gamma \vdash t : \opn[b]{M} \sigma }{ \Gamma, x : \sigma \vdash s : \opn[b]{M} \tau }{ \Gamma \vdash \fn[n]{ \opn[b]{bind} }{ t; x . s } : \opn[b]{M} \tau } \qquad \prftree[r]{\( \opn[b]{M} \text{I} \)}{ \Gamma \vdash t : \tau }{ \Gamma \vdash \fn[n]{ \opn[b]{return} }{ t }: \opn[b]{M} \tau }\]

8.2.8 Measure-specific rules↩︎

All the rules above would remain the same regardless of which monadic effect one wishes to implement in the source language. The ones that follow are specific to the FAD monad and describe the construction of terms using \(\opn[b]{categorical}\) with real-valued log-weights and the handler algebra \(\opn[b]{E}_{\nu}\), \[\prftree[r]{\mathrm{\small Categ. }}{ \Gamma \vdash t_{i}: \tau \quad \left( i \in \left\{ 1, \dots, n \right\} \right) }{ \Gamma \vdash w_{i}: \srcreal \quad \left( i \in \left\{ 1, \dots, n \right\} \right) }{ \Gamma \vdash \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{t_{i}, w_{i}} }_{i = 1}^{n} } : \opn[b]{M} \tau } \quad \prftree[r]{\mathrm{\small Handler }}{ \Gamma \vdash t : \opn[b]{M} \nu }{ \Gamma \vdash \fn[n]{ \opn[b]{E}_{\nu} }{ t } : \nu }\] Here the terms \(w_i\) are log-weights rather than weights themselves; in the \(\Set\)-interpretation below, \(\opn[b]{categorical}\) exponentiates them, so no positivity premise is needed at the source type level.

8.3 Equational theory↩︎

The pure fragment has the usual \(\beta\)- and \(\eta\)-laws for \(1\), \(\times\), and \(\sqcup\), which are omitted (see Section II.2 of [25]). For the monadic and the handler fragments, we add the following rules. No additional definitional equations are imposed on the \(\opn[b]{categorical}\) constructors in the FAD source language treated here.

8.3.1 Kleisli laws↩︎

We consider the usual equational theory of strong monads. In what follows, \(\equiv\) denotes definitional equality, \[\label{eq:appendix-source-language-kleisli-laws} \begin{gather} \fn[n]{ \opn[b]{bind} }{ \fn[n]{ \opn[b]{return} }{ t }; x \; . \; s } \equiv s[t/x] \qquad \fn[n]{ \opn[b]{bind} }{ t; x \; . \; \fn[n]{ \opn[b]{return} }{ x } } \equiv t \\ \fn[n]{ \opn[b]{bind} }{ \fn[n]{ \opn[b]{bind} }{ t; x \; . \; s }; y \; . \; r } \equiv \fn[n]{ \opn[b]{bind} }{ t; x \; . \; \fn[n]{ \opn[b]{bind} }{ s ; y \; . \; r } } \end{gather}\tag{61}\] where we use \(s[t/x]\) for the capture-avoiding substitution of the term \(t\) for the variable \(x\) in the expression \(s\).

8.3.2 Eilenberg–Moore algebra laws↩︎

We also add laws ensuring that the handler algebra \(\opn[b]{E}_{\nu}\) acts as an Eilenberg–Moore algebra. \[\label{eq:appendix-source-language-eilenberg-moore-algebra-laws} \fn[n]{ \opn[b]{E}_{\nu} }{ \fn[n]{ \opn[b]{return} }{ v } } \equiv v \qquad \fn[n]{ \opn[b]{E}_{\nu} }{ \fn[n]{ \opn[b]{bind} }{ t; x \; . \; u } } \equiv \fn[n]{ \opn[b]{E}_{\nu} }{ \fn[n]{ \opn[b]{bind} }{ t; x \; . \; \fn[n]{ \opn[b]{return} }{ \fn[n]{ \opn[b]{E}_{\nu} }{ u } } } }\tag{62}\]

8.4 Denotational semantics↩︎

We denote by \(\mathsf{Syn}\) the term model of the source language, i.e., the syntactic category of the source language. Universally, it is the freely generated distributive category on the objects \(\srcreal[k]\), \(k \in \mathbb{N}\), with operations \(\mathsf{op}\) ranging in the family of built-in operations, equipped with a strong monad, a specific algebra on \(\srcreal[p]\) for some \(p \in \mathbb{N}\), and a categorical operation. This means that any category with this structure admits a structure-preserving functor from \(\mathsf{Syn}\), as is spelled out in Lemma 8.

In particular, we may interpret \(\mathsf{Syn}\) in \(\Set\), thus providing set-based semantics to the source language with a functor \(\llbracket \cdot\rrbracket_{\Set}: \mathsf{Syn}\to \Set\). This semantics is defined recursively on the term and type structures, and is spelled out as follows.

8.4.1 Types↩︎

\[\begin{align} \llbracket 0\rrbracket_{\Set} &= \emptyset \\ \llbracket 1\rrbracket_{\Set} &= \left\{ \ast \right\} \\ \llbracket \srcreal[n]\rrbracket_{\Set} &= \R[n] \end{align} \qquad \begin{align} \llbracket \tau \times \sigma\rrbracket_{\Set} &= \llbracket \tau\rrbracket_{\Set} \times \llbracket \sigma\rrbracket_{\Set} \\ \llbracket \tau \sqcup \sigma\rrbracket_{\Set} &= \llbracket \tau\rrbracket_{\Set} \amalg \llbracket \sigma\rrbracket_{\Set} \\ \llbracket \opn[b]{M} \tau\rrbracket_{\Set} &= \Fun[nn]{ M }{ \llbracket \tau\rrbracket_{\Set} } \end{align}\]

8.4.2 Contexts and typing judgements↩︎

Using the interpretations of types, we may define interpretations of contexts and of typing judgements: a context \(\Gamma = x_{1} : \tau_{1}, \cdots, x_{n} : \tau_{n}\) is interpreted as \(\llbracket \Gamma\rrbracket_{\Set} = \prod_{i = 1}^{n} \llbracket \tau_{i}\rrbracket_{\Set}\), and a typing judgement \(\Gamma \vdash t : \tau\) as a function \(\llbracket \Gamma \vdash t : \tau\rrbracket_{\Set}: \llbracket \Gamma\rrbracket_{\Set} \to \llbracket \tau\rrbracket_{\Set}\), whose definition is provided inductively on the term structure.

8.4.3 Deterministic fragment↩︎

Variables, constants, and constructors of the deterministic fragment are interpreted as follows, for all \(\gamma = \tuple{\gamma_{1}, \dots, \gamma_{n}} \in \llbracket \Gamma\rrbracket_{\Set}\), \[\begin{align} \fn[n]{ \llbracket \Gamma \vdash x_{j} : \tau_{j}\rrbracket_{\Set} }{\tuple{\gamma_{1}, \dots, \gamma_{n}}} &= \gamma_{j} && \left( j \in \left\{ 1, \dots, n \right\} \right) \\ \fn[n]{ \llbracket \Gamma \vdash \tuple{\;} : 1\rrbracket_{\Set} }{\gamma} &= \ast \\ \fn[n]{ \llbracket \Gamma \vdash \opn[b]{pr}_{i} u : \tau_{i}\rrbracket_{\Set} }{\gamma} &= \fn[n]{ \pi_{i} }{ \fn[n]{ \llbracket \Gamma \vdash u : \tau_{1} \times \tau_{2}\rrbracket_{\Set} }{ \gamma } } && \left( i \in \left\{ 1, 2 \right\} \right) \\ \fn[n]{ \llbracket \Gamma \vdash \tuple{ t, s } : \tau \times \sigma\rrbracket_{\Set} }{\gamma} &= \tuple{ \fn[n]{ \llbracket \Gamma \vdash t : \tau \rrbracket_{\Set} }{ \gamma }, \fn[n]{ \llbracket \Gamma \vdash s : \sigma \rrbracket_{\Set} }{ \gamma } } \\ \fn[n]{ \llbracket \Gamma \vdash \opn[b]{inl} u : \tau \sqcup \sigma\rrbracket_{\Set} }{\gamma} &= \fn[n]{ \iota_{1} }{ \fn[n]{ \llbracket \Gamma \vdash u: \tau\rrbracket_{\Set} }{ \gamma } } \\ \fn[n]{ \llbracket \Gamma \vdash \opn[b]{inr} u : \tau \sqcup \sigma\rrbracket_{\Set} }{\gamma} &= \fn[n]{ \iota_{2} }{ \fn[n]{ \llbracket \Gamma \vdash u: \sigma\rrbracket_{\Set} }{ \gamma } } \\ \fn[n]{ \llbracket \Gamma \vdash \opn[b]{abort} t : \rho\rrbracket_{\Set} }{\gamma} &= \fn[n]{ \opn{abs}_{\llbracket \rho\rrbracket_{\Set}} }{ \fn[n]{ \llbracket \Gamma \vdash t : 0\rrbracket_{\Set} }{\gamma} } \\ \fn[n]{ \llbracket \Gamma \vdash \mathsf{let}\;x=t\;\mathsf{in}\;s : \rho\rrbracket_{\Set} }{ \gamma } &= \fn[n]{ \llbracket \Gamma, x : \sigma \vdash s : \rho\rrbracket_{\Set} }{ \gamma, \fn[n]{ \llbracket \Gamma \vdash t : \sigma\rrbracket_{\Set} }{ \gamma } } \end{align}\] where, for any set \(S\), \(\opn{abs}_{S}: \emptyset \to S\) is the unique function. Case statements are interpreted as branching functions, \[\begin{align} &\fn[n]{ \llbracket \Gamma \vdash \mathsf{case}\;p\;\mathsf{of}\;\mathsf{inl}(x)\Rightarrow u\;|\;\mathsf{inr}(y)\Rightarrow v : \rho\rrbracket_{\Set} }{ \gamma } \qquad \\ &= \begin{cases} \fn[n]{ \llbracket \Gamma, x : \tau \vdash u : \rho\rrbracket_{\Set} }{ \gamma, a } &\text{if } \fn[n]{ \llbracket \Gamma \vdash p : \tau \sqcup \sigma\rrbracket_{\Set} }{ \gamma } = \fn[n]{ \iota_{1} }{ a } \\ \fn[n]{ \llbracket \Gamma, y : \sigma \vdash v : \rho\rrbracket_{\Set} }{ \gamma, b } &\text{if } \fn[n]{ \llbracket \Gamma \vdash p : \tau \sqcup \sigma\rrbracket_{\Set} }{ \gamma } = \fn[n]{ \iota_{2} }{ b } \end{cases} \end{align}\] and the built-in operations are interpreted as their representing counterparts, that is, for any operation \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\), \[\fn[n]{ \llbracket \Gamma \vdash \fn[n]{ \mathsf{op}}{ t_{1}, \cdots, t_{k} } : \srcreal[m] \rrbracket_{\Set} }{ \gamma } = \fn[n]{ \llbracket \mathsf{op}\rrbracket_{\Set} }{ \fn[n]{ \llbracket \Gamma \vdash t_{1} : \srcreal[n_{1}]\rrbracket_{\Set} }{ \gamma }, \cdots, \fn[n]{ \llbracket \Gamma \vdash t_{k} : \srcreal[n_{k}]\rrbracket_{\Set} }{ \gamma } }\] where \(\llbracket \mathsf{op}\rrbracket_{\Set}\) is required to be a differentiable function, as we are interested in computing derivatives.

8.4.4 Monadic fragment↩︎

For the interpretation of monadic types and their associated terms, we assume that the category on which the semantics is to be interpreted has a strong monad \(\tuple{ M, \eta, \mu, \operatorname{st}}\), with strength components \(\operatorname{st}_{A,B}: A \times M B \to M(A \times B)\). In the case of \(\Set\), all monads are canonically strong. We interpret \(\opn[b]{return}\) and \(\opn[b]{bind}\) as \[\begin{align} \llbracket \Gamma \vdash \fn[n]{ \opn[b]{return} }{ t } : \opn[b]{M} \tau\rrbracket_{\Set} &= \Nat[nn]{ \eta }{ \llbracket \tau\rrbracket_{\Set} } \circ \llbracket \Gamma \vdash t : \tau\rrbracket_{\Set} \\ \llbracket \Gamma \vdash \fn[n]{ \opn[b]{bind} }{ t; x \; . \; s } : \opn[b]{M} \tau\rrbracket_{\Set} &= \Nat[nn]{ \mu }{ \llbracket \tau\rrbracket_{\Set} } \circ \Fun[np]{ M }{ \llbracket \Gamma, x : \sigma \vdash s : \opn[b]{M} \tau\rrbracket_{\Set} } \circ \Nat[nn]{ \operatorname{st}}{ \llbracket \Gamma\rrbracket_{\Set}, \llbracket \sigma\rrbracket_{\Set} } \circ \tupling{ \mathrm{id}_{\llbracket \Gamma\rrbracket_{\Set}}, \llbracket \Gamma \vdash t : \opn[b]{M} \sigma\rrbracket_{\Set} } \end{align}\] Notice that these formulas are also applicable for interpretations on other distributive categories with strong monads.

8.4.5 Measure-specific terms↩︎

Finally, we present the interpretation of terms constructed with \(\opn[b]{categorical}\) or with the handler algebra \(\opn[b]{E}_{\nu}\). For \(\opn[b]{categorical}\), the second components are interpreted as log-weights: \[\begin{align} &\fn[n]{ \llbracket \Gamma \vdash \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } : \opn[b]{M} \tau\rrbracket_{\Set} }{ \gamma } \\ &\qquad = \fn[n]{ \Nat[nn]{ \operatorname{categorical}}{ \llbracket \tau\rrbracket_{\Set} } }{ \tuple{ \tuple{ \fn[n]{ \llbracket \Gamma \vdash t_{i} : \tau\rrbracket_{\Set} }{ \gamma }, \fn[n]{ \llbracket \Gamma \vdash w_{i} : \srcreal\rrbracket_{\Set} }{ \gamma } } }_{i = 1}^{n} } \end{align}\] and \[\begin{align} \fn[n]{ \llbracket \Gamma \vdash \fn[n]{ \opn[b]{E}_{\nu} }{ t } : \nu \rrbracket_{\Set}}{ \gamma } &= \fn[n]{ \opn{alg}^{\Set}_{\nu} }{ \fn[n]{ \llbracket \Gamma \vdash t : \opn[b]{M} \nu\rrbracket_{\Set} }{ \gamma } } \\ &= \sum_{ v \in \opn{supp}\! \left( \fn[n]{ \llbracket \Gamma \vdash t : \opn[b]{M} \nu\rrbracket_{\Set} }{ \gamma } \right) } \fn[n]{ \fn[n]{ \llbracket \Gamma \vdash t : \opn[b]{M} \nu\rrbracket_{\Set} }{ \gamma } }{ v } \cdot v \end{align}\]

8.5 Soundness and initiality↩︎

As claimed, \(\mathsf{Syn}\) can be seen as a free distributive category with a chosen strong monad on the objects \(\srcreal[k]\), \(k \in \mathbb{N}\), and the built-in family of operations \(\mathsf{op}\), as well as \(\opn[b]{categorical}\). As a freely generated category, it has a universal property, described in Lemma 8 (restatement of Lemma 6), which describes the existence of unique interpretation functors. These interpretations are also sound, meaning that any equation from the equational theory holds semantically, as stated in Lemma 7 (restatement of Lemma 5).

The proof follows by routine verification on the ways to obtain a definitional equality \(t \equiv s\) in the source language. The soundness of the (omitted) pure fragment laws follows from the distributive category laws; there are no additional \(\opn[b]{categorical}\)-equations to check.

On the other hand, the soundness of term equalities with \(\opn[b]{return}\) and \(\opn[b]{bind}\) follows from the interpretation of \(\opn[b]{bind}\) with the strength \(\operatorname{st}\) of the monad \(M\), together with the Kleisli equations and coherence laws of strong monads, while the definitional equalities for the handler algebra follow from the Eilenberg–Moore algebra axioms.

The universal property of the syntactic category can be given explicitly as follows.

The proof follows by standard algebraic presentation arguments, using the formulas for interpretation above, with \(\llbracket \srcreal[n]\rrbracket = X_{n}\) for all \(n \in \mathbb{N}\), and induction over the structure of the morphisms in \(\mathsf{Syn}\). Uniqueness follows from the fact that the equations in \(\mathsf{Syn}\) are precisely those required to enforce its freely generated structure.

9 Target language for CHAD↩︎

In this section we present the complete type-theoretic description of the target language for CHAD.

As explained in the abridged version in Section 3.3, the core idea of CHAD is to encode a code transformation as a categorical transformation between syntactic categories. As such code written in the source language (see Section 8) is systematically converted by induction on its structural components, giving their differentiated versions. The target language is the language in which these transformed codes are written.

The target language is a dependently typed functional programming language, and it is split into a cartesian type theory and, for every context of cartesian types (i.e. a cartesian context), an associated linear type theory.

Denoting by \(\CSyn\) the term model of the cartesian fragment, each cartesian context \(\Gamma\) (viewed as an object of \(\CSyn\), up to variable renaming) determines a syntactic category \(\LSyn[\Gamma]\) for the associated linear type theory. Any substitution \(\Gamma \to \Gamma'\) induces a reindexing functor \(\LSyn[\Gamma'] \to \LSyn[\Gamma]\) between these linear term models. Thus \(\LSyn\) is a strict indexed category \(\CSyn[o] \to \Cat\).

The syntactic category of the target language is then given by the Grothendieck construction \(\mathsf{Tgt}\) (see Section 2.3 of [2]).

As in the source language, we write \(\tau, \sigma, \rho\) for the cartesian types, and reserve underlined Greek letters \(\underline{\tau}, \underline{\sigma}, \underline{\rho}\) for linear types. We then have cartesian typing judgements of the form \(\Gamma \vdash t : \tau\) and linear typing judgements of the form \(\Gamma ; v : \underline{\sigma} \vdash s : \underline{\rho}\), with exactly one distinguished linear variable, whose name may be \(\alpha\)-renamed.

We can divide the description of the target language into three layers:

• Non-monadic core: deterministic operations. The description of this part is independent of any monad-related components.

• Oplax strong-monad core: liftings possible for any strong monad in the source category.

• Distribution and handler: language features that are specific to the FAD monad.

By breaking the target language into these layers, we get a generic procedure for introducing (strong) monadic effects into the automatic differentiation framework of CHAD.

9.1 Type construction grammar↩︎

Cartesian types are described by the following grammar: \[\label{eq:appendix-target-language-cartesian-type-construction-grammar} \tau, \sigma, \rho ::= \left( \text{As in Equation~\eqref{eq:source-language-type-construction-grammar}} \right) \mid \underline{\sigma} \multimap\underline{\rho} \mid \Sigma \left( x : \sigma \right) . \tau'\tag{63}\] meaning that, in the cartesian fragment, we also have function types between linear types and dependent sums. For linear types, we use the grammar \[\label{eq:appendix-target-language-linear-type-construction-grammar} \underline{\tau}, \underline{\sigma}, \underline{\rho} ::= \underline{\srcreal[k]} \; \left( k \in \mathbb{N}\right) \mid \underline{1} \mid \underline{\tau} \times \underline{\sigma} \mid \opn[b]{case} \; t \; \opn[b]{of} \; \left\{ \opn[b]{in}_{i} x_{i} \mapsto \underline{\sigma}_{i} \right\}_{i = 1}^{n} \mid \opn[b]{M}^{w}_{p : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ x \mapsto \underline{ \sigma } \right]\tag{64}\] meaning that the base Euclidean-space types have associated linear types (in the sense that the cotangent space of a Euclidean space at any point can be identified with the base space itself), as well as the fibre-wise zero object \(1\). Moreover, we allow fibre-wise binary biproducts \(\times\) and families of linear types indexed by finite disjunctions, and we introduce the cotangent linear type associated with monadic types.

The goal of highlighting the different layers is to make explicit which constructions are general for the implementation of algebraic monadic effects and which ones are specific to the FAD monad.

9.2 Term construction grammar↩︎

As with types, we also distinguish between cartesian and linear terms. The cartesian terms are described by the grammar \[\label{eq:appendix-target-language-cartesian-term-construction-grammar} t, s, r ::= \left( \text{As in Equation~\eqref{eq:appendix-source-language-term-construction-grammar}} \right) \mid \tuple{ a, t } \mid \opn[b]{case} \; p \; \opn[b]{of} \; \tuple{ x, y } \mapsto s \mid \lambda v \; . \; s\tag{65}\] where the constructor \(\tuple{ a, t }\) and the eliminator \(\opn[b]{case} \; p \; \opn[b]{of} \; \tuple{ x, y } \mapsto s\) are used in the typing rules for the cartesian \(\Sigma\)-type, and \(\lambda v \; . \; s\) is used for abstraction over the unique linear variable \(v\).

Since all term constructions from the source language (Equation 60 ) are present in the cartesian fragment of the target language, with the same typing rules and equational theory, the source syntactic category \(\mathsf{Syn}\) can be seen as a full subcategory of \(\CSyn\).

For linear terms, we use the following grammar: \[\label{eq:appendix-target-language-linear-term-construction-grammar} \begin{align} t, s, r ::= v &\mid \mathsf{let}\;v=t\;\mathsf{in}\;s \mid \tuple{ t, s } \mid \opn[b]{pr}_{i} t \mid \opn[b]{0} \mid t + s \mid \linapp{r}{t} \\ & \mid \fn[n]{ \fn[n]{ \mathsf{op}}{ t_{1}, \cdots, t_{k} } }{v} \mid \fn[n]{ \opn[b]{return}^{w}_{t} }{s} \mid \fn[n]{ \opn[b]{bind}^{w}_{s; x \; . \; t} }{r; u} \\ & \mid \fn[n]{ \opn[b]{categorical}^{w}_{\tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n}} }{ \tuple{ \tuple{ t_{i}', w_{i}' } }_{i = 1}^{n} } \mid \fn[n]{ \opn[b]{E}^{w}_{\nu, t} }{ s } \end{align}\tag{66}\] The three lines correspond to the three layers: the first line consists of the non-monadic core operations (linear variables, let bindings, pairings and projections, the zero vector, sums, and \(\linapp{r}{t}\) for linear application with \(r\) a cartesian term of type \(\underline{\sigma} \multimap\underline{\rho}\)). The second line has the derivatives of the built-in operations \(\mathsf{op}\) of the source language, as well as the linear fibre components of the monad operations \(\opn[b]{return}\) and \(\opn[b]{bind}\). Finally, the third line has the linear fibre components of \(\opn[b]{categorical}\) and of the handler algebra \(\opn[b]{E}_{\nu}\).

9.3 Typing rules↩︎

9.3.1 Non-monadic core↩︎

The cartesian typing rules are the same as the ones of the source language (see Section 8.2) for the common term constructors. For the dependent sums, we have the rules \[\prftree{ \Gamma \vdash a : \sigma }{ \Gamma \vdash t : \tau[a/x] }{ \Gamma \vdash \tuple{ a, t }: \Sigma \left( x : \sigma \right) . \tau } \qquad \prftree{ \Gamma \vdash \rho \text{ type} }{ \Gamma \vdash p : \Sigma \left( x : \sigma \right) . \tau }{ \Gamma, x : \sigma, y : \tau \vdash s : \rho }{ \Gamma \vdash \opn[b]{case} \; p \; \opn[b]{of} \; \tuple{ x, y } \to s : \rho }\] and for the linear function types, we introduce the typing rules \[\prftree{ \Gamma; v : \underline{\sigma} \vdash s : \underline{\rho} }{ \Gamma \vdash \lambda v \; . \; s : \underline{\sigma} \multimap\underline{\rho} } \qquad \prftree{ \Gamma \vdash r : \underline{\sigma} \multimap\underline{\rho} }{ \Gamma; v : \underline{\tau} \vdash t : \underline{\sigma} }{ \Gamma; v : \underline{\tau} \vdash \linapp{r}{t} : \underline{\rho} }\]

There are also non-monadic-layer typing rules for linear terms.

9.3.1.1 Linear variable

\[\prftree{}{ \Gamma; v : \underline{\sigma} \vdash v : \underline{\sigma} } \qquad\qquad \textit{(one distinguished linear variable)}\]

9.3.1.2 Linear let

\[\prftree{ \Gamma; v : \underline{\tau} \vdash t : \underline{\sigma} }{ \Gamma; v : \underline{\sigma} \vdash s : \underline{\rho} }{ \Gamma; v : \underline{\tau} \vdash \mathsf{let}\;v=t\;\mathsf{in}\;s : \underline{\rho} }\]

9.3.1.3 Linear biproduct

\[\begin{gather} \prftree{}{ \Gamma ; v : \underline{\rho} \vdash \opn[b]{0} : \underline{\sigma} } \qquad \prftree{ \Gamma; v : \underline{\rho} \vdash s : \underline{\sigma} }{ \Gamma; v : \underline{\rho} \vdash t : \underline{\sigma} }{ \Gamma; v : \underline{\rho} \vdash s + t : \underline{\sigma} } \\ \prftree{ \Gamma; v : \underline{\rho} \vdash s : \underline{\tau} }{ \Gamma; v : \underline{\rho} \vdash t : \underline{\sigma} }{ \Gamma; v : \underline{\rho} \vdash \tuple{ s, t } : \underline{\tau} \times \underline{\sigma} } \qquad \prftree{ \Gamma; v : \underline{\rho} \vdash u : \underline{\tau}_{1} \times \underline{\tau}_{2} }{ \Gamma; v : \underline{\rho} \vdash \opn[b]{pr}_{i} u : \underline{\tau}_{i} } \end{gather}\] where \(i \in \left\{ 1, 2 \right\}\).

9.3.1.4 Linear case over cartesian sums

\[\prftree{ \Gamma \vdash t : \tau_{1} \sqcup \cdots \sqcup \tau_{n} }{ \left\{ \Gamma, x_{i} : \tau_{i}; v : \underline{\sigma}_{i} \vdash r_{i} : \underline{\rho}_{i} \right\}_{i=1}^{n} } { \Gamma; v: \casenst{t}{i}{n}{x_{i}}{\underline{\sigma}_{i}} \vdash \casenst{t}{i}{n}{x_{i}}{r_{i}} : \casenst{t}{i}{n}{x_{i}}{\underline{\rho}_{i}} }\]

9.3.1.5 Built-in operations

For any built-in operation \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\) in the source language, we associate a term constructor \(\opn[b]{Dop}\), whose intended interpretation is as the transposed Jacobian of \(\mathsf{op}\). It has the following associated typing rule. \[\prftree{ \left\{ \Gamma \vdash t_{j} : \srcreal[n_{j}] \right\}_{j = 1}^{k} }{ \Gamma ; v : \underline{\srcreal[m]} \vdash \fn[n]{ \fn[n]{ \opn[b]{Dop} }{ t_{1}, \cdots, t_{k} } }{ v } : \underline{\srcreal[n_{1}]} \times \cdots \times \underline{\srcreal[n_{k}]} }\]

9.3.2 Oplax strong-monad core↩︎

The following typing rules concern a generic strong monad, together with its chosen cotangent lift. In every rule below, all displayed cartesian contexts are assumed to be well formed. A judgement such as \(\Gamma, x : \tau \vdash \underline{\sigma} \text{ type}\) means that \(\underline{\sigma}\) is a linear type in the fibre \(\LSyn[\Gamma,x:\tau]\); a judgement \(\Gamma \vdash \underline{\rho} \text{ type}\) means that \(\underline{\rho}\) is a linear type in \(\LSyn[\Gamma]\). The linear variables occurring in the linear arguments of \(\opn[b]{return}^{w}\), \(\opn[b]{bind}^{w}\), \(\opn[b]{categorical}^{w}\), and \(\opn[b]{E}_{\nu}^{w}\) are local binding occurrences; following the single-linear-variable convention, we write each of them as \(v\), with scope determined by the surrounding judgement and with the usual capture-avoiding substitution convention.

9.3.2.1 Type formation

\[\prftree{ \Gamma \vdash p : \Fun[nn]{ \opn[b]{M} }{ \tau } }{ \Gamma,x:\tau \vdash \underline{\sigma} \text{ type} }{ \Gamma \vdash \opn[b]{M}^{w}_{p : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ x \mapsto \underline{\sigma} \right] \text{ type} }\] Here \(x\) is the returned value on which the continuation cotangent type may depend.

9.3.2.2 \(w\)-\(\opn[b]{return}\)

\[\prftree{ \Gamma \vdash t: \tau }{ \Gamma, x : \tau \vdash \underline{\sigma} \text{ type} }{ \Gamma \vdash \underline{\rho} \text{ type} }{ \Gamma; v : \underline{\sigma} [t/x] \vdash s : \underline{\rho} }{ \Gamma; v : \opn[b]{M}^{w}_{\fn[n]{ \opn[b]{return} }{ t } : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ x\mapsto \underline{\sigma} \right] \vdash \fn[n]{ \opn[b]{return}^{w}_{t} }{ s }: \underline{\rho} }\] Thus the argument \(s\) is typed in the returned-value fibre \(\underline{\sigma}[t/x]\), while the result consumes a cotangent of \(\fn[n]{\opn[b]{return}}{t}\) whose continuation fibre is \(x\mapsto\underline{\sigma}\).

9.3.2.3 \(w\)-\(\opn[b]{bind}\)

\[\begin{gather} \Gamma \vdash s : \Fun[nn]{ \opn[b]{M} }{ \sigma } \qquad \Gamma, x : \sigma \vdash t : \Fun[nn]{ \opn[b]{M} }{ \tau } \\ \Gamma, x : \sigma \vdash \underline{\sigma} \text{ type} \qquad \Gamma, y : \tau \vdash \underline{\tau} \text{ type} \qquad \Gamma \vdash \underline{\rho} \text{ type} \\ \prftree{ \Gamma; v : \opn[b]{M}^{w}_{s : \Fun[nn]{ \opn[b]{M} }{ \sigma }} \left[ x\mapsto \underline{\sigma} \right] \vdash r : \underline{\rho} }{ \Gamma, x : \sigma ; v : \opn[b]{M}^{w}_{t : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ y \mapsto \underline{\tau} \right] \vdash u : \underline{\rho} \times \underline{\sigma} }{ \Gamma; v : \opn[b]{M}^{w}_{\fn[n]{ \opn[b]{bind} }{ s; x \; . \; t } : \Fun[nn]{ \opn[b]{M} }{ \tau }} \left[ y \mapsto \underline{\tau} \right] \vdash \fn[n]{ \opn[b]{bind}^{w}_{s; x \; . \; t} }{ r; u } : \underline{\rho} } \end{gather}\] Here \(\underline{\sigma}\) is the continuation-cotangent type for the first computation, depending on its returned value \(x:\sigma\), and \(\underline{\tau}\) is the continuation-cotangent type for the composite computation, depending on its final returned value \(y:\tau\). The second linear argument \(u\) returns both a contribution to the external cotangent \(\underline{\rho}\) and a cotangent in the fibre \(\underline{\sigma}\) of the bound value \(x\).

9.3.3 Measure-specific layer↩︎

These typing rules relate specifically to the FAD monad, providing rules for the transposed derivative of \(\opn[b]{categorical}\) and for the handler algebra \(\opn[b]{E}_{\nu}\).

9.3.3.1 \(w\)-\(\opn[b]{categorical}\)

For all \(n \in \mathbb{N}\), the finite atomic distribution constructor has the following cotangent rule. \[\begin{gather} \Gamma, x : \tau \vdash \underline{\sigma} \text{ type} \qquad \Gamma \vdash \underline{\rho} \text{ type} \\ \prftree{ \left\{ \Gamma \vdash t_{i} : \tau \right\}_{i = 1}^{n} }{ \left\{ \Gamma \vdash w_{i} : \srcreal \right\}_{i = 1}^{n} }{ \left\{ \Gamma; v : \underline{\sigma}[t_{i}/x] \vdash t_{i}' : \underline{\rho} \right\}_{i = 1}^{n} }{ \left\{ \Gamma; v : \underline{\srcreal} \vdash w_{i}' : \underline{\rho} \right\}_{i = 1}^{n} }{ \Gamma; v : \opn[b]{M}^{w}_{ \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } : \Fun[nn]{ \opn[b]{M} }{ \tau } } \left[ x \mapsto \underline{\sigma} \right] \vdash \fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{ \tuple{ t_{i}, w_{i} } }_{i = 1}^{n} } }{ \tuple{ \tuple{ t_{i}', w_{i}' } }_{i = 1}^{n} } : \underline{\rho} } \end{gather}\] Thus \(\underline{\sigma}\) is the cotangent family for the sampled value, while \(\underline{\rho}\) is the ambient cotangent type to which each atom and weight backpropagator contributes.

9.3.3.2 Transposed derivative of the handler

Since \(\nu = \srcreal[p]\) for some \(p \in \mathbb{N}\), write \(\underline{\tau}_{\nu} := \underline{\srcreal[p]}\) for its associated linear type. The constant family \(y:\nu\mapsto\underline{\tau}_{\nu}\) is a well-formed linear type in \(\LSyn[\Gamma,y:\nu]\). We impose the following typing rule. \[\prftree{ \Gamma \vdash t : \Fun[nn]{ \opn[b]{M} }{ \nu } }{ \Gamma \vdash \underline{\rho} \text{ type} }{ \Gamma; v : \opn[b]{M}^{w}_{t : \Fun[nn]{ \opn[b]{M} }{ \nu }} \left[ y \mapsto \underline{\tau}_{\nu} \right] \vdash s : \underline{\rho} }{ \Gamma; v : \underline{\tau}_{\nu} \vdash \fn[n]{ \opn[b]{E}^{w}_{\nu, t} }{ s } : \underline{\rho} }\]

9.4 Equational theory↩︎

In addition to the equations from the equational theory of the source language (Section 8.3), which also apply for the cartesian fragment, the following equations describe the equational theory of the target language. They are stated up to the usual capture-avoiding substitution on the bound variables.

9.4.0.1 Cartesian layer (\(\Sigma\), products, and sums)

\[\begin{array}{rrl} \left( \Sigma \beta \right) &\quad \caseprod{\tuple{ a, t }}{x, y}{s} &\equiv s \left[ x := a, y := t \right] \\ \left( \Sigma \eta \right) &\quad \caseprod{p}{x, y}{\tuple{ x, y }} &\equiv p \end{array}\]

9.4.0.2 Linear biproduct (fibrewise) and additivity

\[\begin{array}{rrl} \left( \oplus \beta \right) &\quad \opn[b]{pr}_{1} \tuple{ s, t } &\equiv s \\ &\quad \opn[b]{pr}_{2} \tuple{ s, t } &\equiv t \\ \left( \oplus \eta \right) &\quad \tuple{ \opn[b]{pr}_{1} t, \opn[b]{pr}_{2} t } &\equiv t \\ \left( \text{Add-unit} \right) &\quad s + \opn[b]{0} &\equiv s \\ &\quad \opn[b]{0} + s &\equiv s \\ \left( \text{Add-assoc/comm} \right) &\quad \left( r + s \right) + t &\equiv r + \left( s + t \right) \\ &\quad s + t &\equiv t + s \\ \left( \text{Add-lin} \right) &\quad r[\opn[b]{0}/v] &\equiv \opn[b]{0} \\ &\quad r[s + t/v] &\equiv r[s/v] + r[t/v] \end{array}\]

9.4.0.3 Linear structural

\[\begin{array}{rr} \left( \text{Lin. struct.} \right) &\quad \mathsf{let}\;v=t\;\mathsf{in}\;s \equiv s[t/v] \end{array}\]

9.4.0.4 Linear function space

\[\begin{array}{rrl} \left( \multimap\beta \right) &\quad \fn[p]{ \lambda v \; . \; s }{ t } &\equiv s[t/v] \\ \left( \multimap\eta \right) &\quad \lambda v \; . \; \fn[n]{ r }{ v } &\equiv r \qq{when} \left( v \not\in \fn[n]{ \opn{FV} }{ r } \right) \end{array}\]

9.4.0.5 Built-in operations

For every built-in operation \(\mathsf{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\), its associated transposed Jacobian term constructor \(\opn[b]{Dop}\) is additive in its linear argument and respects cartesian substitution.

9.4.0.6 Oplax fibred strong-monad and algebra equations

We impose the following equations whenever all displayed terms are well-typed. They are stated modulo the Kleisli equations and Eilenberg–Moore algebra equations in the cartesian fragment, so that the linear types on the two sides are identified by the corresponding cartesian equality. All occurrences of the distinguished linear variable in arguments of \[\fn[n]{\opn[b]{return}^{w}_{t}}{r}, \qquad \fn[n]{\opn[b]{bind}^{w}_{s;x.t}}{r;u}, \qquad \fn[n]{\opn[b]{E}^{w}_{\nu,t}}{r}\] are binding occurrences and are implicitly \(\alpha\)-renamed as necessary to avoid capture. The equations we impose are precisely those of an oplax indexed strong monad and an oplax indexed Eilenberg–Moore algebra. They are also the equations that are required to ensure that the interpretation of the source language in the target language is well-defined, i.e. the transposed derivatives of the Kleisli and Eilenberg–Moore algebra equations for the source language.

The \(w\)-left-unit law is the transpose of \[\fn[n]{\opn[b]{bind}}{ \fn[n]{\opn[b]{return}}{a}; x.t } \equiv t[a/x].\] Suppose \[\Gamma\vdash a:\sigma, \qquad \Gamma,x:\sigma\vdash t:\Fun[nn]{\opn[b]{M}}{\tau},\] \[\Gamma,x:\sigma\vdash \underline{\sigma}\;\mathrm{type}, \qquad \Gamma,y:\tau\vdash \underline{\tau}\;\mathrm{type}, \qquad \Gamma\vdash \underline{\rho}\;\mathrm{type},\] and \[\Gamma; v:\underline{\sigma}[a/x] \vdash r:\underline{\rho}, \qquad \Gamma,x:\sigma; v:\opn[b]{M}^{w}_{t:\Fun[nn]{\opn[b]{M}}{\tau}} \left[y\mapsto\underline{\tau}\right] \vdash u:\underline{\rho}\times\underline{\sigma}.\] Then \[\label{eq:appendix-target-language-monad-w-left-unit} \begin{align} & \fn[n]{ \opn[b]{bind}^{w}_{ \fn[n]{\opn[b]{return}}{a}; x.t } }{ \fn[n]{\opn[b]{return}^{w}_{a}}{r}; u } \\ &\qquad\equiv \mathsf{let}\;v=u[a/x]\;\mathsf{in}\; \opn[b]{pr}_{1}v + r[\opn[b]{pr}_{2}v/v] . \end{align}\tag{67}\]

The \(w\)-right-unit law is the transpose of \[\fn[n]{\opn[b]{bind}}{ s; x.\fn[n]{\opn[b]{return}}{x} } \equiv s.\] Suppose \[\Gamma\vdash s:\Fun[nn]{\opn[b]{M}}{\sigma}, \qquad \Gamma,x:\sigma\vdash \underline{\sigma}\;\mathrm{type}, \qquad \Gamma\vdash \underline{\rho}\;\mathrm{type},\] and \[\Gamma; v:\opn[b]{M}^{w}_{s:\Fun[nn]{\opn[b]{M}}{\sigma}} \left[x\mapsto\underline{\sigma}\right] \vdash r:\underline{\rho}.\] Then \[\label{eq:appendix-target-language-monad-w-right-unit} \fn[n]{ \opn[b]{bind}^{w}_{ s; x.\fn[n]{\opn[b]{return}}{x} } }{ r; \tuple{ \opn[b]{0}, \fn[n]{\opn[b]{return}^{w}_{x}}{v} } } \equiv r .\tag{68}\] In the second argument, the term \[\fn[n]{\opn[b]{return}^{w}_{x}}{v}\] is typed using the returned-value fibre \[y:\sigma\mapsto \underline{\sigma}[y/x].\] Thus it produces a term of type \(\underline{\sigma}\) in the context \(\Gamma,x:\sigma\).

The \(w\)-associativity law is the transpose of \[\fn[n]{\opn[b]{bind}}{ \fn[n]{\opn[b]{bind}}{s;x.t}; y.q } \equiv \fn[n]{\opn[b]{bind}}{ s; x.\fn[n]{\opn[b]{bind}}{t;y.q} }.\] Suppose \[\Gamma\vdash s:\Fun[nn]{\opn[b]{M}}{\sigma}, \qquad \Gamma,x:\sigma\vdash t:\Fun[nn]{\opn[b]{M}}{\tau}, \qquad \Gamma,y:\tau\vdash q:\Fun[nn]{\opn[b]{M}}{\upsilon},\] \[\Gamma,x:\sigma\vdash \underline{\sigma}\;\mathrm{type}, \qquad \Gamma,y:\tau\vdash \underline{\tau}\;\mathrm{type}, \qquad \Gamma,z:\upsilon\vdash \underline{\upsilon}\;\mathrm{type}, \qquad \Gamma\vdash \underline{\rho}\;\mathrm{type},\] and \[\Gamma; v:\opn[b]{M}^{w}_{s:\Fun[nn]{\opn[b]{M}}{\sigma}} \left[x\mapsto\underline{\sigma}\right] \vdash r:\underline{\rho},\] \[\Gamma,x:\sigma; v:\opn[b]{M}^{w}_{t:\Fun[nn]{\opn[b]{M}}{\tau}} \left[y\mapsto\underline{\tau}\right] \vdash u:\underline{\rho}\times\underline{\sigma},\] \[\Gamma,y:\tau; v:\opn[b]{M}^{w}_{q:\Fun[nn]{\opn[b]{M}}{\upsilon}} \left[z\mapsto\underline{\upsilon}\right] \vdash e:\underline{\rho}\times\underline{\tau}.\] Then \[\label{eq:appendix-target-language-monad-w-associativity} \begin{align} & \fn[n]{ \opn[b]{bind}^{w}_{ \fn[n]{\opn[b]{bind}}{s;x.t}; y.q } }{ \fn[n]{\opn[b]{bind}^{w}_{s;x.t}}{r;u}; e } \\ &\qquad\equiv \fn[n]{ \opn[b]{bind}^{w}_{ s; x.\fn[n]{\opn[b]{bind}}{t;y.q} } }{ r; \fn[n]{\opn[b]{bind}^{w}_{t;y.q}}{ u; \mathsf{let}\;v=e\;\mathsf{in}\; \tuple{ \tuple{\opn[b]{pr}_{1}v,\opn[b]{0}}, \opn[b]{pr}_{2}v } } } . \end{align}\tag{69}\] On the right-hand side, \(q\) and \(e\) are weakened from \(\Gamma,y:\tau\) to \(\Gamma,x:\sigma,y:\tau\), and the inserted zero has type \(\underline{\sigma}\).

Side note: the oplax comparison term for cartesian reindexing is definable. It is defined by the pure-continuation special case of \(\opn[b]{bind}^{w}\). If \[\Gamma\vdash p:\Fun[nn]{\opn[b]{M}}{\sigma}, \qquad \Gamma,x:\sigma\vdash h:\tau, \qquad \Gamma,y:\tau\vdash \underline{\tau}\;\mathrm{type},\] then \[\label{eq:appendix-target-language-monad-w-alpha} \begin{align} & \Gamma; v: \opn[b]{M}^{w}_{ \fn[n]{\opn[b]{bind}}{ p; x.\fn[n]{\opn[b]{return}}{h} } : \Fun[nn]{\opn[b]{M}}{\tau} } \left[y\mapsto\underline{\tau}\right] \vdash \alpha_{p,x.h,\underline{\tau}}(v) : \opn[b]{M}^{w}_{p:\Fun[nn]{\opn[b]{M}}{\sigma}} \left[ x\mapsto\underline{\tau}[h/y] \right], \\ & \alpha_{p,x.h,\underline{\tau}}(v) := \fn[n]{ \opn[b]{bind}^{w}_{ p; x.\fn[n]{\opn[b]{return}}{h} } }{ v; \tuple{ \opn[b]{0}, \fn[n]{\opn[b]{return}^{w}_{h}}{v} } }. \end{align}\tag{70}\] The identity and composition equations for these comparison maps follow from 68 and 69 , respectively. Thus the displayed equations present an oplax fibred monad over the source monad. Since they are schematic in an arbitrary cartesian context, the usual contextual strength is inherited from the given strength of the source monad.

The handler \(w\)-unit law is the transpose of \[\fn[n]{\opn[b]{E}_{\nu}}{ \fn[n]{\opn[b]{return}}{a} } \equiv a.\] Suppose \[\Gamma\vdash a:\nu, \qquad \Gamma\vdash \underline{\rho}\;\mathrm{type}, \qquad \Gamma;v:\underline{\tau}_{\nu}\vdash r:\underline{\rho}.\] Then \[\label{eq:appendix-target-language-handler-w-unit} \fn[n]{ \opn[b]{E}^{w}_{ \nu, \fn[n]{\opn[b]{return}}{a} } }{ \fn[n]{\opn[b]{return}^{w}_{a}}{r} } \equiv r .\tag{71}\]

The handler \(w\)-multiplication law is the transpose of \[\fn[n]{\opn[b]{E}_{\nu}}{ \fn[n]{\opn[b]{bind}}{t;x.u} } \equiv \fn[n]{\opn[b]{E}_{\nu}}{ \fn[n]{\opn[b]{bind}}{ t; x.\fn[n]{\opn[b]{return}}{ \fn[n]{\opn[b]{E}_{\nu}}{u} } } }.\] Suppose \[\Gamma\vdash t:\Fun[nn]{\opn[b]{M}}{\sigma}, \qquad \Gamma,x:\sigma\vdash u:\Fun[nn]{\opn[b]{M}}{\nu},\] \[\Gamma,x:\sigma\vdash \underline{\sigma}\;\mathrm{type}, \qquad \Gamma\vdash \underline{\rho}\;\mathrm{type},\] and \[\Gamma; v: \opn[b]{M}^{w}_{t:\Fun[nn]{\opn[b]{M}}{\sigma}} \left[x\mapsto\underline{\sigma}\right] \vdash r:\underline{\rho},\] \[\Gamma,x:\sigma; v: \opn[b]{M}^{w}_{u:\Fun[nn]{\opn[b]{M}}{\nu}} \left[y\mapsto\underline{\tau}_{\nu}\right] \vdash s:\underline{\rho}\times\underline{\sigma}.\] Then \[\label{eq:appendix-target-language-handler-w-multiplication} \begin{align} & \fn[n]{ \opn[b]{E}^{w}_{ \nu, \fn[n]{\opn[b]{bind}}{t;x.u} } }{ \fn[n]{\opn[b]{bind}^{w}_{t;x.u}}{ r; s } } \\ &\qquad\equiv \fn[n]{ \opn[b]{E}^{w}_{ \nu, \fn[n]{\opn[b]{bind}}{ t; x.\fn[n]{\opn[b]{return}}{ \fn[n]{\opn[b]{E}_{\nu}}{u} } } } }{ \fn[n]{ \opn[b]{bind}^{w}_{ t; x.\fn[n]{\opn[b]{return}}{ \fn[n]{\opn[b]{E}_{\nu}}{u} } } }{ r; \fn[n]{ \opn[b]{return}^{w}_{ \fn[n]{\opn[b]{E}_{\nu}}{u} } }{ \fn[n]{\opn[b]{E}^{w}_{\nu,u}}{s} } } } . \end{align}\tag{72}\]

9.5 Syntactic category of the target language↩︎

With the proper type theory spelled out, we may now provide the proper definition of \(\LSyn[\tau]\) for each cartesian type \(\tau\):

• Objects: linear types \(\underline{\sigma}\) well-formed in the context \(p: \tau\).

• Morphisms: a morphism \(\underline{\rho} \to \underline{\sigma}\) is a linear term \(p: \tau; v : \underline{\rho} \vdash s : \underline{\sigma}\), considered modulo \(\alpha\)-equivalence. Composition is performed by linear let bindings, and identities are given by the linear variable \(v\) itself (see the Linear variable typing rule).

Given a term \(p: \tau \vdash t : \sigma\) in \(\CSyn\), write \(q:\sigma\) for the distinguished variable in the fibre over \(\sigma\). Its reindexing \(t^{\ast}: \LSyn[\sigma] \to \LSyn[\tau]\) acts by substitution on objects and terms: \[\begin{array}{rrl} \left( \text{Objects} \right) &\quad \Fun[nn]{ t^{\ast} }{ \underline{\sigma} } &:= \underline{\sigma}[t/q] \\ \left( \text{Morphisms} \right) &\quad \Fun[nn]{ t^{\ast} }{ \left( q : \sigma; v : \underline{\rho} \vdash s : \underline{\sigma} \right) } &:= p : \tau; v : \underline{\rho}[t/q] \vdash s[t/q] : \underline{\sigma}[t/q] \end{array}\]

10 \(\mathsf{Fam}(\Vect[o])\) denotational semantics↩︎

In Section 3.4, we discuss briefly the functor \(\llbracket \cdot\rrbracket_{\cat[]{F}}: \mathsf{Tgt}\to \mathsf{Fam}(\Vect[o])\), which interprets the target language in the category \(\mathsf{Fam}(\Vect[o])\). The goal of this Appendix is to describe this functor in further detail.

Thus cartesian types are interpreted as sets (the base part of the objects in \(\mathsf{Fam}(\Vect[o])\)), whereas linear judgements over cartesian contexts are interpreted in \(\Vect[o]\). For convenience, given a morphism \(f: U \to V\) in \(\Vect[o]\), we denote by \(\opmor{f}: V \to U\) the corresponding linear map in \(\Vect\). This lets us display fibre maps in the ordinary \(\Vect\) orientation.

10.1 Type semantics↩︎

The semantics of the \(\CSyn\) fragment establishes the base cartesian semantics, interpreting cartesian contexts as the base sets of objects of \(\mathsf{Fam}(\Vect[o])\). The term constructors shared with the source language have the same interpretation as in the source-language semantics in \(\Set\) (Sections 3.2 and 8.4), \[\begin{align} \llbracket 0\rrbracket_{\Set} &= \emptyset \\ \llbracket 1\rrbracket_{\Set} &= \left\{ \ast \right\} \\ \llbracket \srcreal[n]\rrbracket_{\Set} &= \R[n] \end{align} \qquad \begin{align} \llbracket \tau \times \sigma\rrbracket_{\Set} &= \llbracket \tau\rrbracket_{\Set} \times \llbracket \sigma\rrbracket_{\Set} \\ \llbracket \tau \sqcup \sigma\rrbracket_{\Set} &= \llbracket \tau\rrbracket_{\Set} \amalg \llbracket \sigma\rrbracket_{\Set} \\ \llbracket \opn[b]{M} \tau\rrbracket_{\Set} &= \Fun[nn]{ M }{ \llbracket \tau\rrbracket_{\Set} } \end{align}\] with deterministic primitives interpreted as their chosen total functions (which we assume to be differentiable), and \(\tuple{ \opn[b]{M}, \opn[b]{return}, \opn[b]{bind} }\) as \(\tuple{ M, \eta, \mu }\). For a cartesian judgement under the context \(p : \tau\), dependent sums are interpreted as follows: for every \(\gamma \in \llbracket \tau\rrbracket_{\Set}\), \[\label{eq:appendix-target-language-csyn-dependent-sum-interpretation} \fn[n]{ \llbracket \Sigma x : \sigma \; . \; \tau'\rrbracket_{\Set} }{ \gamma } := \coprod_{a \in \fn[n]{\llbracket \sigma\rrbracket_{\Set}}{\gamma}} \fn[n]{ \llbracket \tau'\rrbracket_{\Set} }{ \gamma, a }\tag{73}\] and the linear function types are interpreted as the set of linear morphisms between the linear interpretations of the respective linear types, \[\label{eq:appendix-target-language-csyn-linear-function-interpretation} \fn[n]{ \llbracket \underline{\sigma} \multimap\underline{\rho} \rrbracket_{\Set} }{ \gamma } := \fn[n]{ \Vect }{ \fn[n]{ \llbracket \underline{\sigma}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma }, \fn[n]{ \llbracket \underline{\rho}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } }\tag{74}\] where the linear interpretation \(\llbracket \cdot\rrbracket_{p : \tau}^{\mathrm{lin}}\) of linear types dependent on a cartesian variable \(p : \tau\) provides a real vector space interpretation for those linear types, making them into (non-topological) vector bundles over the base set \(\llbracket \tau\rrbracket_{\Set}\). As such, for every linear type \(\underline{\sigma}\) well-formed in the context \(p: \tau\), we get a functor (i.e. a family) \[\llbracket \underline{\sigma}\rrbracket_{p : \tau}^{\mathrm{lin}}: \llbracket \tau\rrbracket_{\Set} \to \Vect[o]\] The linear interpretation is also described inductively over the structure of linear types, for every \(\gamma \in \llbracket \tau\rrbracket_{\Set}\), \[\label{eq:appendix-target-language-interpretation-inductive-definition} \begin{align} \Fun[np]{ \llbracket \underline{\srcreal[k]}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \R[k] \\ \Fun[np]{ \llbracket \underline{1}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= 0 \\ \Fun[np]{ \llbracket \underline{\sigma} \times \underline{\rho}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \Fun[np]{ \llbracket \underline{\sigma}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } \oplus \Fun[np]{ \llbracket \underline{\rho}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } \\ \Fun[np]{ \llbracket \casenst{t}{i}{n}{x_{i}}{\underline{\sigma}_{i}} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \begin{cases} \Fun[np]{ \llbracket \underline{\sigma}_{i}\rrbracket_{p : \tau, x_{i} : \tau_{i}}^{\mathrm{lin}} }{ \gamma, a } &\text{if } \fn[n]{ \llbracket t\rrbracket_{\Set} }{ \gamma } = \opn[b]{in}_{i} a \end{cases} \\ \Fun[np]{ \llbracket \opn[b]{M}^{w}_{q : \Fun[nn]{ \opn[b]{M} }{ \rho }} \left[ x \mapsto \underline{\sigma} \right] \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \bigoplus_{x \in \operatorname{supp}\fn[n]{ \llbracket q\rrbracket_{\Set} }{\gamma} } \left( \fn[n]{ \llbracket \underline{\sigma}\rrbracket_{p : \tau, x : \rho}^{\mathrm{lin}} }{ \gamma, x } \oplus \R \right) \end{align}\tag{75}\] Compare the definition of the monadic cotangent type semantics to the definition of the fibre part of the monad on objects in \(\mathsf{Fam}(\Vect[o])\) given in Equation 14 .

With these definitions, we get a functor \(\llbracket \cdot\rrbracket_{\cat[]{F}}: \mathsf{Tgt}\to \mathsf{Fam}(\Vect[o])\) defined on objects, which, as described in Section 9.5, are given by a pair \(\tuple{ p : \tau, \underline{\sigma} }\) where \(p : \tau\) is a cartesian context and \(\underline{\sigma}\) a well-formed linear type under this context, as \[\label{eq:appendix-target-language-interpretation-functor-on-objects} \llbracket \tuple{ p : \tau, \underline{\sigma} }\rrbracket_{\cat[]{F}} := \tuple{ \llbracket \tau\rrbracket_{\Set}, \llbracket \underline{\sigma}\rrbracket_{p : \tau}^{\mathrm{lin}} }\tag{76}\] and on any morphism \(\tuple{ t, s }: \tuple{ p : \tau, \underline{\rho} } \to \tuple{ q : \sigma, \underline{\chi} }\), which consists of a cartesian arrow \(t : p \to q\) and a linear arrow in the fibre \(s: \underline{\chi}[t/q] \to \underline{\rho}\) in \(\LSyn[ p : \tau ]\), as \[\label{eq:appendix-target-language-fam-vectop-interpretation-morphisms} \llbracket \tuple{ t, s }\rrbracket_{\cat[]{F}} := \tuple{ \llbracket t\rrbracket_{\Set}, \left( \gamma \mapsto \left( \opmor[p]{ \fn[n]{ \llbracket s\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \in \fn[n]{ \Vect }{ \fn[n]{ \llbracket \underline{\chi}\rrbracket_{q : \sigma}^{\mathrm{lin}} }{ \fn[n]{ \llbracket t\rrbracket_{\Set} }{ \gamma } }, \fn[n]{ \llbracket \underline{\rho}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \right) \right) }\tag{77}\]

10.2 Term semantics↩︎

Next, we provide interpretations for the terms themselves, which will be mapped to morphisms in \(\mathsf{Fam}(\Vect[o])\). Again, this is done inductively over the structure of the terms, and the cartesian part of the language is interpreted in \(\Set\) as the base part of the morphisms.

10.2.0.1 Cartesian fragment

For the term constructors in common with the source language, the interpretation is the same as the one described in Section 8.4. For the other cartesian term constructors, in a cartesian context \(\Gamma=p:\tau\), we define, for all \(\gamma \in \llbracket \tau\rrbracket_{\Set}\), \[\label{eq:appendix-target-language-new-cartesian-term-constructors-semantics} \begin{align} \fn[n]{ \llbracket \Gamma \vdash \tuple{ a, t }: \Sigma \left( x : \sigma \right) . \tau' \rrbracket_{\Set} }{ \gamma } &= \tuple{ \fn[n]{ \llbracket a\rrbracket_{\Set} }{ \gamma }, \fn[n]{ \llbracket t\rrbracket_{\Set} }{ \gamma } } \\ \fn[n]{ \llbracket \Gamma \vdash \opn[b]{case} \; p \; \opn[b]{of} \; \tuple{x,y} \to s : \rho \rrbracket_{\Set} }{ \gamma } &= \fn[n]{ \llbracket s\rrbracket_{\Set} }{ \gamma, a, b } \qq{where} \tuple{ a, b } = \fn[n]{ \llbracket p\rrbracket_{\Set} }{ \gamma } \\ \fn[n]{ \llbracket \Gamma \vdash \lambda v \; . \; s: \underline{\sigma} \multimap\underline{\rho} \rrbracket_{\Set} }{ \gamma } &= \opmor[p]{ \fn[n]{ \llbracket s\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \end{align}\tag{78}\] Linear application is a linear term constructor; its semantics is given below.

10.2.0.2 Linear core

Linear terms are interpreted fibrewise in \(\Vect[o]\), or, equivalently, as linear maps in \(\Vect\), but with domain and codomain reversed. The first display below is written in the \(\Vect[o]\)-orientation; when we need the ordinary \(\Vect\)-orientation, we apply \(\opmor{\mathord{-}}\). In what follows, we consider the linear terms under a cartesian context \(\Gamma = p : \tau\), \[\begin{align} {2} \fn[n]{ \llbracket \Gamma; v : \underline{\sigma} \vdash v : \underline{\sigma} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \mathrm{id}_{ \Fun[np]{ \llbracket \underline{\sigma} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \\ \fn[n]{ \llbracket \opn[b]{0} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= 0 \\ \fn[n]{ \llbracket s + t \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \fn[n]{ \llbracket s \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } + \fn[n]{ \llbracket t \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } \\ \fn[n]{ \llbracket \tuple{ s, t } \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \cotupling{ \fn[n]{ \llbracket s \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma }, \fn[n]{ \llbracket t \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } \\ \fn[n]{ \llbracket \opn[b]{pr}_{i} u \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \fn[n]{ \llbracket u \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } \circ \iota_{i} && \quad \left( i \in \left\{ 1, 2 \right\} \right) \\ \fn[n]{ \llbracket \mathsf{let}\;v=t\;\mathsf{in}\;s \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \fn[n]{ \llbracket t\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } \circ \fn[n]{ \llbracket s\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } && \quad \text{(Composition in \( \Vect[o] \))} \\ \fn[n]{ \llbracket \casenst{t}{i}{n}{x_{i}}{r_{i}} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } &= \fn[n]{ \llbracket r_{i}\rrbracket_{p : \tau, x_{i} : \tau_{i}}^{\mathrm{lin}} }{ \gamma, a } && \quad \text{if } \fn[n]{ \llbracket t\rrbracket_{\Set} }{ \gamma } = \fn[n]{ \opn[b]{in}_{i} }{ a } \end{align}\] For linear application, where \(\Gamma \vdash r : \underline{\sigma} \multimap\underline{\rho}\) and \(\Gamma; v : \underline{\alpha} \vdash t : \underline{\sigma}\), we use the ordinary \(\Vect\)-orientation \[\opmor[p]{ \fn[n]{ \llbracket \linapp{r}{t}\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } = \fn[n]{ \llbracket r\rrbracket_{\Set} }{ \gamma } \circ \opmor[p]{ \fn[n]{ \llbracket t\rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } }.\] Finally, for any built-in operation \(\opn[b]{op}: \srcreal[n_{1}] \times \cdots \times \srcreal[n_{k}] \to \srcreal[m]\), \[\opmor[p]{ \fn[n]{ \llbracket \fn[n]{ \fn[n]{ \opn[b]{Dop} }{ t_{1}, \cdots, t_{k} } }{v} \rrbracket_{p : \tau}^{\mathrm{lin}} }{ \gamma } } = \fn[p]{ \grad \llbracket \opn[b]{op}\rrbracket_{\Set} }{ \fn[n]{ \llbracket t_{1}\rrbracket_{\Set} }{\gamma}, \cdots, \fn[n]{ \llbracket t_{k}\rrbracket_{\Set} }{\gamma} }: \R[m] \to \R[n_{1}] \oplus \cdots \oplus \R[n_{k}]\] where \(\llbracket \opn[b]{op}\rrbracket_{\Set}\) is the intended set-based interpretation of the operation \(\opn[b]{op}\), assumed to be differentiable, so that taking its transpose derivative \(\grad \llbracket \opn[b]{op}\rrbracket_{\Set}\) is well-defined.

10.2.0.3 Transposed derivative of \(\opn[b]{return}\)

Let the cartesian context be \(\Gamma=p:\chi\). Suppose \(p:\chi\vdash t:\tau\), \(p:\chi,x:\tau\vdash \underline{\sigma}\text{ type}\), and \(p:\chi; v:\underline{\sigma}[t/x]\vdash s:\underline{\rho}\). For \(\gamma\in\llbracket \chi\rrbracket_{\Set}\), put \[a:=\fn[n]{\llbracket t\rrbracket_{\Set}}{\gamma}, \qquad A_{a}:=\fn[n]{\llbracket \underline{\sigma}\rrbracket_{p:\chi,x:\tau}^{\mathrm{lin}}}{\gamma,a}, \qquad R_{\gamma}:=\fn[n]{\llbracket \underline{\rho}\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma}.\] Then the ordinary \(\Vect\)-orientation of the fibre map is \[\label{eq:appendix-target-language-transpose-derivative-of-return} \opmor[p]{ \fn[n]{ \llbracket \fn[n]{ \opn[b]{return}^{w}_{t} }{ s } \rrbracket_{p : \chi}^{\mathrm{lin}} }{ \gamma } } = \opmor[p]{ \fn[n]{ \llbracket s\rrbracket_{p : \chi}^{\mathrm{lin}} }{ \gamma } } \circ \pi_{A_{a}} : A_{a}\oplus\R \to R_{\gamma}.\tag{79}\] Here \(\pi_{A_{a}}:A_{a}\oplus\R\to A_{a}\) is the projection onto the returned-value cotangent. The log-weight component is discarded, as expected for the derivative of a unit distribution with constant log-weight.

10.2.0.4 Transposed derivative of \(\opn[b]{bind}\)

Let the cartesian context again be \(\Gamma=p:\chi\), and consider a well-typed term \(\fn[n]{ \opn[b]{bind}^{w}_{s; x \; . \; t}}{ r ; u }\) with \(p:\chi\vdash s:\Fun[nn]{\opn[b]{M}}{\sigma}\) and \(p:\chi,x:\sigma\vdash t:\Fun[nn]{\opn[b]{M}}{\tau}\). Fix \(\gamma\in\llbracket \chi\rrbracket_{\Set}\), and define \[\begin{align} w &:= \fn[n]{\llbracket s\rrbracket_{\Set}}{\gamma} &&\in \Fun[nn]{M}{\llbracket \sigma\rrbracket_{\Set}}, \\ \mu_{x} &:= \fn[n]{\llbracket t\rrbracket_{\Set}}{\gamma,x} &&\in \Fun[nn]{M}{\llbracket \tau\rrbracket_{\Set}}, \\ f_{\gamma}(x) &:= \mu_{x} &&:\llbracket \sigma\rrbracket_{\Set}\to\Fun[nn]{M}{\llbracket \tau\rrbracket_{\Set}}, \\ \widehat w &:= \fn[n]{\Fun[nn]{M}{f_{\gamma}}}{w} &&\in \Fun[nn]{M}{\Fun[nn]{M}{\llbracket \tau\rrbracket_{\Set}}}, \\ \bar w &:= \fn[n]{\Nat[nn]{\mu}{\llbracket \tau\rrbracket_{\Set}}}{\widehat w} &&\in \Fun[nn]{M}{\llbracket \tau\rrbracket_{\Set}}. \end{align}\] Thus \(\bar w\) is the set-level denotation of \(\fn[n]{\opn[b]{bind}}{s;x.t}\) at \(\gamma\). Put \[A_{x}:=\fn[n]{\llbracket \underline{\sigma}\rrbracket_{p:\chi,x:\sigma}^{\mathrm{lin}}}{\gamma,x}, \qquad B_{y}:=\fn[n]{\llbracket \underline{\tau}\rrbracket_{p:\chi,y:\tau}^{\mathrm{lin}}}{\gamma,y}, \qquad R_{\gamma}:=\fn[n]{\llbracket \underline{\rho}\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma}.\] The two linear arguments give ordinary-orientation maps \[\begin{align} r^{\sharp}_{\gamma} &:={} \opmor[p]{ \fn[n]{\llbracket r\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma} } : \bigoplus_{x\in\operatorname{supp}w} \left(A_{x}\oplus\R\right) \to R_{\gamma}, \\ u^{\sharp}_{\gamma,x} &:={} \opmor[p]{ \fn[n]{\llbracket u\rrbracket_{p:\chi,x:\sigma}^{\mathrm{lin}}}{\gamma,x} } : \bigoplus_{y\in\operatorname{supp}\mu_{x}} \left(B_{y}\oplus\R\right) \to R_{\gamma}\oplus A_{x}. \end{align}\] The family \(\{u^{\sharp}_{\gamma,x}\}_{x\in\llbracket \sigma\rrbracket_{\Set}}\) is the fibre part of a morphism in \(\mathsf{Fam}(\Vect[o])\) \[\widehat f_{\gamma}: \tuple{ \llbracket \sigma\rrbracket_{\Set}, x\mapsto R_{\gamma}\oplus A_{x} } \to \Fun[nn]{M^{\cat[]{F}}}{ \tuple{ \llbracket \tau\rrbracket_{\Set}, y\mapsto B_{y} } },\] whose base map is \(f_{\gamma}\). Applying the lifted FAD monad to this morphism and evaluating the fibre map at \(w\) gives \[\Theta^{u}_{w} := \fn[n]{ \Nat[nn]{ \left(\Fun[nn]{M^{\cat[]{F}}}{\widehat f_{\gamma}}\right)_{\mathrm{fbr}} }{w} }: \bigoplus_{\xi\in\operatorname{supp}\widehat w} \left( \left( \bigoplus_{y\in\operatorname{supp}\xi} \left(B_{y}\oplus\R\right) \right) \oplus\R \right) \to \bigoplus_{x\in\operatorname{supp}w} \left( \left(R_{\gamma}\oplus A_{x}\right) \oplus\R \right).\] Similarly, let \[\varpi^{B}_{\widehat w} := \fn[n]{ \Nat[nn]{ \left( \Nat[nn]{\mu^{\cat[]{F}}}{ \tuple{ \llbracket \tau\rrbracket_{\Set}, y\mapsto B_{y} } } \right)_{\mathrm{fbr}} }{\widehat w} }: \bigoplus_{y\in\operatorname{supp}\bar w} \left(B_{y}\oplus\R\right) \to \bigoplus_{\xi\in\operatorname{supp}\widehat w} \left( \left( \bigoplus_{y\in\operatorname{supp}\xi} \left(B_{y}\oplus\R\right) \right) \oplus\R \right)\] be the ordinary-orientation fibre map of the lifted monad multiplication at \(\widehat w\). For \(C_{x}:=R_{\gamma}\oplus A_{x}\), define \[\begin{align} \Omega_{\gamma} &: \bigoplus_{x\in\operatorname{supp}w}\left(C_{x}\oplus\R\right) \to R_{\gamma}, \\ \Omega_{\gamma} &:={} \left( \nabla_{R_{\gamma}} \circ \bigoplus_{x\in\operatorname{supp}w} \left(\pi_{R_{\gamma}}\circ\pi_{1}\right) \right) + \left( r^{\sharp}_{\gamma} \circ \bigoplus_{x\in\operatorname{supp}w} \tupling{ \pi_{A_{x}}\circ\pi_{1}, \pi_{2} } \right). \end{align}\] Here \(\pi_{1}:C_{x}\oplus\R\to C_{x}\), \(\pi_{2}:C_{x}\oplus\R\to\R\), and \(\pi_{R_{\gamma}}\), \(\pi_{A_x}\) are the projections from \(C_x=R_\gamma\oplus A_x\). The first summand of \(\Omega_\gamma\) adds the direct ambient cotangents produced by the continuation backpropagators; the second feeds the cotangent of the bound value and the outer log-weight cotangent through \(r\).

The semantics of \(\opn[b]{bind}^{w}\) is then \[\label{eq:appendix-target-language-transpose-derivative-of-bind} \opmor[p]{ \fn[n]{ \llbracket \fn[n]{ \opn[b]{bind}^{w}_{s; x \; . \; t}}{ r ; u } \rrbracket_{p : \chi}^{\mathrm{lin}} }{ \gamma } } = \Omega_{\gamma} \circ \Theta^{u}_{w} \circ \varpi^{B}_{\widehat w} : \bigoplus_{y\in\operatorname{supp}\bar w}\left(B_{y}\oplus\R\right) \to R_{\gamma}.\tag{80}\] Because \(\Theta^{u}_{w}\) and \(\varpi^{B}_{\widehat w}\) are the actual fibre maps of the lifted FAD functor and multiplication, this definition also covers collisions of atoms and collisions among the intermediate FADs \(\mu_x\).

10.2.0.5 Transposed derivative of \(\opn[b]{categorical}\)

Again work in a cartesian context \(p:\chi\). Consider a well-typed term \[\fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{\tuple{t_{i},w_{i}}}_{i=1}^{n} } }{ \tuple{\tuple{t_{i}',w_{i}'}}_{i=1}^{n} }\] generated by the rule in the target language. Fix \(\gamma\in\llbracket \chi\rrbracket_{\Set}\), and put \[a_{i}:=\fn[n]{\llbracket t_{i}\rrbracket_{\Set}}{\gamma}, \qquad \ell_{i}:=\fn[n]{\llbracket w_{i}\rrbracket_{\Set}}{\gamma}, \qquad q:= \fn[n]{\llbracket \fn[n]{\opn[b]{categorical}}{ \tuple{\tuple{t_{i},w_{i}}}_{i=1}^{n} } \rrbracket_{\Set}}{\gamma}.\] Let \[A_{x}:=\fn[n]{\llbracket \underline{\sigma}\rrbracket_{p:\chi,x:\tau}^{\mathrm{lin}}}{\gamma,x}, \qquad R_{\gamma}:=\fn[n]{\llbracket \underline{\rho}\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma}.\] The lifted categorical operation \(\operatorname{categorical}^{\cat[]{F}}\) has an ordinary-orientation fibre map at the list \(\tuple{\tuple{a_{i},\ell_{i}}}_{i=1}^{n}\), which we write as \[\kappa_{\gamma}: \bigoplus_{x\in\operatorname{supp}q}\left(A_{x}\oplus\R\right) \to \bigoplus_{i=1}^{n}\left(A_{a_i}\oplus\R\right).\] This is the fibre part of the left factor in the FAD factorisation; in particular, it performs the correct splitting when several displayed atoms denote the same value. Define \[\begin{align} t^{\sharp}_{i,\gamma} &:={} \opmor[p]{ \fn[n]{\llbracket t_{i}'\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma} } : A_{a_i}\to R_{\gamma}, \\ w^{\sharp}_{i,\gamma} &:={} \opmor[p]{ \fn[n]{\llbracket w_{i}'\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma} } : \R\to R_{\gamma}. \end{align}\] Then \[\label{eq:appendix-target-language-transpose-derivative-of-categorical} \begin{align} &\opmor[p]{ \fn[n]{ \llbracket \fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{\tuple{t_{i},w_{i}}}_{i=1}^{n} } }{ \tuple{\tuple{t_{i}',w_{i}'}}_{i=1}^{n} } \rrbracket_{p:\chi}^{\mathrm{lin}} }{\gamma} } \\ &\qquad= \nabla_{R_{\gamma}} \circ \left( \bigoplus_{i=1}^{n} \left( t^{\sharp}_{i,\gamma}\circ\pi_{A_{a_i}} + w^{\sharp}_{i,\gamma}\circ\pi_{\R} \right) \right) \circ \kappa_{\gamma}. \end{align}\tag{81}\]

10.2.0.6 Transposed derivative of the handler algebra

Let \(\nu=\srcreal[d]\), and write \(V_{\nu}:=\R[d]\) for the fibre of \(\underline{\tau}_{\nu}\). Suppose \(p:\chi\vdash t:\Fun[nn]{\opn[b]{M}}{\nu}\), and fix \(\gamma\in\llbracket \chi\rrbracket_{\Set}\). Put \(m:=\fn[n]{\llbracket t\rrbracket_{\Set}}{\gamma}\). The lifted handler algebra \(\opn[b]{E}^{\cat[]{F}}_{\nu}:\Fun[nn]{M^{\cat[]{F}}}{\TR[d]}\to\TR[d]\) has ordinary-orientation fibre map \[\epsilon_{m}:V_{\nu}\to \bigoplus_{y\in\operatorname{supp}m}\left(V_{\nu}\oplus\R\right),\] which, for the expectation handler, is explicitly \[\fn[n]{\epsilon_{m}}{c} = \sum_{y\in\operatorname{supp}m} \left( \fn[n]{m}{y}\,c\cdot\mathbf{e}_{y} + \fn[n]{m}{y}\,\left\langle c,y\right\rangle\cdot\mathbf{\varepsilon}_{y} \right), \qquad c\in V_{\nu}.\] Here \(\mathbf{e}_{y}\) denotes the atom-cotangent summand and \(\mathbf{\varepsilon}_{y}\) denotes the log-weight summand. Therefore, for \(p:\chi;v:\opn[b]{M}^{w}_{t}\left[y\mapsto\underline{\tau}_{\nu}\right]\vdash s:\underline{\rho}\), \[\label{eq:appendix-target-language-transpose-derivative-of-handler} \opmor[p]{ \fn[n]{ \llbracket \fn[n]{\opn[b]{E}^{w}_{\nu,t}}{s} \rrbracket_{p:\chi}^{\mathrm{lin}} }{\gamma} } = \opmor[p]{ \fn[n]{\llbracket s\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma} } \circ \epsilon_{m} : V_{\nu}\to \fn[n]{\llbracket \underline{\rho}\rrbracket_{p:\chi}^{\mathrm{lin}}}{\gamma}.\tag{82}\]

11 Code transformation rules↩︎

This section gives the full reverse-mode CHAD translation used in Equation 47 . For every source judgement \(\Gamma \vdash t : \tau\), the translation is a target-language term \[\label{eq:appendix-chad-term-typing-shape} \Fun[np]{\chad}{\Gamma}_{1} \vdash \Fun[np]{\chad[\Gamma]}{t} : \Sigma \left( p : \Fun[np]{\chad}{\tau}_{1} \right) . \left( \Fun[np]{\chad}{\tau}_{2}[p/p_{\tau}] \multimap \Fun[np]{\chad}{\Gamma}_{2} \right).\tag{83}\] The first component is the primal computation and the second component is the backpropagator. To avoid capture in dependent cotangent types, we write \(p_{\tau} : \Fun[np]{\chad}{\tau}_{1}\) for the distinguished primal variable of the type \(\tau\). Thus \(\Fun[np]{\chad}{\tau}_{2}\) is read in the cartesian context \(p_{\tau} : \Fun[np]{\chad}{\tau}_{1}\), and every substitution into a cotangent type below is understood after this \(\alpha\)-renaming.

11.0.0.1 Convention for the distinguished linear variable.

The target language has one distinguished linear variable. In clauses involving \(\opn[b]{return}^{w}\), \(\opn[b]{bind}^{w}\), \(\opn[b]{categorical}^{w}\), or \(\opn[b]{E}^{w}\), the variable bound by the outer expression \(\lambda v\; .\;(-)\) has the conclusion type of the corresponding \(w\)-typing rule. The occurrences of \(v\) inside the arguments of the \(w\)-constructor are checked at the premise types of that rule. This is exactly the typing convention used in the rules of Section 9.

11.1 Type and context translation↩︎

11.1.0.1 Base types, products, and sums.

The primal translation \(\Fun[np]{\chad}{-}_{1}\) is cartesian, while \(\Fun[np]{\chad}{-}_{2}\) gives a linear cotangent type indexed by the corresponding distinguished primal variable: \[\label{eq:appendix-chad-type-translation-core-basis} \begin{array}{rcl} \Fun[np]{\chad}{1}_{1} &:=& 1 \\[2pt] \Fun[np]{\chad}{0}_{1} &:=& 0 \\[2pt] \Fun[np]{\chad}{\srcreal[k]}_{1} &:=& \srcreal[k] \\[6pt] \Fun[np]{\chad}{\tau \times \sigma}_{1} &:=& \Fun[np]{\chad}{\tau}_{1} \times \Fun[np]{\chad}{\sigma}_{1} \\[6pt] \Fun[np]{\chad}{\tau \sqcup \sigma}_{1} &:=& \Fun[np]{\chad}{\tau}_{1} \sqcup \Fun[np]{\chad}{\sigma}_{1} \end{array}\tag{84}\] and \[\label{eq:appendix-chad-type-translation-core-cotangents} \begin{array}{rcl} \Fun[np]{\chad}{1}_{2} &:=& \underline{1} \\[2pt] \Fun[np]{\chad}{0}_{2} &:=& \underline{1} \\[2pt] \Fun[np]{\chad}{\srcreal[k]}_{2} &:=& \underline{\srcreal[k]} \\[6pt] \Fun[np]{\chad}{\tau \times \sigma}_{2} &:=& \Fun[np]{\chad}{\tau}_{2}[\opn[b]{pr}_{1}p_{\tau\times\sigma}/p_{\tau}] \times \Fun[np]{\chad}{\sigma}_{2}[\opn[b]{pr}_{2}p_{\tau\times\sigma}/p_{\sigma}]. \\[6pt] \Fun[np]{\chad}{\tau \sqcup \sigma}_{2} &:=& \opn[b]{case}\;p_{\tau\sqcup\sigma}\;\opn[b]{of} \left\{ \begin{array}{l} \opn[b]{inl}\,x \mapsto \Fun[np]{\chad}{\tau}_{2}[x/p_{\tau}] \\ \opn[b]{inr}\,y \mapsto \Fun[np]{\chad}{\sigma}_{2}[y/p_{\sigma}] \end{array} \right\}. \end{array}\tag{85}\]

For the linear case-type former we use the usual dependent \(\beta\)-rules. For example, \[\Fun[np]{\chad}{\tau \sqcup \sigma}_{2}[\opn[b]{inl}a/p_{\tau\sqcup\sigma}] \equiv \Fun[np]{\chad}{\tau}_{2}[a/p_{\tau}], \quad \Fun[np]{\chad}{\tau \sqcup \sigma}_{2}[\opn[b]{inr}b/p_{\tau\sqcup\sigma}] \equiv \Fun[np]{\chad}{\sigma}_{2}[b/p_{\sigma}].\] The same convention applies to the \(n\)-ary coproduct notation introduced in Section 8.

11.1.0.2 The monadic type.

For the FAD monad, the primal remains a monadic computation and the cotangent is the lifted monadic cotangent type of Section 9: \[\label{eq:appendix-chad-type-translation-monad} \begin{align} \Fun[np]{\chad}{\Fun[nn]{\opn[b]{M}}{\tau}}_{1} &:= \Fun[nn]{\opn[b]{M}}{\Fun[np]{\chad}{\tau}_{1}}, \\ \Fun[np]{\chad}{\Fun[nn]{\opn[b]{M}}{\tau}}_{2} &:= \opn[b]{M}^{w}_{p_{\opn[b]{M}\tau} : \Fun[nn]{\opn[b]{M}}{\Fun[np]{\chad}{\tau}_{1}}} \left[ x \mapsto \Fun[np]{\chad}{\tau}_{2}[x/p_{\tau}] \right]. \end{align}\tag{86}\] Here \(p_{\opn[b]{M}\tau}\) is the primal variable for the source type \(\Fun[nn]{\opn[b]{M}}{\tau}\), while \(p_{\tau}\) is the primal variable for the returned value type \(\tau\). This distinction is necessary because the continuation cotangent family is indexed by returned values, not by the whole monadic computation.

11.1.0.3 Contexts.

Contexts are translated componentwise on primals and by fibrewise biproducts on cotangents: \[\label{eq:appendix-chad-context-translation} \begin{align} \Fun[np]{\chad}{\cdot}_{1} &:= \cdot, & \Fun[np]{\chad}{\Gamma, x : \tau}_{1} &:= \Fun[np]{\chad}{\Gamma}_{1},\; x : \Fun[np]{\chad}{\tau}_{1}, \\ \Fun[np]{\chad}{\cdot}_{2} &:= \underline{1}, & \Fun[np]{\chad}{\Gamma, x : \tau}_{2} &:= \Fun[np]{\chad}{\Gamma}_{2} \times \Fun[np]{\chad}{\tau}_{2}[x/p_{\tau}]. \end{align}\tag{87}\] The first factor \(\Fun[np]{\chad}{\Gamma}_{2}\) is implicitly reindexed along the cartesian weakening \(\Fun[np]{\chad}{\Gamma}_{1},x:\Fun[np]{\chad}{\tau}_{1}\to \Fun[np]{\chad}{\Gamma}_{1}\). Thus \(\opn[b]{pr}_{1}\) selects the cotangent tuple for the previous context and \(\opn[b]{pr}_{2}\) selects the cotangent for the newest variable.

For \(\Gamma = x_{1}:\tau_{1},\dots,x_{n}:\tau_{n}\), define the derived injection of a single variable cotangent into the context cotangent by recursion. Writing \(\fn[n]{\opn{idx}}{x_{j};\Gamma}=j\), if \(\Gamma=\Gamma',x_{n}:\tau_{n}\), then \[\label{eq:appendix-chad-coproj-definition} \begin{align} \opn[b]{coproj}^{\Gamma}_{n}(v) &:= \tuple{ \opn[b]{0}_{\Fun[np]{\chad}{\Gamma'}_{2}}, v }, \\ \opn[b]{coproj}^{\Gamma}_{j}(v) &:= \tuple{ \opn[b]{coproj}^{\Gamma'}_{j}(v), \opn[b]{0}_{\Fun[np]{\chad}{\tau_{n}}_{2}[x_{n}/p_{\tau_n}]} } \qquad (1\leq j<n). \end{align}\tag{88}\] We write \(\opn[b]{coproj}_{\fn[n]{\opn{idx}}{x;\Gamma}}v\) when \(\Gamma\) is clear from context.

11.2 Term translation↩︎

All zero terms below are the typed zeroes of the relevant fibrewise biproduct. The term translation is by structural recursion on the source typing derivation.

11.2.0.1 Nullary terms and variables.

\[\begin{align} \tag{89} \Fun[np]{\chad[\Gamma]}{\tuple{\;}} &:= \tuple{ \tuple{\;}, \lambda v\; .\; \opn[b]{0}_{\Fun[np]{\chad}{\Gamma}_{2}} }, \\ \tag{90} \Fun[np]{\chad[\Gamma]}{\fn[n]{\opn[b]{abort}}{t}} &:= \opn[b]{case}\; \Fun[np]{\chad[\Gamma]}{t}\; \opn[b]{of}\; \tuple{p,b}\to \fn[n]{\opn[b]{abort}}{p}, \\ \tag{91} \Fun[np]{\chad[\Gamma]}{x} &:= \tuple{ x, \lambda v\; .\; \opn[b]{coproj}_{\fn[n]{\opn{idx}}{x;\Gamma}}v }. \end{align}\]

11.2.0.2 Products.

\[\begin{align} \tag{92} \Fun[np]{\chad[\Gamma]}{\tuple{t,s}} &:= \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{s}\;\opn[b]{of}\;\tuple{q,c}\to \notag\\[-2pt] &\qquad \tuple{ \tuple{p,q}, \lambda v\; .\; \fn[n]{b}{\opn[b]{pr}_{1}v} + \fn[n]{c}{\opn[b]{pr}_{2}v} }, \\ \tag{93} \Fun[np]{\chad[\Gamma]}{\opn[b]{pr}_{1}t} &:= \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \notag\\[-2pt] &\qquad \tuple{ \opn[b]{pr}_{1}p, \lambda v\; .\; \fn[n]{b}{ \tuple{ v, \opn[b]{0}_{\Fun[np]{\chad}{\sigma}_{2}[\opn[b]{pr}_{2}p/p_{\sigma}]} } } } \quad (t: \tau\times\sigma), \\ \tag{94} \Fun[np]{\chad[\Gamma]}{\opn[b]{pr}_{2}t} &:= \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \notag\\[-2pt] &\qquad \tuple{ \opn[b]{pr}_{2}p, \lambda v\; .\; \fn[n]{b}{ \tuple{ \opn[b]{0}_{\Fun[np]{\chad}{\tau}_{2}[\opn[b]{pr}_{1}p/p_{\tau}]}, v } } } \quad (t: \tau\times\sigma). \end{align}\]

11.2.0.3 Coproducts.

For \(n\)-ary sums, the following introduction rule specializes to \(\opn[b]{inl}\) and \(\opn[b]{inr}\) in the binary grammar: \[\label{eq:appendix-chad-sum-intro} \Fun[np]{\chad[\Gamma]}{\opn[b]{in}_{i}t} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{x,x'}\to \tuple{\opn[b]{in}_{i}x,x'}.\tag{95}\] For case analysis, suppose \(\Gamma\vdash t:\tau_{1}\sqcup\cdots\sqcup\tau_{n}\) and \(\Gamma,x_{i}:\tau_{i}\vdash s_{i}:\rho\). Then \[\begin{align} \label{eq:appendix-chad-sum-elim} &\Fun[np]{\chad[\Gamma]}{ \opn[b]{case}\;t\;\opn[b]{of}\; \left\{\opn[b]{in}_{i}x_{i}\to s_{i}\right\}_{i=1}^{n} } := \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{y,y'}\to \notag\\ &\quad \opn[b]{case}\;y\;\opn[b]{of} \left\{ \begin{array}{l} \opn[b]{in}_{i}x_{i}\mapsto \opn[b]{case}\;\Fun[np]{\chad[\Gamma,x_{i}:\tau_{i}]}{s_{i}}\;\opn[b]{of}\;\tuple{z_{i},z_{i}'}\to \\ \qquad \tuple{ z_{i}, \lambda v\; .\; \mathsf{let}\;v=\fn[n]{z_{i}'}{v}\;\mathsf{in}\; \opn[b]{pr}_{1}v + \fn[n]{y'}{\opn[b]{pr}_{2}v} } \end{array} \right\}_{i=1}^{n}. \end{align}\tag{96}\] The application of \(y'\) in the branch uses the case-type \(\beta\)-rule above, so \(\opn[b]{pr}_{2}v\) has precisely the cotangent type of the active summand.

11.2.0.4 Let bindings and deterministic operations.

If \(\Gamma\vdash t:\sigma\) and \(\Gamma,x:\sigma\vdash s:\rho\), then \[\begin{align} \label{eq:appendix-chad-let} \Fun[np]{\chad[\Gamma]}{\mathsf{let}\;x=t\;\mathsf{in}\;s} &:= \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \opn[b]{case}\;\Fun[np]{\chad[\Gamma,x:\sigma]}{s}[p/x]\;\opn[b]{of}\;\tuple{q,c}\to \notag\\[-2pt] &\qquad \tuple{ q, \lambda v\; .\; \mathsf{let}\;v=\fn[n]{c}{v}\;\mathsf{in}\; \opn[b]{pr}_{1}v + \fn[n]{b}{\opn[b]{pr}_{2}v} }. \end{align}\tag{97}\] For a built-in operation \(\mathsf{op}: \srcreal[n_{1}]\times\cdots\times\srcreal[n_{k}]\to\srcreal[m]\), write \(\mathsf{op}\) for the chosen transposed Jacobian term constructor. Its CHAD clause is \[\begin{align} \label{eq:appendix-chad-op} &\Fun[np]{\chad[\Gamma]}{\fn[n]{\mathsf{op}}{t_{1},\dots,t_{k}}} := \notag\\ &\quad \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t_{1}}\;\opn[b]{of}\;\tuple{x_{1},x_{1}'}\to \cdots \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t_{k}}\;\opn[b]{of}\;\tuple{x_{k},x_{k}'}\to \notag\\[-2pt] &\quad\qquad \tuple{ \fn[n]{\mathsf{op}}{x_{1},\dots,x_{k}}, \lambda v\; .\; \mathsf{let}\;v=\fn[n]{\fn[n]{\mathsf{op}}{x_{1},\dots,x_{k}}}{v}\;\mathsf{in}\; \sum_{j=1}^{k} \fn[n]{x_{j}'}{\opn[b]{pr}_{j}v} }. \end{align}\tag{98}\] Here \(\opn[b]{pr}_{j}\) denotes the derived projection from the cotangent product returned by \(\mathsf{op}\).

11.2.0.5 Monad unit and multiplication.

The unit clause is the one displayed in Equation 49 : \[\label{eq:appendix-chad-return} \Fun[np]{\chad[\Gamma]}{\fn[n]{\opn[b]{return}}{t}} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \tuple{ \fn[n]{\opn[b]{return}}{p}, \lambda v\; .\; \fn[n]{\opn[b]{return}^{w}_{p}}{\fn[n]{b}{v}} }.\tag{99}\] For bind, let \(\Gamma\vdash t:\Fun[nn]{\opn[b]{M}}{\sigma}\) and \(\Gamma,x:\sigma\vdash s:\Fun[nn]{\opn[b]{M}}{\tau}\). In the scope of \(x\), put \[r_{x} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma,x:\sigma]}{s}\;\opn[b]{of}\;\tuple{q,c}\to q, \qquad c_{x} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma,x:\sigma]}{s}\;\opn[b]{of}\;\tuple{q,c}\to c.\] Then the bind clause is \[\label{eq:appendix-chad-bind} \Fun[np]{\chad[\Gamma]}{\fn[n]{\opn[b]{bind}}{t; x\; .\; s}} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \tuple{ \fn[n]{\opn[b]{bind}}{p; x\; .\; r_{x}}, \lambda v\; .\; \fn[n]{\opn[b]{bind}^{w}_{p; x\; .\; r_{x}}}{ \fn[n]{b}{v}; \fn[n]{c_{x}}{v} } }.\tag{100}\] The first argument of \(\opn[b]{bind}^{w}\) accumulates the cotangent contribution through the sampled computation \(t\), while the second argument accumulates the cotangent contribution through the continuation \(s\), including the cotangent of the bound variable \(x\).

11.2.0.6 FAD constructors and handler.

For the finite atomic distribution constructor, assume \(\Gamma\vdash t_{i}:\tau\) and \(\Gamma\vdash w_{i}:\srcreal\) for \(1\leq i\leq n\). Then \[\begin{align} \label{eq:appendix-chad-categorical} &\Fun[np]{\chad[\Gamma]}{ \fn[n]{\opn[b]{categorical}}{\tuple{\tuple{t_{i},w_{i}}}_{i=1}^{n}} } := \notag\\ &\quad \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t_{1}}\;\opn[b]{of}\;\tuple{u_{1},u_{1}'}\to \cdots \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t_{n}}\;\opn[b]{of}\;\tuple{u_{n},u_{n}'}\to \notag\\ &\quad \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{w_{1}}\;\opn[b]{of}\;\tuple{a_{1},a_{1}'}\to \cdots \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{w_{n}}\;\opn[b]{of}\;\tuple{a_{n},a_{n}'}\to \notag\\[-2pt] &\quad\qquad \tuple{ \fn[n]{\opn[b]{categorical}}{\tuple{\tuple{u_{i},a_{i}}}_{i=1}^{n}}, \lambda v\; .\; \fn[n]{ \opn[b]{categorical}^{w}_{\tuple{\tuple{u_{i},a_{i}}}_{i=1}^{n}} }{ \tuple{\tuple{\fn[n]{u_{i}'}{v},\fn[n]{a_{i}'}{v}}}_{i=1}^{n} } }. \end{align}\tag{101}\] Finally, for the handler algebra \(\opn[b]{E}_{\nu}\), where \(\nu=\srcreal[p]\), we define \[\label{eq:appendix-chad-handler} \Fun[np]{\chad[\Gamma]}{\fn[n]{\opn[b]{E}_{\nu}}{t}} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \tuple{ \fn[n]{\opn[b]{E}_{\nu}}{p}, \lambda v\; .\; \fn[n]{\opn[b]{E}^{w}_{\nu,p}}{\fn[n]{b}{v}} }.\tag{102}\]

11.3 Typing and universal characterization↩︎

The typing statement follows by induction on the source typing derivation. The only non-structural cases are the monadic and FAD constructors, where the typing rules for \(\opn[b]{return}^{w}\), \(\opn[b]{bind}^{w}\), \(\opn[b]{categorical}^{w}\), and \(\opn[b]{E}^{w}\) are exactly the target rules used in the displayed clauses. The uniqueness statement is the universal property of the freely generated source syntactic category (Lemma 8) applied to the corresponding structure in \(\mathsf{Tgt}\). The target equations of Section 9 ensure that the source Kleisli, distributive, and Eilenberg–Moore equations are preserved.

12 Example of differentiated code↩︎

This section gives the complete differentiated example code from Equation 58 , following the inductive rules of code differentiation outlined in Appendix 11. As in Section 5, we adopt the following abbreviations: \[\begin{align} a_{0} &:= 0 \\ a_{1} &:= P \cdot \fn[n]{ \opn[b]{sig} }{ -s_{p} \cdot \left( p - p_{0} \right) } \end{align} \qquad \begin{align} \gamma_{0} &:= \fn[n]{ \opn[b]{lsig} }{ -s_{m} \cdot \left( m - m_{0} \right) } \\ \gamma_{1} &:= \fn[n]{ \opn[b]{lsig} }{ s_{m} \cdot \left( m - m_{0} \right) } \end{align}\] The full code transformation is then given by \[\begin{align} {2} &\opn[b]{case} \\ &\qquad \opn[b]{case} \; \tuple{ a_{0}, \grad a_{0} } \; \opn[b]{of} \tuple{ \theta_{0}, \theta_{0}' } \to \\ &\qquad \opn[b]{case} \; \tuple{ a_{1}, \grad a_{1} } \; \opn[b]{of} \tuple{ \theta_{1}, \theta_{1}' } \to \\ &\qquad \opn[b]{case} \; \tuple{ \gamma_{0}, \grad \gamma_{0} } \; \opn[b]{of} \tuple{ \rho_{0}, \rho_{0}' } \to \\ &\qquad \opn[b]{case} \; \tuple{ \gamma_{1}, \grad \gamma_{1} } \; \opn[b]{of} \tuple{ \rho_{1}, \rho_{1}' } \to \\ &\qquad \qquad \tuple{ \fn[n]{ \opn[b]{categorical} }{ \tuple{ \tuple{ \theta_{j}, \rho_{j} } }_{j = 0}^{1} }, \lambda v \; . \; \fn[n]{ \opn[b]{categorical}^{w}_{ \tuple{ \tuple{ \theta_{j}, \rho_{j} } }_{j = 0}^{1} } }{ \tuple{ \tuple{ \fn[n]{ \theta_{j}' }{ v }, \fn[n]{ \rho_{j}' }{ v } } }_{j = 0}^{1} } } \\ &\opn[b]{of} \tuple{ \alpha, \alpha' } \to \\ &\qquad \opn[b]{case} \\ &\qquad \qquad \opn[b]{case} \\ &\qquad \qquad \qquad \opn[b]{case} \tuple{ s, \lambda v \; . \; \opn[b]{coproj}_{\fn[n]{ \opn{idx} }{ s; \Gamma, s }} v } \\ &\qquad \qquad \qquad \opn[b]{of} \tuple{ \eta, \eta' } \to \\ &\qquad \qquad \qquad \qquad \opn[b]{let} \xi = \\ &\qquad \qquad \qquad \qquad \qquad \opn[b]{case} \\ &\qquad \qquad \qquad \qquad \qquad \qquad \opn[b]{case} \tuple{ \left( p \cdot x - m \right), \grad_{p, m, x} \left( p \cdot x - m \right) } \\ &\qquad \qquad \qquad \qquad \qquad \qquad \opn[b]{of} \tuple{ \mu, \mu' } \to \tuple{ \fn[n]{ \opn[b]{return} }{ \mu }, \lambda v \; . \; \fn[n]{ \opn[b]{return}^{w}_{\mu} }{ \fn[n]{ \mu' }{ v } } } \\ &\qquad \qquad \qquad \qquad \qquad \opn[b]{of} \tuple{ \kappa, \kappa' } \to \kappa \\ &\qquad \qquad \qquad \qquad \opn[b]{in} \\ &\qquad \qquad \qquad \qquad \qquad \left\langle \opn[b]{bind} \left( \eta; x \; . \; \xi \right), \lambda v \; . \; \opn[b]{bind}_{\eta; x \; . \; \xi}^{w} \left( \fn[n]{ \eta' }{ v }; \right. \right. \\ &\qquad \qquad \qquad \qquad \qquad \qquad \left( \opn[b]{case} \right. \\ &\qquad \qquad \qquad \qquad \qquad \qquad \qquad \opn[b]{case} \tuple{ \left( p \cdot x - m \right), \grad_{p, m, x} \left( p \cdot x - m \right) } \\ &\qquad \qquad \qquad \qquad \qquad \qquad \qquad \opn[b]{of} \tuple{ \mu, \mu' } \to \tuple{ \fn[n]{ \opn[b]{return} }{ \mu }, \lambda v \; . \; \fn[n]{ \opn[b]{return}^{w}_{\mu} }{ \fn[n]{ \mu' }{ v } } } \\ &\qquad \qquad \qquad \qquad \qquad \qquad \left. \left.\left. \opn[b]{of} \tuple{ \kappa, \kappa' } \to \kappa' \right) \left( v \right) \right) \right\rangle \\ &\qquad \qquad \opn[b]{of} \tuple{ \delta, \delta' } \to \\ &\qquad \qquad \qquad \opn[b]{case} \\ &\qquad \qquad \qquad \qquad \opn[b]{case} \tuple{ t, \lambda v \; . \; \opn[b]{coproj}_{\fn[n]{ \opn{idx} }{ t; \Gamma, s, t }} v} \\ &\qquad \qquad \qquad \qquad \opn[b]{of} \tuple{ \zeta, \zeta' } \to \tuple{ \fn[n]{ \opn[b]{E}_{\srcreal} }{ \zeta }, \lambda v \; . \; \fn[n]{ \opn[b]{E}_{\srcreal, \zeta}^{w} }{ \fn[n]{ \zeta' }{ v } } } \\ &\qquad \qquad \qquad \opn[b]{of} \tuple{ \epsilon, \epsilon' } \to \tuple{ \epsilon[\delta/t], \lambda v \; . \; \mathsf{let}\;w= \fn[n]{ \epsilon'[\delta/t] }{ v } \;\mathsf{in}\; \opn[b]{pr}_{1} w + \fn[n]{ \delta' }{ \opn[b]{pr}_{2} w } } \\ &\qquad \opn[b]{of} \tuple{ \beta, \beta' } \to \tuple{ \beta[\alpha/s], \lambda v \; . \; \mathsf{let}\; w = \fn[n]{ \beta'[\alpha/s] }{v} \;\mathsf{in}\; \opn[b]{pr}_{1} w + \fn[n]{ \alpha' }{ \opn[b]{pr}_{2} w } } \end{align}\]

13 Deriving concrete gradient estimation formulas↩︎

We now explain how the usual gradient rule for weighted particles is already contained in the lifted FAD monad and in the lifted canonical algebra.

In this section, we write the second component of a morphism in \(\mathsf{Fam}(\Vect[o])\) using the symbol \(\nabla\). Thus a morphism \(\hat{h}:T^*X\to T^*Y\) is written as \[\hat{h}=(h,\nabla h),\] where \(h:X\to Y\) is its base part and \[(\nabla h)_x:T^*_{h(x)}Y\to T^*_xX\] is the \(x\)-fibre of its fibre part. This notation is suggestive: \(\nabla h\) denotes the chosen cotangent map in \(\mathsf{Fam}(\Vect[o])\), whether or not \(X\) and \(Y\) carry genuine smooth structure.

Write \[\opn{E}^{\cat[]{F}}_{\R}:M^{\cat[]{F}}\TR\to\TR\] for the lifted canonical algebra on real values. Its base part is the unnormalised expectation map \[u\longmapsto \sum_{r\in\operatorname{supp}u}u(r)r, \qquad u\in\Fun[nn]{M}{\R},\] and its fibre part is \[\label{eq:appendix-lifted-real-expectation-fibre} \left(\nabla\opn{E}^{\cat[]{F}}_{\R}\right)_u(t) = \sum_{r\in\operatorname{supp}u} \left( u(r)t\cdot\mathbf{e}_{r} + u(r)r t\cdot\mathbf{\varepsilon}_{r} \right).\tag{103}\] Here the \(\mathbf{e}_{r}\)-component is the cotangent of the real value \(r\), while the \(\mathbf{\varepsilon}_{r}\)-component is the cotangent of the log-weight attached to \(r\).

Let \[T^*X=\tuple{X,\{T^*_xX\}_{x\in X}}\] be an object of \(\mathsf{Fam}(\Vect[o])\), and let \[\hat{g}=(g,\nabla g):T^*X\to\TR\] be a real-valued test morphism.

For \(w\in\Fun[nn]{M}{X}\), set \[u:=\fn[n]{\Fun[nn]{M}{g}}{w}.\] The Frobenius, or pushforward, identity says on bases that \[\int_X g(x)\,dw = \int_{\R} r\,d\fn[n]{\Fun[nn]{M}{g}}{w},\] or explicitly \[\sum_{x\in\operatorname{supp}w}w(x)g(x) = \sum_{r\in\operatorname{supp}u}u(r)r.\] In categorical form, this is the base part of the composite \[M^{\cat[]{F}}T^*X \xrightarrow{M^{\cat[]{F}}\hat{g}} M^{\cat[]{F}}\TR \xrightarrow{\opn{E}^{\cat[]{F}}_{\R}} \TR.\]

The fibre part of this composite is obtained by composing Equation 103 with the fibre part of \(M^{\cat[]{F}}\hat{g}\). For \(x\in\operatorname{supp}w\), the denominator is nonzero, and the corresponding summand is \[\frac{w(x)}{u(g(x))} \left( (\nabla g)_x \left( u(g(x))t \right) \cdot\mathbf{e}_{x} + u(g(x))g(x)t\cdot\mathbf{\varepsilon}_{x} \right).\] Since \((\nabla g)_x\) is linear, the factor \(u(g(x))\) in the atom-cotangent part cancels with the denominator; the log-weight part cancels directly. Hence \[\label{eq:appendix-frobenius-cotangent-rule} \left( \nabla\!\left( \opn{E}^{\cat[]{F}}_{\R} \circ M^{\cat[]{F}}\hat{g} \right) \right)_w(t) = \sum_{x\in\operatorname{supp}w} \left( w(x)(\nabla g)_x(t)\cdot\mathbf{e}_{x} + w(x)g(x)t\cdot\mathbf{\varepsilon}_{x} \right).\tag{104}\] Thus the lifted algebra and the lifted FAD monad already produce the two basic cotangent contributions: a value cotangent weighted by \(w(x)\), and a log-weight cotangent weighted by \(w(x)g(x)\).

Now apply this to a parameterised family of particles. Let \[T^*\Theta=\tuple{\Theta,\{T^*_\theta\Theta\}_{\theta\in\Theta}}, \qquad T^*Z=\tuple{Z,\{T^*_zZ\}_{z\in Z}},\] and let \[\hat{q}=(q,\nabla q):T^*\Theta\to M^{\cat[]{F}}T^*Z\] be a parameterised FAD, equipped with its chosen cotangent map. Locally, away from collisions of particles, choose a labelling of the support and write the base part of \(\hat{q}\) as \[\fn[n]{ q}{\theta} = \sum_{i\in I}w_i(\theta)\delta_{z_i(\theta)}.\] The corresponding strengthened FAD over \(T^*\Theta\times T^*Z\) is \[\hat{\underline q}=\tuple{\underline q,\nabla\underline q} := \Nat[nn]{\operatorname{st}^{\cat[]{F}}}{T^*\Theta,T^*Z} \circ \tupling{\mathrm{id}_{T^*\Theta},\hat{q}} : T^*\Theta \to M^{\cat[]{F}}(T^*\Theta\times T^*Z),\] whose base part is \[\fn[n]{\underline q}{\theta} = \sum_{i\in I}w_i(\theta)\delta_{(\theta,z_i(\theta))}.\]

In the chosen local labelling, the fibre of \(M^{\cat[]{F}}(T^*\Theta\times T^*Z)\) at \(\fn[n]{\underline q}{\theta}\) identifies with \[\bigoplus_{i\in I} \left( \left( T^*_\theta\Theta\oplus T^*_{z_i(\theta)}Z \right) \oplus\R \right).\] We write elements of this fibre as finite sums \[\sum_{i\in I} \left( \tuple{d_i,a_i}\cdot\mathbf{e}_{(\theta,z_i(\theta))} + s_i\cdot\mathbf{\varepsilon}_{(\theta,z_i(\theta))} \right),\] where \[d_i\in T^*_\theta\Theta, \qquad a_i\in T^*_{z_i(\theta)}Z, \qquad s_i\in\R.\] Let \[\ell_i:\Theta\to\R, \qquad \ell_i(\theta)=\log w_i(\theta).\] In these coordinates, the fibre part of \(\hat{\underline q}\) is \[\label{eq:appendix-suggestive-fibre-parametrised-fad} \begin{align} & (\nabla\underline q)_\theta \left( \sum_{i\in I} \left( \tuple{d_i,a_i}\cdot\mathbf{e}_{(\theta,z_i(\theta))} + s_i\cdot\mathbf{\varepsilon}_{(\theta,z_i(\theta))} \right) \right) \\ &\qquad = \sum_{i\in I}d_i + \sum_{i\in I}(\nabla z_i)_\theta(a_i) + \sum_{i\in I}(\nabla\ell_i)_\theta(s_i). \end{align}\tag{105}\] Here \[(\nabla z_i)_\theta:T^*_{z_i(\theta)}Z\to T^*_\theta\Theta\] is the \(\theta\)-fibre of the fibre part of the chosen morphism \(\hat{z}_i=(z_i,\nabla z_i):T^*\Theta\to T^*Z\) representing the parameterised particle, and \[(\nabla\ell_i)_\theta:\R\to T^*_\theta\Theta\] is the \(\theta\)-fibre of the fibre part of the chosen morphism \(\hat{\ell}_i=(\ell_i,\nabla\ell_i):T^*\Theta\to\TR\) representing the parameterised log-weight. Thus an atom cotangent is sent backwards through the moving particle \(z_i\), while a log-weight cotangent is sent backwards through \(\ell_i=\log w_i\).

Finally, let \[\widehat f=(f,\nabla f): T^*\Theta\times T^*Z\to\TR\] be a parameter-dependent test morphism, with base part \[(\theta,z)\longmapsto f(\theta,z).\] We write the \((\theta,z)\)-fibre of its fibre part as \[(\nabla f)_{(\theta,z)}(t) = \tuple{ (\nabla f)^\Theta_{(\theta,z)}(t), (\nabla f)^Z_{(\theta,z)}(t) },\] where \[(\nabla f)^\Theta_{(\theta,z)}(t)\in T^*_\theta\Theta, \qquad (\nabla f)^Z_{(\theta,z)}(t)\in T^*_zZ.\] Consider the composite \[T^*\Theta \xrightarrow{\hat{\underline q}} M^{\cat[]{F}}(T^*\Theta\times T^*Z) \xrightarrow{M^{\cat[]{F}}\widehat f} M^{\cat[]{F}}\TR \xrightarrow{\opn{E}^{\cat[]{F}}_{\R}} \TR.\] By the Frobenius identity, its base part is \[\theta \longmapsto \sum_{i\in I} w_i(\theta)f(\theta,z_i(\theta)).\]

Applying Equation 104 to \(\widehat f:T^*\Theta\times T^*Z\to\TR\), an output cotangent \(t\in\R\) is first sent to the following cotangent of the FAD over \(\Theta\times Z\): \[\sum_{i\in I} \left( w_i(\theta) \tuple{ (\nabla f)^\Theta_{(\theta,z_i(\theta))}(t), (\nabla f)^Z_{(\theta,z_i(\theta))}(t) } \cdot\mathbf{e}_{(\theta,z_i(\theta))} + w_i(\theta)f(\theta,z_i(\theta))t \cdot\mathbf{\varepsilon}_{(\theta,z_i(\theta))} \right).\] Composing with Equation 105 therefore gives \[\label{eq:appendix-three-term-weighted-particle-gradient} \begin{align} & \left( \nabla\!\left( \opn{E}^{\cat[]{F}}_{\R} \circ M^{\cat[]{F}}\widehat f \circ \hat{\underline q} \right) \right)_\theta(t) \\ &\qquad = \sum_{i\in I} w_i(\theta)(\nabla f)^\Theta_{(\theta,z_i(\theta))}(t) + \sum_{i\in I} w_i(\theta) (\nabla z_i)_\theta \left( (\nabla f)^Z_{(\theta,z_i(\theta))}(t) \right) \\ &\qquad\qquad + \sum_{i\in I} (\nabla\ell_i)_\theta \left( w_i(\theta)f(\theta,z_i(\theta))t \right). \end{align}\tag{106}\]

The three summands in Equation 106 correspond to three different cotangent routes. The first term comes from the \(T^*_\theta\Theta\)-component of the test morphism \(\widehat f:T^*\Theta\times T^*Z\to\TR\), and represents the explicit dependence of \(f\) on the parameter. The second term comes from the \[T^*_{z_i(\theta)}Z\] part of the atom component inside the fibre of \(M^{\cat[]{F}}(T^*\Theta\times T^*Z)\), and represents how moving the particle \(z_i(\theta)\) changes the value of the test function. The third term comes from the real component \[\R\] inside the same fibre, and represents how changing the log-weight \(\ell_i(\theta)=\log w_i(\theta)\) changes the weighted sum. This third contribution is the usual score, or likelihood-ratio, term.

If the weights are already normalised, \(w_i(\theta)=p_i(\theta)\), and we write \[\ell_i(\theta)=\log p_i(\theta),\] then Equation 106 gives \[\begin{align} & \left( \nabla (\theta\mapsto \mathbb{E}_{i\sim p_\theta} \left[ f(\theta,z_i(\theta)) \right]) \right)_\theta(t) \\ &\qquad = \mathbb{E}_{i\sim p_\theta} \left[ (\nabla f)^\Theta_{(\theta,z_i(\theta))}(t) + (\nabla z_i)_\theta \left( (\nabla f)^Z_{(\theta,z_i(\theta))}(t) \right) + f(\theta,z_i(\theta))(\nabla\ell_i)_\theta(t) \right]. \end{align}\]

If normalised probabilities are instead obtained from unnormalised weights by a separate normalisation map, let \[p_i(\theta)=\frac{w_i(\theta)}{\sum_{j\in I}w_j(\theta)}, \qquad \ell_i(\theta)=\log w_i(\theta).\] Then the lifted fibre map of that normalisation gives \[\begin{align} & \left( \nabla (\theta\mapsto \mathbb{E}_{i\sim p_\theta} \left[ f(\theta,z_i(\theta)) \right]) \right)_\theta(t) \\ &\qquad = \mathbb{E}_{i\sim p_\theta} \left[ (\nabla f)^\Theta_{(\theta,z_i(\theta))}(t) + (\nabla z_i)_\theta \left( (\nabla f)^Z_{(\theta,z_i(\theta))}(t) \right) \right. \\ &\qquad\qquad\left. + \left( f(\theta,z_i(\theta)) - \mathbb{E}_{j\sim p_\theta} [f(\theta,z_j(\theta))] \right) (\nabla\ell_i)_\theta(t) \right]. \end{align}\]

14 Algebras and strong morphisms into the continuation monad↩︎

We give a proof of Lemma 1. Throughout this appendix, let \[K := \mathsf{K}{V}^{\cat[]{X}} = \inthom{\inthom{\mathord{-}}{V}}{V}.\] We use the internal language of the cartesian closed category \(\cat[]{X}\). Thus the elementwise formulas below denote the corresponding morphism equalities obtained by currying and evaluation.

Recall first the continuation monad structure. For a morphism \(f:A\to B\), the functorial action of \(K\) is contravariant precomposition twice: \[\Fun{K}{f}(c)(\ell)=c(\ell\circ f), \qquad c:\Fun{K}{A}, \quad \ell:\inthom{B}{V}.\] Its unit and multiplication are \[\Nat{\eta^{K}}{A}(x)(k)=k(x), \qquad x:A, \quad k:\inthom{A}{V},\] and \[\Nat{\mu^{K}}{A}(\Phi)(k) = \Phi(c\mapsto c(k)), \qquad \Phi:\Fun{K}{\Fun{K}{A}}, \quad k:\inthom{A}{V}.\] The canonical strength of \(K\) is \[\Nat{\operatorname{st}^{K}}{A,B}(x,c)(h) = c(b\mapsto h(x,b)),\] where \(x:A\), \(c:\Fun{K}{B}\), and \(h:\inthom{A\times B}{V}\).

A \(T\)-algebra structure on \(V\) is a morphism \[a:\Fun{T}{V}\to V\] satisfying \[a\circ\Nat{\eta^{T}}{V}=\mathrm{id}_{V} \qquad \text{and} \qquad a\circ\Nat{\mu^{T}}{V} = a\circ\Fun{T}{a}.\]

Given such an algebra \(a\), define a natural transformation \[\theta^{a}:T\to K\] as follows. Its component at \(A\) is the transpose of the composite \[\Fun{T}{A}\times\inthom{A}{V} \xrightarrow{\simeq} \inthom{A}{V}\times\Fun{T}{A} \xrightarrow{\Nat{\operatorname{st}^{T}}{\inthom{A}{V},A}} \Fun[np]{T}{\inthom{A}{V}\times A} \xrightarrow{\Fun{T}{\operatorname{ev}_{A,V}}} \Fun{T}{V} \xrightarrow{a} V,\] where \(\operatorname{ev}_{A,V}:\inthom{A}{V}\times A\to V\) is evaluation. Equivalently, \[\Nat{\theta^{a}}{A}(t)(k) = a(\Fun{T}{k}(t)), \qquad t:\Fun{T}{A}, \quad k:\inthom{A}{V}.\] Naturality in \(A\) follows immediately from functoriality of \(T\): for \(f:A\to B\) and \(\ell:\inthom{B}{V}\), \[\Nat{\theta^{a}}{B}(\Fun{T}{f}(t))(\ell) = a(\Fun{T}{\ell\circ f}(t)) = \Fun{K}{f}(\Nat{\theta^{a}}{A}(t))(\ell).\]

The unit law for a monad morphism is \[\begin{align} \Nat{\theta^{a}}{A}(\Nat{\eta^{T}}{A}(x))(k) &= a(\Nat{\eta^{T}}{V}(k(x))) \\ &= k(x) \\ &= \Nat{\eta^{K}}{A}(x)(k). \end{align}\] For multiplication, let \(u:\Fun{T}{\Fun{T}{A}}\). Then \[\begin{align} \Nat{\theta^{a}}{A}(\Nat{\mu^{T}}{A}(u))(k) &= a(\Fun{T}{k}(\Nat{\mu^{T}}{A}(u))) \\ &= a(\Nat{\mu^{T}}{V}(\Fun{T}{\Fun{T}{k}}(u))) \\ &= a(\Fun{T}{a}(\Fun{T}{\Fun{T}{k}}(u))) \\ &= a(\Fun{T}{(a\circ\Fun{T}{k})}(u)) \\ &= \Nat{\theta^{a}}{\Fun{T}{A}}(u) (s\mapsto \Nat{\theta^{a}}{A}(s)(k)) \\ &= \bigl( \Nat{\mu^{K}}{A} \circ \Fun{K}{\Nat{\theta^{a}}{A}} \circ \Nat{\theta^{a}}{\Fun{T}{A}} \bigr)(u)(k). \end{align}\] Thus \(\theta^{a}\) is a monad morphism.

It remains to check that it is strong. Let \(x:A\), \(t:\Fun{T}{B}\), and \(h:\inthom{A\times B}{V}\). Using the coherence of the strength of \(T\), applying \(h\) after adding the pure parameter \(x\) to the computation \(t\) is the same as applying the function \(b\mapsto h(x,b)\) to the computation \(t\). Hence \[\begin{align} \Nat{\theta^{a}}{A\times B} (\Nat{\operatorname{st}^{T}}{A,B}(x,t))(h) &= a(\Fun{T}{h}(\Nat{\operatorname{st}^{T}}{A,B}(x,t))) \\ &= a(\Fun{T}{(b\mapsto h(x,b))}(t)) \\ &= \Nat{\theta^{a}}{B}(t)(b\mapsto h(x,b)) \\ &= \Nat{\operatorname{st}^{K}}{A,B} (x,\Nat{\theta^{a}}{B}(t))(h). \end{align}\] Therefore \(\theta^{a}\) preserves strength.

Conversely, suppose that \[\theta:T\to K\] is a strong monad morphism. Define \[a^{\theta}:\Fun{T}{V}\to V\] by evaluating \(\Nat{\theta}{V}\) at the identity continuation: \[a^{\theta} := \operatorname{ev}_{\mathrm{id}_{V}}\circ\Nat{\theta}{V}.\] In internal notation, \[a^{\theta}(t)=\Nat{\theta}{V}(t)(\mathrm{id}_{V}).\] Here \(\operatorname{ev}_{\mathrm{id}_{V}}\) denotes evaluation at the point \(\mathrm{id}_{V}:1\to\inthom{V}{V}\) corresponding to the identity morphism on \(V\).

The algebra unit law follows from the unit law for the monad morphism \(\theta\): \[\begin{align} a^{\theta}(\Nat{\eta^{T}}{V}(v)) &= \Nat{\theta}{V}(\Nat{\eta^{T}}{V}(v))(\mathrm{id}_{V}) \\ &= \Nat{\eta^{K}}{V}(v)(\mathrm{id}_{V}) \\ &= v. \end{align}\] For multiplication, let \(u:\Fun{T}{\Fun{T}{V}}\). Then \[\begin{align} a^{\theta}(\Nat{\mu^{T}}{V}(u)) &= \Nat{\theta}{V}(\Nat{\mu^{T}}{V}(u))(\mathrm{id}_{V}) \\ &= \Nat{\mu^{K}}{V} \bigl( \Fun{K}{\Nat{\theta}{V}} (\Nat{\theta}{\Fun{T}{V}}(u)) \bigr)(\mathrm{id}_{V}) \\ &= \Nat{\theta}{\Fun{T}{V}}(u) (s\mapsto \Nat{\theta}{V}(s)(\mathrm{id}_{V})) \\ &= \Nat{\theta}{\Fun{T}{V}}(u)(a^{\theta}) \\ &= \Nat{\theta}{V}(\Fun{T}{a^{\theta}}(u))(\mathrm{id}_{V}) \\ &= a^{\theta}(\Fun{T}{a^{\theta}}(u)). \end{align}\] Thus \(a^{\theta}\) is a \(T\)-algebra structure on \(V\).

Finally, the two constructions are inverse. Starting from an algebra \(a\), we have \[a^{\theta^{a}}(t) = \Nat{\theta^{a}}{V}(t)(\mathrm{id}_{V}) = a(\Fun{T}{\mathrm{id}_{V}}(t)) = a(t).\] Starting from a strong monad morphism \(\theta\), let \(t:\Fun{T}{A}\) and \(k:\inthom{A}{V}\). By the definition of \(\theta^{a^{\theta}}\), by naturality of \(\theta\), and then by preservation of strength, we obtain \[\begin{align} \Nat{\theta^{a^{\theta}}}{A}(t)(k) &= a^{\theta} \left( \Fun{T}{\operatorname{ev}_{A,V}} \bigl( \Nat{\operatorname{st}^{T}}{\inthom{A}{V},A}(k,t) \bigr) \right) \\ &= \Nat{\theta}{\inthom{A}{V}\times A} \bigl( \Nat{\operatorname{st}^{T}}{\inthom{A}{V},A}(k,t) \bigr) (\operatorname{ev}_{A,V}) \\ &= \Nat{\operatorname{st}^{K}}{\inthom{A}{V},A} (k,\Nat{\theta}{A}(t)) (\operatorname{ev}_{A,V}) \\ &= \Nat{\theta}{A}(t)(k). \end{align}\] Hence \(\theta^{a^{\theta}}=\theta\). Therefore \(T\)-algebra structures on \(V\) are in one-to-one correspondence with strong monad morphisms \(T\to\mathsf{K}{V}^{\cat[]{X}}\), as claimed.

15 Other algebraic effects↩︎

This appendix gives the generic CHAD technique for discrete-output effects including the ones mentioned in Section 6. The main body develops the case of the FAD monad in detail. Here we identify the reusable part: if an algebraic effect only has values returned to continuations through discrete output types, then the raw AST monad has a fibred polynomial lift and the CHAD proof extends uniformly as soon as we give a compatible lifted generic effects and a compatible lifted handler algebra. We assume some standard background about (many-sorted) Lawvere theories, algebraic effects, and continuation monads [18], [26], [27].

The presentation is organised as follows. We first explain what we mean by discrete-output effects and we specify what data need to be supplied to apply CHAD to them. Then we state the source and target language extensions and give a generic CHAD code transformation and correctness theorem. We next specialise to the three concrete instances from Section 6. Finally, we demonstrate why we impose the limitation to non-discrete outputs.

15.1 Discrete-output setup and admissible data↩︎

15.1.1 Polynomial ASTs with discrete outputs↩︎

Let \(\Sigma\) be an algebraic signature with a small set of sorts. A generic effect symbol is written \[\omega:I_{\omega}\to T O_{\omega}.\] Here, we call \(I_{\omega}\) the effect parameter type or input type and \(O_{\omega}\) the effect output type.

Equivalently, with a \(\Set\) interpretation, it determines by Yoneda algebraic-operation families \[\llbracket I_{\omega}\rrbracket_{\Set}\times X^{\llbracket O_{\omega}\rrbracket_{\Set}} \to TX,\] and \[\llbracket I_{\omega}\rrbracket_{\Set}\times (TX)^{\llbracket O_{\omega}\rrbracket_{\Set}} \to TX .\] We see that \(O_{\omega}\) gets the role of the type of values returned to the continuation.

Write \[\widehat I_{\omega}=\llbracket I_{\omega}\rrbracket_{\cat[]{F}}\] for the chosen interpretation in \(\mathsf{Fam}(\Vect[o])\), with base set \(\llbracket I_{\omega}\rrbracket_{\Set}\). We allow \(I_{\omega}\) to be a genuine differentiable type in the sense that it may have non-trivial cotangent fibres. We require \(O_{\omega}\) to be a discrete type, in the sense that it has zero cotangent fibres: \[\llbracket O_{\omega}\rrbracket_{\cat[]{F}}=\Delta\llbracket O_{\omega}\rrbracket_{\Set}=\tuple{\llbracket O_{\omega}\rrbracket_{\Set},\underline{1}}.\] For the CHAD rules below we use the finitary case, so \(\llbracket O_{\omega}\rrbracket_{\Set}\) is finite; more generally the same formula applies whenever the required products and direct sums exist.

This induces, before quotienting by equations, the raw syntax tree monad with the polynomial form \[\mathsf{AST}_{\Sigma}X \simeq \mu Y.\; X\sqcup \coprod_{\omega\in\Sigma} \llbracket I_{\omega}\rrbracket_{\Set}\times Y^{\llbracket O_{\omega}\rrbracket_{\Set}} \simeq \mu Y.\; X\sqcup \coprod_{\omega\in\Sigma} \llbracket I_{\omega}\rrbracket_{\Set}\times \prod_{o\in\llbracket O_{\omega}\rrbracket_{\Set}}Y .\] The corresponding fibred polynomial lift to \(\mathsf{Fam}(\Vect[o])\), for an object \(A\in\mathsf{Fam}(\Vect[o])\) is \[\widehat{\mathsf{AST}}_{\Sigma}A \simeq \mu B.\; A\sqcup \coprod_{\omega\in\Sigma} \widehat I_{\omega}\times B^{\Delta\llbracket O_{\omega}\rrbracket_{\Set}} .\] Its base is canonically \(\mathsf{AST}_{\Sigma}(\left(A\right)_{\mathrm{set}})\). More explicitly, over a tree-node base element \(\omega(i,k)\), where \(i\in\llbracket I_{\omega}\rrbracket_{\Set}\) and \(k:\llbracket O_{\omega}\rrbracket_{\Set}\to\left(B\right)_{\mathrm{set}}\), the fibre is \[\Fun[np]{\left( \widehat{I_{\omega}} \right)_{\mathrm{fbr}}}{i} \oplus \bigoplus_{o\in\llbracket O_{\omega}\rrbracket_{\Set}} \Fun[np]{\left(B\right)_{\mathrm{fbr}}}{k(o)} .\] This is the point at which the discreteness of \(O_{\omega}\) is used: a continuation branch is indexed only by the set-level value \(o\), and the variable \(o:O_{\omega}\) contributes no cotangent of its own. The same polynomial construction applies componentwise to \(\cat[]{LR}\).

Thus continuous operation-inputs are harmless, provided they are treated as parameters with their own cotangent fibres; what matters for this direct fibred-polynomial construction is discreteness of the operation-output object \(O_{\omega}\). This covers finite choice, exception operations \(\opn{raise}:E\to T_E0\), writer operations such as \(\opn{tell}:\R[n]\to W_{\R[n]}1\), and state/reader operations over finite discrete state/reader objects. It does not directly cover continuous state, continuous reader, or continuous probability, where the canonical operations have genuine output objects with nonzero fibres. For such an output object \(\widehat O\), a continuation is a morphism \(\widehat O\to B\), not merely a function \(\llbracket O\rrbracket_{\Set}\to\left(B\right)_{\mathrm{set}}\); hence the simple set-indexed product \(B^{\Delta\llbracket O\rrbracket_{\Set}}\) above no longer describes the desired continuation object.

15.1.2 Admissible lifted data↩︎

A CHAD-admissible discrete-output presentation consists of the following data.

First, we assume a signature \(\Sigma\) as above and a set of equations \(E\) presenting a strong monad \(T\) on \(\Set\) as a quotient of the raw tree monad, \[q:\mathsf{AST}_{\Sigma}\twoheadrightarrow T,\] with \(q\) a morphism of strong monads. Equivalently, \(T\) is the monad presented by the generic effects of \(\Sigma\) modulo \(E\). (In fact, the equations \(E\) are often precisely obtained from the algebra/handler below.)

Second, we assume a chosen handler carrier type \(\nu\) and a \(T\)-algebra \[h:T\nu\to \nu .\] This algebra is what discharges the effect in real-output programs. (We can also allow a family of handler carriers and handler algebras, in which case we obtain a strong monad morphism from ASTs to a product of continuation monads, but we avoid this generality to simplify the presentation.)

Third, we assume chosen strong monad lifts \[T^{\cat[]{F}},\quad T^{\cat[]{LR}}\] over \(T\), in the sense that their projections to the underlying set semantics are naturally isomorphic to \(T\) and the isomorphisms respect unit, multiplication, and strength. For each operation symbol \(\omega\), these lifted monads are equipped with lifted generic effects, for example in \(\mathsf{Fam}(\Vect[o])\) \[\widehat\omega^{\cat[]{F}}_A: \widehat I_{\omega}\times \left(T^{\cat[]{F}}A\right)^{\Delta\llbracket O_{\omega}\rrbracket_{\Set}} \to T^{\cat[]{F}}A,\] and analogously in \(\cat[]{LR}\). Under the base projection, \(\widehat\omega^{\cat[]{F}}_A\) is required to be the set-level generic-effect constructor \[\llbracket I_{\omega}\rrbracket_{\Set}\times \left(\Fun[np]{T}{\left(A\right)_{\mathrm{set}}}\right)^{\llbracket O_{\omega}\rrbracket_{\Set}} \to \Fun[np]{T}{\left(A\right)_{\mathrm{set}}}\] induced by \(q\); the \(\cat[]{LR}\)-operation satisfies the corresponding logical-relations compatibility. The lifted operations must satisfy the equations \(E\) and be coherent with the lifted unit, multiplication, and strength.

Finally, we assume lifted handler algebras \[h^{\cat[]{F}}:T^{\cat[]{F}}\nu^{\cat[]{F}}\to \nu^{\cat[]{F}}, \qquad h^{\cat[]{LR}}:T^{\cat[]{LR}}\nu^{\cat[]{LR}}\to \nu^{\cat[]{LR}}.\] They are required to be Eilenberg–Moore algebras for the lifted monads and to project to the set-level handler algebra \(h\), with the evident logical-relations compatibility for \(h^{\cat[]{LR}}\).

These assumptions are exactly the data consumed by the CHAD proof. The proof does not use how the lifted monad was constructed. It may be obtained by a direct quotient/coimage construction, by an image construction, or directly from products and coproducts, provided the resulting operations and handler satisfy the compatibility conditions just listed.

15.2 Language extension and CHAD theorem↩︎

Fix a CHAD-admissible discrete-output presentation as in Section 15.1. The monadic type \(\opn[b]{M}\) is now interpreted by the presented monad \(T\), together with its chosen lifts \(T^{\cat[]{F}}\) and \(T^{\cat[]{LR}}\). The rules below do not construct these lifts from the set-level operations alone; they expose the backwards fibre maps of the chosen lifted generic effects and lifted handler algebras. This point is essential when an operation parameter \(I_{\omega}\) has nonzero cotangent fibres. All equations involving the added operations are interpreted modulo the equations \(E\) of the presentation and, in the target, through the corresponding lifted operations.

15.2.1 Source operations, equations, and handlers↩︎

For every generic effect symbol \(\omega:I_{\omega}\to T O_{\omega}\), the source language is extended by the operation form \[\prftree[r]{\mathrm{\small Op}\(_\omega\)}{ \Gamma\vdash i:I_{\omega} \qquad \Gamma,o:O_{\omega}\vdash k_o:\Fun[nn]{\opn[b]{M}}{\tau} }{ \Gamma\vdash \fn[n]{\opn[b]{op}_{\omega}}{i;\,o.k_o} :\Fun[nn]{\opn[b]{M}}{\tau} }.\] The binder \(o.k_o\) denotes a finite family of continuation terms indexed by \(\llbracket O_{\omega}\rrbracket_{\Set}\). Since \(O_{\omega}\) is discrete, extending the context by \(o:O_{\omega}\) changes the base context but adds no cotangent component. Write \[\omega_X(i,k) := \mu_X\bigl(T(k)(\omega(i))\bigr), \qquad k:\llbracket O_{\omega}\rrbracket_{\Set}\to TX,\] for the quotient-induced Kleisli constructor associated with the generic effect \(\omega:I_{\omega}\to T O_{\omega}\). The set-level semantics is \[\llbracket \fn[n]{\opn[b]{op}_{\omega}}{i;\,o.k_o}\rrbracket_{\Set}(\gamma) = \omega_{\llbracket \tau\rrbracket_{\Set}} \left( \llbracket i\rrbracket_{\Set}(\gamma), o\mapsto \llbracket k_o\rrbracket_{\Set}(\gamma,o) \right).\]

The source equational theory is extended in three ways. First, the equations \(E\) presenting \(T\) are imposed on the generic effects. Secondly, the operation form is algebraic, i.e. it commutes with Kleisli extension. Explicitly, for \(\Gamma,x:\tau\vdash s:\Fun[nn]{\opn[b]{M}}{\sigma}\), we impose the equation schema \[\label{eq:appendix-generic-operation-algebraicity-law} \fn[n]{\opn[b]{bind}}{ \fn[n]{\opn[b]{op}_{\omega}}{i;\,o.k_o}; x.s } \equiv \fn[n]{\opn[b]{op}_{\omega}}{ i;\,o.\fn[n]{\opn[b]{bind}}{k_o;x.s} } .\tag{107}\] The right-hand side is well typed because the variable \(o:O_{\omega}\) is in scope in the continuation branch. Finally, the handler associated with the chosen algebra \(h:T\nu\to\nu\) is added as \[\prftree[r]{\mathrm{\small Handle}\(_h\)}{ \Gamma\vdash t:\Fun[nn]{\opn[b]{M}}{\nu} }{ \Gamma\vdash \fn[n]{\opn[b]{handle}_{h}}{t}:\nu },\] and satisfies the Eilenberg–Moore algebra laws. In term notation, for \(\Gamma\vdash v:\nu\), \(\Gamma\vdash t:\Fun[nn]{\opn[b]{M}}{\sigma}\), and \(\Gamma,x:\sigma\vdash u:\Fun[nn]{\opn[b]{M}}{\nu}\), these are the equation schemas \[\label{eq:appendix-generic-handler-em-laws} \begin{align} \fn[n]{\opn[b]{handle}_{h}}{\fn[n]{\opn[b]{return}}{v}} &\equiv v, \\ \fn[n]{\opn[b]{handle}_{h}}{ \fn[n]{\opn[b]{bind}}{t;x.u} } &\equiv \fn[n]{\opn[b]{handle}_{h}}{ \fn[n]{\opn[b]{bind}}{ t;x.\fn[n]{\opn[b]{return}}{ \fn[n]{\opn[b]{handle}_{h}}{u} } } } . \end{align}\tag{108}\] These laws are exactly the assertion that \(\opn[b]{handle}_{h}\) interprets the \(T\)-algebra \(h\); the second equation is \(h\circ\mu=h\circ T h\) written in Kleisli form.

15.2.2 Linear target operations↩︎

The target language contains, for each generic effect \(\omega\), a linear constructor \[\fn[n]{\opn[b]{op}^{w}_{\omega,i;\,o.k_o}}{ u.i' \; ; \; o.z.k'_o } .\] The notation \(\underline{I_\omega}\) in the rule below denotes the fibre of the chosen parameter object \(\widehat I_\omega\); if \(I_\omega\) is discrete, this is the zero linear type. The constructor represents the reverse fibre component of the chosen lifted operation (the transposed derivative in differentiable examples) \[\widehat\omega^{\cat[]{F}}_A: \widehat I_{\omega}\times \left(T^{\cat[]{F}}A\right)^{\Delta\llbracket O_{\omega}\rrbracket_{\Set}} \to T^{\cat[]{F}}A,\] and similarly in \(\cat[]{LR}\). Thus, for \(i\in\llbracket I_\omega\rrbracket_{\Set}\) and a branch family \(p:\llbracket O_\omega\rrbracket_{\Set}\to\left(T^{\cat[]{F}}A\right)_{\mathrm{set}}\), writing \(p_o=p(o)\), its reverse fibre map is \[\Fun[np]{\left( T^{\cat[]{F}}A \right)_{\mathrm{fbr}}}{\omega(i,p)} \longrightarrow \Fun[np]{\left( \widehat{ I_\omega } \right)_{\mathrm{fbr}}}{i} \oplus \bigoplus_{o\in\llbracket O_\omega\rrbracket_{\Set}} \Fun[np]{\left( T^{\cat[]{F}}A \right)_{\mathrm{fbr}}}{p(o)},\] where \(\omega(i,p)=\mu_{\left(A\right)_{\mathrm{set}}}(T(p)(\omega(i)))\) is the induced Kleisli constructor. In one-sorted notation it has the following rule, in the same continuation-passing style as \(\opn[b]{categorical}^{w}\): \[\label{eq:appendix-generic-operation-w-typing} \begin{gather} \Gamma,y:I_{\omega}\vdash \underline{I_{\omega}} \text{ type} \qquad \Gamma,x:\tau\vdash \underline{\tau} \text{ type} \qquad \Gamma\vdash \underline{\rho} \text{ type} \\ \prftree{ \Gamma\vdash i:I_{\omega} }{ \Gamma,o:O_{\omega}\vdash k_o:\Fun[nn]{\opn[b]{M}}{\tau} }{ \Gamma;u:\underline{I_{\omega}}[i/y] \vdash i':\underline{\rho} }{ \Gamma,o:O_{\omega}; z: \opn[b]{M}^{w}_{k_o:\Fun[nn]{\opn[b]{M}}{\tau}} \left[x\mapsto \underline{\tau}\right] \vdash k'_o:\underline{\rho} }{ \Gamma; v: \opn[b]{M}^{w}_{ \fn[n]{\opn[b]{op}_{\omega}}{i;\,o.k_o} : \Fun[nn]{\opn[b]{M}}{\tau} } \left[x\mapsto \underline{\tau}\right] \vdash \fn[n]{\opn[b]{op}^{w}_{\omega,i;\,o.k_o}}{ u.i' \; ; \; o.z.k'_o } : \underline{\rho} } \end{gather}\tag{109}\] The variable \(u\) receives the cotangent contribution for the operation parameter, while \(z\) receives the cotangent contribution for the corresponding continuation branch. Since \(O_{\omega}\) is discrete, extending the context by \(o:O_{\omega}\) does not add a linear component; hence each \(k'_o\) returns a cotangent for the same ambient context \(\Gamma\). The finite direct sum above is where the finiteness of \(O_{\omega}\) enters the target syntax. Semantically, if the reverse fibre map sends the incoming monadic cotangent \(v\) to \(\tuple{u,(z_o)_{o\in\llbracket O_{\omega}\rrbracket_{\Set}}}\), then the constructor denotes the finite sum \[\fn[n]{i'}{u} + \sum_{o\in\llbracket O_{\omega}\rrbracket_{\Set}} \fn[n]{k'_o}{z_o}.\] For \(O_{\omega}=0\), this branch sum is empty.

With this operation former, the CHAD clause for algebraic operations is the following. We use the canonical identification \[\Fun[np]{\chad}{\Gamma,o:O_{\omega}}_{2} \cong \Fun[np]{\chad}{\Gamma}_{2},\] since \(O_{\omega}\) is discrete and has zero cotangent fibres (syntactically, \(\underline{O_{\omega}}\) is the terminal zero-cotangent type \(\underline{1}\)). Then \[\label{eq:appendix-generic-operation-chad} \begin{align} &\Fun[np]{\chad[\Gamma]}{ \fn[n]{\opn[b]{op}_{\omega}}{i;\,o.k_o} } := \\ &\quad \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{i}\;\opn[b]{of}\; \tuple{i_0,i'}\to \\ &\quad \opn[b]{let}\;\tuple{r_o,r'_o} = \Fun[np]{\chad[\Gamma,o:O_{\omega}]}{k_o} \;\opn[b]{in} \\ &\quad\quad \tuple{ \fn[n]{\opn[b]{op}_{\omega}}{i_0;\,o.r_o}, \lambda v\; .\; \fn[n]{\opn[b]{op}^{w}_{\omega,i_0;\,o.r_o}}{ u.\fn[n]{i'}{u} \; ; \; o.z.\fn[n]{r'_o}{z} } } . \end{align}\tag{110}\] The \(\opn[b]{let}\)-binding of \(\tuple{r_o,r'_o}\) is meta-notation for the family obtained by translating the continuation in the context \(\Gamma,o:O_{\omega}\). The incoming cotangent \(v\) is consumed by \(\opn[b]{op}^{w}_{\omega}\), which separates it into the parameter cotangent and the branch cotangents; the backpropagators \(i'\) and \(r'_o\) then return these contributions to the ambient context. No cotangent for \(o\) is omitted: there is none.

15.2.3 Linear handler operations↩︎

Similarly, the target language contains a linear constructor \(\opn[b]{handle}^{w}_{h,p}\). To avoid overloading the incoming output cotangent with the monadic cotangent produced by the handler algebra, we write this constructor with an explicit binder \(w.s\): \[\label{eq:appendix-generic-handler-w-typing} \begin{gather} \Gamma,x:\nu\vdash \underline{\nu} \text{ type} \qquad \Gamma\vdash \underline{\rho} \text{ type} \\ \prftree{ \Gamma\vdash p:\Fun[nn]{\opn[b]{M}}{\nu} }{ \Gamma; w: \opn[b]{M}^{w}_{p:\Fun[nn]{\opn[b]{M}}{\nu}} \left[x\mapsto \underline{\nu}\right] \vdash s:\underline{\rho} }{ \Gamma; v: \underline{\nu}[\fn[n]{\opn[b]{handle}_{h}}{p}/x] \vdash \fn[n]{\opn[b]{handle}^{w}_{h,p}}{w.s} :\underline{\rho} } \end{gather}\tag{111}\] Semantically, if \(B=\nu^{\cat[]{F}}\), the lifted handler algebra \(h^{\cat[]{F}}:T^{\cat[]{F}}B\to B\) has reverse fibre map \[\Fun[np]{\left(B\right)_{\mathrm{fbr}}}{h(p)} \longrightarrow \Fun[np]{\left( T^{\cat[]{F}}B \right)_{\mathrm{fbr}}}{p}.\] The constructor \(\opn[b]{handle}^{w}_{h,p}\) binds \(w\) to the image of \(v\) under this reverse fibre map and then evaluates \(s\). Thus it uses the supplied lifted algebra, not merely the underlying set map \(h\). In the FAD instance \(\nu=\srcreal[p]\), this is the same continuation-passing rule as Equation 43 ; the explicit binder only makes the two cotangent types visible.

As for FAD, we again impose the transposed derivatives of the equations for the operations and handlers from the source language, e.g., the Eilenberg–Moore laws for the handler algebra.

The CHAD clause is the exact analogue of the canonical FAD handler clause: \[\label{eq:appendix-generic-handler-chad} \Fun[np]{\chad[\Gamma]}{\fn[n]{\opn[b]{handle}_{h}}{t}} := \opn[b]{case}\;\Fun[np]{\chad[\Gamma]}{t}\;\opn[b]{of}\;\tuple{p,b}\to \tuple{ \fn[n]{\opn[b]{handle}_{h}}{p}, \lambda v\; .\; \fn[n]{\opn[b]{handle}^{w}_{h,p}}{w.\fn[n]{b}{w}} }.\tag{112}\] Here \(b\) is a backpropagator expecting a monadic cotangent over \(p\). The incoming output cotangent \(v\) is first pulled back along the lifted handler algebra to the monadic cotangent \(w\), and only then passed to \(b\).

15.2.4 The generic CHAD result↩︎

15.3 Concrete instances↩︎

The following examples instantiate the admissible data of Section 15.1. In each case the point to check is not merely that the set-level monad is algebraic, but that the chosen fibred monad, the lifted generic effects, and the lifted handler are genuine Eilenberg–Moore data over the set-level ones. We display the \(\mathsf{Fam}(\Vect[o])\)-fibre maps; the \(\cat[]{LR}\)-lifts used in the logical relations proof are obtained by the same constructions with the relation data carried componentwise.

15.3.1 Finite non-determinism via the finite multiset monad↩︎

We model multiplicity-sensitive finite non-determinism by the finite multiset monad \[\mathsf{MFin}X = \coprod_{k \in \mathbb{N}} \left(X^{k}/\mathfrak{S}_{k}\right),\] where \(\mathfrak{S}_{k}\) acts by permuting the entries of a \(k\)-tuple. Multiplicity is retained but order is forgotten. If one additionally quotients by idempotence, one obtains the finite-powerset variant of non-determinism, whose algebras are join-semilattices with bottom rather than commutative monoids. We use the multiset, not the idempotent, variant below; in particular, the additive handler \(\opn{sum}_n\) below would not factor through the idempotent quotient unless the chosen target algebra were idempotent.

The unit sends \(x\) to the singleton multiset \(\left\{\!\left\{x\right\}\!\right\}\), and bind maps each occurrence to a finite multiset and then takes multiset union. An algebra \(a:\mathsf{MFin}A\to A\) is equivalently a commutative monoid object on \(A\). Given \(a\), the unit is \(a(\varnothing)\) and multiplication is \((a_1,a_2)\mapsto a(\left\{\!\left\{a_1,a_2\right\}\!\right\})\). Conversely, a commutative monoid folds finite multisets by iterated multiplication; the algebra laws are exactly unit, associativity, and commutativity.

The generic-operation presentation may be taken to have operations \[\opn{choose}_k:1\to \mathsf{MFin}\{1,\dots,k\}, \qquad \ast\mapsto \left\{\!\left\{1,\dots,k\right\}\!\right\},\] for \(k\in\mathbb{N}\), modulo the usual unit, associativity, and permutation equations. The induced Kleisli constructor sends a finite family of branch computations to their multiset union; \(k=0\) gives the empty choice.

The quotient by permutations must also be reflected in the fibres. Let \(\tuple{X,A}\in\mathsf{Fam}(\Vect[o])\), let \(m\in\mathsf{MFin}X\), and choose a representative tuple \(\vec{x}=(x_1,\dots,x_k)\) of \(m\). The free commutative-monoid lift obtained by quotienting the lifted list monad has fibre \[\label{eq:appendix-multiset-fibre-stabiliser} \Fun[np]{\left( \mathsf{MFin}^{\cat[]{F}}\tuple{X,A} \right)_{\mathrm{fbr}}}{m} = \left( \bigoplus_{i=1}^{k}\Fun[nn]{A}{x_i} \right)^{\operatorname{Stab}(\vec{x})},\tag{113}\] where \(\operatorname{Stab}(\vec{x})\leq\mathfrak{S}_k\) is the subgroup permuting equal entries of \(\vec{x}\), acting by the corresponding symmetries of the finite direct sum. This description is independent of the representative. Equivalently, if \(m\) is written as a finite-support multiplicity function, then the fibre is canonically the direct sum over the support, \[\Fun[np]{\left( \mathsf{MFin}^{\cat[]{F}}\tuple{X,A} \right)_{\mathrm{fbr}}}{m} \cong \bigoplus_{x\in\operatorname{supp}(m)}\Fun[nn]{A}{x},\] with the quotient map from ordered lists copying the support cotangent to every indistinguishable occurrence. This last copying is essential: the tempting formula with an independent summand for every occurrence is a labelled-trace semantics before quotienting by stabilisers, not the canonical free commutative-monoid lift on the finite-multiset quotient.

For a morphism \(\tuple{f,\alpha}:\tuple{X,A}\to\tuple{Y,B}\), the base map is multiset push-forward. Under the support description, the fibre map over \(m\in\mathsf{MFin}X\) sends \[(b_y)_{y\in\operatorname{supp}(\mathsf{MFin}f(m))} \longmapsto \bigl(\Nat[nn]{\alpha}{x}(b_{f(x)})\bigr)_{x\in\operatorname{supp}(m)} .\] Thus a cotangent at a target support point is copied to all source support points that map to it, and the quotient map then copies it to the corresponding occurrences. The unit, multiplication, and strength are induced from the lifted list monad and descend because the permutation equations act by symmetries of finite direct sums.

For every \(n\), the additive commutative monoid \(\tuple{\R[n],0,+}\) gives the handler \[\opn{sum}_{n}:\mathsf{MFin}\R[n]\to\R[n], \qquad \left\{\!\left\{v_{1},\dots,v_{k}\right\}\!\right\}\mapsto \sum_{i=1}^{k} v_i .\] Its lifted fibre map is characterised by the list quotient: after pulling back to an ordered representative, it sends an incoming cotangent \(\xi\in\R[n]\) to one copy of \(\xi\) at every occurrence. Under the support identification above, this is the invariant cotangent with component \(\xi\) at every support point. Hence, when two syntactic branches produce the same base value, the same support cotangent is copied back to both branches and the surrounding context addition accounts for the multiplicity.

The source operation \(\opn[b]{choose}_k\) is therefore interpreted by multiset union of the \(k\) branch computations. Its linear operation \(\opn[b]{choose}^{w}_k\) is the reverse fibre map of the descended union map: it restricts the incoming invariant multiset cotangent to the branch multisets, copying it along collisions. The handler clause for \(\opn[b]{fold}_{\nu}\), and in particular \(\opn[b]{sum}_n\), is the generic algebra clause using the lifted commutative-monoid fold. These data satisfy the commutative-monoid equations in \(\mathsf{Fam}(\Vect[o])\), so the generic CHAD theorem applies.

15.3.2 Exception monads↩︎

Fix an exception object \(\widehat E=\tuple{E,B}\) in \(\mathsf{Fam}(\Vect[o])\); for a purely discrete exception set one takes \(B_e=0\). The underlying set-level exception monad is \[T_E X=X\sqcup E .\] Its signature has one generic effect \[\opn{raise}:E\to T_E0,\] equivalently one nullary operation for each point of \(E\). Since the output type is \(0\), this is a discrete-output effect even when the parameter object \(\widehat E\) has nonzero cotangent fibres. The raw AST monad is already \[\mathsf{AST}_E X \simeq \mu Y.\;X\sqcup E\times Y^0 \simeq X\sqcup E,\] so there is no non-trivial quotient to impose.

A \(T_E\)-algebra on a set \(Y\) is a map \(a:Y\sqcup E\to Y\) whose restriction to \(Y\) is the identity. Equivalently it is a recovery map \[h_0:E\to Y, \qquad a=[\mathrm{id}_Y,h_0].\] For instance, taking \(Y=A\sqcup E'\) allows a handler to recover some exceptions as values in \(A\) and translate others into values tagged by a new exception set \(E'\); it is still an algebra for the original monad \(T_E\) on the carrier \(Y\).

The exception monad lifts by the coproduct formula \[T_E^{\cat[]{F}}\tuple{X,A} = \tuple{ X\sqcup E, z\mapsto \begin{cases} \Fun[nn]{A}{x}, & z=\opn{inl}(x),\\ \Fun[nn]{B}{e}, & z=\opn{inr}(e). \end{cases} }.\] The lifted unit, multiplication, and strength are the coproduct maps induced by those of \(T_E\). The lifted \(\opn{raise}\) operation \(\widehat E\to T_E^{\cat[]{F}}\tuple{X,A}\) has reverse fibre map the identity \[\Fun[nn]{B}{e}\longrightarrow \Fun[nn]{B}{e},\] because there are no continuation branches.

Let \(\tuple{Y,C}\) be a lifted carrier and let \(\widehat h:\widehat E\to\tuple{Y,C}\) be a morphism over the recovery map \(h_0:E\to Y\). The lifted algebra over \([\mathrm{id}_Y,h_0]\) is case analysis. Its reverse fibre map sends \(\gamma\in\Fun[np]{C}{y}\) over a returned value \(\opn{inl}(y)\) to \(\gamma\), while over an exception \(\opn{inr}(e)\) it sends \(\gamma\in\Fun[np]{C}{h_0(e)}\) to the fibre map of \(\widehat h\), \[\left( \left( \widehat{h} \right)_{\mathrm{fbr}} \right)_e^{\ast}(\gamma) \in \Fun[np]{B}{e} .\] When \(h_0\) is an ordinary differentiable recovery map and \(\widehat h\) is its cotangent lift, this fibre map is \((Dh_0(e))^{\mathsf T}\gamma\). The Eilenberg–Moore laws are immediate from \([\mathrm{id}_Y,h_0]\circ\eta=\mathrm{id}_Y\) and from the fact that multiplication for \(X\sqcup E\) is case analysis that leaves the first encountered exception unchanged. Thus exceptions fit the admissible pattern directly; no quotient is needed for the direct coproduct lift. The free-algebra/product-continuation construction also recovers \(T_E\) because free exception algebras separate returned values and exception labels, but a single chosen recovery handler need not be separating if it identifies exceptions or misses some cotangent direction.

15.3.3 Writer monads and \(\R[n]\)-accumulation↩︎

Let \(M=\tuple{M,e,\cdot}\) be a monoid. We use the right-action convention for the monad on \(X\times M\): bind sends \(\tuple{x,m}\), together with \(f:X\to Y\times M\), to \(\tuple{y,n\cdot m}\) when \(\fn[n]{f}{x}=\tuple{y,n}\). Equivalently, this is the usual writer monad for the opposite monoid; for the additive monoids \(\R[n]\) used below, the two conventions coincide. With this convention, an Eilenberg–Moore algebra for \(W_M\) is precisely a right \(M\)-action \[a:A\times M\to A\] with \(a(a_0,e)=a_0\) and \(a(a(a_0,m),n)=a(a_0,m\cdot n)\).

The generic effect is \[\opn{tell}:M\to W_M1, \qquad r\mapsto\tuple{\ast,r} .\] The induced Kleisli constructor sends a parameter \(r\in M\) and a branch \(k(\ast)=\tuple{x,s}\in W_MX\) to \(\tuple{x,s\cdot r}\). Hence the operation-output type is \(1\), so the writer effect is in the discrete-output fragment even when the parameter object \(M\) is continuous.

For the accumulation example take \(M=\tuple{\R[n],0,+}\) with its standard cotangent bundle. The lifted writer monad maps \[\tuple{X,A} \mapsto \tuple{ X\times\R[n], \left(\tuple{x,r}\mapsto \Fun[nn]{A}{x}\oplus\R[n]\right) }.\] For a morphism \(\tuple{f,\alpha}:\tuple{X,A}\to\tuple{Y,B}\), the base map is \(\tuple{x,r}\mapsto\tuple{\fn[n]{f}{x},r}\), and the reverse fibre map is \[\Fun[np]{B}{\fn[n]{f}{x}}\oplus\R[n] \longrightarrow \Fun[nn]{A}{x}\oplus\R[n], \qquad \tuple{\beta,\rho}\mapsto\tuple{\Nat[nn]{\alpha}{x}(\beta),\rho}.\] The lifted multiplication is the transpose derivative of \[\tuple{\tuple{x,s},r}\mapsto\tuple{x,s+r},\] so an incoming cotangent \(\tuple{x^{\ast},u^{\ast}}\) is sent to \(\tuple{\tuple{x^{\ast},u^{\ast}},u^{\ast}}\). This proves the lifted monad laws because the displayed maps are the transpose derivatives of the unit, associative addition, and strength maps.

For the lifted \(\opn{tell}\) operation, if the branch computation has base \(\tuple{x,s}\), then the operation has base \(\tuple{x,s+r}\). Its reverse fibre map sends an incoming cotangent \(\tuple{x^{\ast},u^{\ast}}\) to the parameter cotangent \(u^{\ast}\) and the branch cotangent \(\tuple{x^{\ast},u^{\ast}}\). Thus the newly written value and the writer value produced by the continuation both receive the same accumulator cotangent.

For the commutative additive accumulation example, the useful carrier is \(\sigma\times\R[n]\) with action \[\label{eq:canonical-accumulation-action} (\sigma\times\R[n])\times\R[n] \to \sigma\times\R[n], \qquad \tuple{\tuple{a,u},r}\mapsto\tuple{a,u+r}.\tag{114}\] If \(\sigma\) is lifted with fibre \(A_a\), the transposed derivative of Equation 114 sends a cotangent \(\tuple{a^{\ast},u^{\ast}}\) at \(\tuple{a,u+r}\) to \[\tuple{\tuple{a^{\ast},u^{\ast}},u^{\ast}} \in \left(A_a\oplus\R[n]\right)\oplus\R[n],\] where the two copies of \(u^{\ast}\) correspond to the old accumulated value and the newly written value. The action laws are exactly the Eilenberg–Moore algebra laws for the writer monad, so this gives the lifted handler used by \(\opn[b]{run}_{\tau}\) and its linear operation \(\opn[b]{run}^{w}_{\tau,p}\). This primitive handles computations whose returned values already lie in \(\tau\times\srcreal[n]\); the usual \(\opn{runWriter}:\opn[b]{M}\tau\to\tau\times\srcreal[n]\) is obtained by first mapping \(a:\tau\) to \(\tuple{a,0}\) and then applying this algebra. For a general differentiable monoid object, the same statements hold with addition replaced by the transpose derivative of the chosen multiplication-order map.

15.4 Scope and non-discrete output limitations↩︎

The constructions above are deliberately discrete-output constructions. They allow differentiable operation-input objects, such as continuous exception labels or continuous writer values, because these are parameters \(I_{\omega}\) and their cotangents are accounted for at the operation node. They also allow arbitrary zero-fibre output indices \(\Delta O_{\omega}\), provided the required set-indexed products and direct sums exist; in the finitary CHAD theorem these output sets are finite.

They do not cover genuine operation-output objects with nonzero cotangent fibres. For such an object \(\widehat O_{\omega}\), a continuation is a morphism \(\widehat O_{\omega}\to B\), not just a function \(P\widehat O_{\omega}\to PB\), so its base point can contain reverse-linear fibre data. If syntax or handlers can inspect that data, the base of the lifted construction need not be the intended set-level monad. The obstruction remains even if genuine powers \(B^{\widehat O_{\omega}}\) are supplied.

References↩︎

---

1. A normalised barycentre would require division by the total mass. Here we deliberately omit this normalisation.↩︎

2. As in \(\Set\), one may instead take \(\TR[p]\) for some \(p\in\mathbb{N}\).↩︎

3. For the description of the source language and its interpretation requirements, see Section 3.1.↩︎

4. For the corresponding \(p\)-dimensional barycentre semimodule, one takes \(\tuple{ \opn{Diff} \left( \mathbb{E},\R[p]\right), \tuple{\R[p],\TR[p]},\operatorname{dv}_{\R[p]} }\).↩︎

5. This assumes that the theory has product types. Without product types, one can instead work with the usual construction based on equivalence classes of contexts.↩︎