

* Verified mathematical code rewriting

  We introduce a library capable of mathematically sound source code transformations with a focus on automatic and symbolic differentiation.


* Symbolic differentiation

  We will start with symbolic differentiation. There difference between symbolic and automatic differentiation is very small and blurry. In this section we will not worry about generating efficient code, let bindings and focus only on computing a derivative of a function in a given direction.
  

  The core code transformation is computation of differential. Give a program/function $f : X \rightarrow Y$ and a point $x : X$ and direction $v : X$, compute the differential $\partial f(x;v)$. Given by a limit
  $$
  \partial f(x;v) = \lim_{h \rightarrow 0} \frac{f(x + t v) - f(x)}{h}
  $$
  we the Lean notation $\partial$ is a function of type $(X \rightarrow Y) \rightarrow X \rightarrow X \rightarrow Y$.

  Taking derivative makes sense only for differentiable functions. However, it would be quite cumbersome to demand a proof of differentiability of $f$ each time we write $\partial f$. Thus we take a similar approach that is taken with division in Mathlib. You can write $x / y$ for any two real numbers $x,y$ even for $y$ equal to zero. The proof of $y\neq 0$ is only required when some rewriting rules are applied. For example the rule $(x / y) * y = x$ can be applied only if $y$ is non-zero. We take a similar approach with the differential $\partial$. You can apply it to any function but rules like $\partial (f + g) = \partial f + \partial g$ are applicable only if $f$ and $g$ are differentiable.

  To perform automatic of symbolic differentiation in Lean we need two mechanisms. Rewrite rules like $\partial (f + g) = \partial f + \partial g$ and a way to prove that $f$ and $g$ are differentiable functions.

  We use Lean's inbuilt simplifier to do the rewriting. An example of the mentioned rule is

  $$
  @[simp]
  theorem differential_of_add {X Y Z : Type} [Vec X] [Vec Y] [Vec Z]
    (f : Y -> Z) [IsSmooth f] (g : X -> Y) [IsSmooth g]
    : ∂ (f + g) = ∂ f + ∂ g := ...
  $$
  This rule states that for any types $X, Y, Z$ that are vector spaces, $[Vec X], and for any two function $f, g$ that are smooth, $[IsSmooth f]$, the statement after colon is true. In this case $∂ (f + g) = ∂ f + ∂ g$. The proof of this statement goes to the place of $...$.

  There are three kinds of argument brackets. The normal bracket $( )$ indicate explicit argument, ${ }$ implicit arguments and $[ ]$ type class arguments. For readers unfamiliar with typeclasses, you can think about them as implicit arguments but with a way more powerful inference rules. Most importantly, these rules can be user specified by providing so called typeclass instances as you will see shortly.

  Traditionally type classes are used to do ad-hoc polymorphism in more principled way. For example in this case the type class $[Vec X]$ fetches implementation of addition $. + . : X -> X -> X$ and scalar multiplication $. * . : R -> X -> X$ on the type $X$. But in Lean it also provides proofs that $X$ actually forms proper vector space\footnote{Not completely true because of rounding errors.}. The typeclass $[IsSmooth f]$ does not carry any data but only a proof that the function $f$ is indeed smooth. The most important rule is that composition of smooth functions is smooth, this is expressed as the following instance:
  $$
  instance composition_is_smooth {X Y Z : Type} [Vec X] [Vec Y] [Vec Z]
    (f : Y -> Z) [IsSmooth f] (g : X -> Y) [IsSmooth g]
    : IsSmooth (fun x => f (g x)) := ...
  $$
  

* Differentiation of lambda functions and SKI combinators

  Any lambda function can be rewritten in terms of SKI combinators. The identity combinator $I x = x$, the constant combinator $K x y = x$ and the generalized composition combinator $S f g x = f x (g x)$. In some way any program is just one lambda function thus to perform a source code transformation we need to know how does the transformation act on these combinators.

  In the case of differentiation and smoothness we are in luck as it plays nice with all these three combinators. As mentioned earlier, there are two main components to our source code transformations. The simplification rules and rules to prove function smoothness.

  The smoothness rules are as follows
  $$
  instance I.arg_x.isSmooth 
    : IsSmooth λ x : X => x := ...
    
  instance K.arg_y.isSmooth (x : X) 
    : IsSmooth λ y : Y => x := ...

  instance S.arg_x.isSmooth 
    (f : X → Y → Z) [IsSmooth f] [∀ x, IsSmooth (f x)] 
    (g : X → Y) [IsSmooth g] :
    IsSmooth (λ x => f x (g x)) := ...
  $$

  The first two instance just claims that the identity function is smooth and the constant function, i.e. function taking an argument $y$ and always returning $x$, is smooth.

  The only confusing part is how the smoothness of $f$ is expressed. We will not go into the details here but just in words the $[IsSmooth f]$ states that the function $f$ is as a function from $X$ to $Y -> Z$ is smooth and $[∀ x, IsSmooth (f x)]$ states that for all $x$ the function $(f x)$ from $Y$ to $Z$. One might be tempted to reformulate the first one as $[∀ y, IsSmooth (fun x => f x y)]$ but this does not work for some quite subtle reasons coming from the underlining functional analysis and will be discussed in the following optional section.

  The corresponding simplification rules are:

  $$
  theorem id.arg_x.diff_simp
    : ∂ (λ x : X => x) = λ x dx => dx := ...

  theorem const.arg_y.diff_simp (x : X)
    : ∂ (λ y : Y => x) = λ y dy => (0 : X) := ...
    
  theorem scomb.arg_y.diff_simp (x : X)
    (f : X → Y → Z) [IsSmooth f] [∀ x, IsSmooth (f x)] 
    (g : X → Y) [IsSmooth g] :
    : ∂ (λ x => f x (g x)) = λ x dx => ∂ f x dx (g x) + ∂ (f x) (g x) (∂ g x dx) := ...
  $$

  The differential of identity function is identity function, this is just a special case of the fact that differential of linear function it that linear function itself. The next one is that differential of a constant function is zero.
  The last one is a mix of chain rule and product rule. Chain rule is obtained if $f$ does not depend on the $x$ argument. The product rule if $f x y = (h x) * y$ for some function h, then $∂ (λ x => (h x)*(g x)) = (∂ h x dx) * (g x) + (h x) * (∂ g x dx)$.


  !!! THIS NEEDS MORE INVESTIGATION !!! I STILL DO NOT UNDERSTAND THE ROLE OF SWAP AND PARM THEOREMS !!!
  There is one additional rule that allows us to work with higher order functions and that is swapping of arguments.
  $$
  instance (priority := low) swap.arg_x.isSmooth 
    (f : α → X → Y) [∀ a, IsSmooth (f a)] 
    : IsSmooth (λ x a => f a x) := ...
  $$
  This say if we have a function parameterized by a parameter $a$ and for every parameter $a$ the function is smooth. Then the function taking $x : X$ and producing a family of functions $α → Y$ is smooth. In some sense, we can delay the choice of the parameter $a$ and still have smooth function.
  This instance is crucial but fairly dangerous as it can introduce new function arguments and lead the type class inference astray. Therefore it has lower priority then all other instances.
  Relate this to infinite products! 

  All of the source code transformations in this work follow very similar pattern. They introduce a class of types, like vector space, finite type etc, function property, like smoothness, linearity, invertibility etc., that is used as a condition for code transformation and finally the code transformation itself, like differentiation, adjunction, inversion etc.

  The code transformations differ slightly in how they handle individual combinators. For example, when dealing with linearity the theorem for S combinator is very restrictive on the function $f$ or the theorem about constant functions does hold at all, i.e. constant function is not linear. We will talk about this later.
  
This will not be the case for example when dealing with linearity where bilinearity of $f$ and linearity of $g$ does not imply linearity of $S f g$.

* Differentiation of curry and uncurry

  Another way of looking at the source code transformations we are interested in is too investigate how they interact with currying.

  Look at curry and uncurry. Product projections and product fmap.


* Technical problems

  There are few technical issues we have to deal with. The hardest problem is high order unification. I.e. given a lambda function, figuring you if it matches for example the pattern $fun x => f (g x)$. and a theorems about function composition can be used.

  *Trailing arguments*

  *Breaking simplifier loop*


* Function analytical and categorical view

Let's look at the math behind automatic differentiation more carefully. We are interested in doing differentiation with higher order functions. For example differentiate function composition with respect to the inner function, $∂ fun g => f (g x)$. Therefore it we can suffice with a simple definition of derivative of a function between two Eucledian spaces. The most common generalization is Frechet or Gatoux derivative between Banach spaces but the problem is that the space of differentiable function between Banach spaces is not a Banach space. Therefore we work with Convenient Vector Spaces, these are locally convex vector spaces with some mild completeness condition. You can define a notion of smooth map between two convenient vector spaces and the space of all these maps form a convenient vector space. The general property we are after is that the category of convenient vector spaces form a cartesian closed category. In practice, this means you can use the language of lambda function to form new functions.

The work of ... propose to use diffeological spaces as the semantics of differentiable programming. We are definitely planning on doing so, but for the sake of simplicity we have decided to work only with vector spaces. And convenient vector spaces are vector spaces with standard diffeology.

A good example of diffeological spaces we would like to work with is sum type or space of all diffeomorphisms. They do not have a vector space structure. For example the sum type =X ⊕ Y= has tangent space =X= at every point =inl x= and tangent space =Y= at every point =inr y=. Or clearly sum of two diffeomorphisms does not have to be a diffeormoprhism.

For example Lean is using internally sum type to implement for loops. And diffeomorphisms naturally show up when you are dealing PDEs describing fluid flow or elastic materials.


- point out functorial nature

* Bundled functions

  Introduce =X ⟿ Y= and =λ x ⟿ f x= notation. or =X -o Y= 

* Interactive and automatic rewriting

  Exploit the interactive environment for proofs. Introduce AutoImpl struct and conv mode and simple =rewrite_by= notation. Function annotations.

* Do we really need proofs?
  
  Do we really need those smoothness proofs? Can't we just apply the rewrite rule without worrying about differentiability of $f$ and $g$? This is a common question from people with not experience with formal methods as they do not see the benefit for the extra hassle. Honestly, we are unsure if the proofs are really necessary here. The question that needs to be answered is: do we get into trouble if we omit all the proofs and just apply the rewrite rules whenever applicable? An example of a rule that definitely needs a proof is $A (x + y) = A x + A y$ for $A$ linear map. This rule can be only applied for linear maps $A$ and would lead to incorrect results.

 For differentiation the situation seems to be different. Lean can print out all the simplification rules that were not applied because a function failed to be differentiable and in our limited test cases there were no such rewrite rules. On the other hand, having proofs that a type is a vector space is important and many rules are not applied because of this. This does not mean that in future we will introduce a rewrite rule that will really need the proof of differentiability. Right now, we always synthesize differentiability proofs to have a formal guarantee about correctness. In future, we can easily short circuit all these proofs by providing an instance that all function are differentiable and gain some performance. This would of course invalidate all our correctness guarantees but one can do this only while developing to gain performance and turn off this short circuiting for the final release to properly check correctness.

The situation is different with other source code transformations like function inversion or transposition/adjunction. There are definitely cases where omitting proof of invertibility or linearity of a function would lead to incorrect results. 

* Gradients and transposition

  Automatic differentiation is commonly used for computing gradients of a function and used in algorithms like gradient descent to minimize a cost function. For scalar valued function $f : R^n -> R$ the gradient function $∇ f : R^n -> R^n$ is defined as
  $$
  (∇ f)_i = ∂ f / ∂ x_i
  $$
  Or in the Lean notation with differential $(∇ f)_i = ∂ f e_i$ where $e_i$ is the canonical basis of $R^n$. Not so commonly stressed assumption is that the definition of the gradient really depends on the choice of the basis, more generally on the choice of the inner product on $R^n$. The gradient can be formed from the differential using adjunction or transposition. We use the term adjunction as this is the term more commonly used term when working with infinitely dimensional vector spaces.

  If we have a linear function $f : X -> Y$ between two spaces with inner product. Its adjoint is defined as function $f†$ satisfying:
  $$
  <f x, y> = <x, f† y>    forall x y
  $$

  Using this notation for adjunction/transposition the gradient is
  $$
  ∇ f x = (∂ f x)† 1
  $$
  i.e. the gradient of function $f$ at point $x$ is the evaluation at 1 of the adjoint differential at point $x$.

  Now the problem is that we want to work with higher order functions. For example we would like to compute the gradient of a function with type $(R ~~> R) -> R$. This comes a lot in physics or in the study of partial differential equations. For example, the gradient of the Dirichlet energy $∫ |∇ u|²$ is the Laplacian equation $Δ u$ or the gradient of action integral $∫ L(x,x',t)$ is the Euler-Lagrange equation ... . Or compute the adjoint for a one dimensional Sobel filter $F u_i = u_(i+1) - u_i$ where $u_i : Z -> R$.

  We will not go to details here as full explanation requires some familiarity with functional analysis and distribution and it is given in the next optional subsection. Let's give a summary here. We define Semi-Hilbert spaces and for those we define the notion of function adjunction. If SemiHilbert space is just a vector space with inner product, adjunction is the well known transposition. However, spaces like $Z -> R$ or $R ~~> R$ are not inner product spaces but are SemiHilbert spaces.

  Each SemiHilbert space has to define a subset of its function at ~test functions~. For $Z -> R$ these are functions with only finitely many non-zero elements. For $R ~~> R$ those are functions that are non-zero only on some finite interval. Next SemiHilbert space defines something like inner product, $<.,.> : X -> Test Functions X -> R$. But as you can see the second argument can accept only test functions. For $Z -> R$ this is $Σ_i u_i v_i$, not that this makes sense only because only finitely many terms are non-zero. For $R ~~> R$ it is $∫ f(x)g(x) dx$, again this makes sense as we can restrict the integration only to the finite interval on which $g$ is non-zero. The last main property is that if $<f,g> = 0$ for every test function $g$ then $f=0$. This is know as fundamental lemma of variational calculus. The abstract definition of Semi Hilbert space is exactly tailored to allow easy computations on variational calculus.

  Now the function $F : X -> Y$ between Semi Hilbert spaces has adjoint if it preserves test functions and is linear and continuous and there is a linear, continuous function $F† : Y -> X$ that preserves test functions and satisfy
  $$
  <F x, y>_Y = <x, F† y>_X  forall x y test functions
  $$

  To give and example, for $F u = u_(i+1) - u_i$
  $$
  Σ_i (u_(i+1) - u_i) v_i = (Σ_i u_(i+1) v_i) - (Σ_i u_i v_i) = (Σ_i u_i v_(i-1)) - (Σ_i u_i v_i) = Σ_i u_i (v_(i-1) - v_i)
  $$
  again all these operations makes sense only because u_i and v_i are nonzero for only finitely many indices i.

  Adjoint of derivative, i.e. $F f = f'$
  $$
  ∫ Φ'(x) Ψ(x) dx = - ∫ Φ(x) Ψ'(x) dx
  $$
  by per partes and the boundary term disappeared because we are working with test functions.

  Adjoint of convolution, i.e. $F(f)(x) = ∫ f(y) h(x-y) dy$
  $$
  ∫ ∫ Φ(y) h(x-y) dy Ψ(x) dx = ∫ ∫ Ψ(x) h(x-y) dx Φ(y) dy
  $$
  so $F†(f)(x) = ∫ f(y) h(y-x) dy$. Again this works out only because Ψ and Φ are test function and we can integrate over finite interval.

  Not every linear function has an adjoint. This is in a start contrast when working with Hilbert spaces, where every linear function has an adjoint. In particular, the constant combinator fails to have adjoint in a generic case.
  
  So the constant function over integers, $F x i := x$ of type $ F: R -> (Z -> R)$, does not have an adjoint. We can compute the candidate adjoint function $F†$P
  $$
  Σ_i (F x)_i u_i = x Σ_i u_i
  $$
  thus the candidate is $F† u = Σ_i u_i$ but that is not well defined for every function $Z -> R$. For example for $u_i = 1$, $sum_i u_i = infinity$. Therefore the adjoint does not exists.


** Semi-Hilbert spaces
  
  But the space $R ~~> R$ does not really have an inner product. The only candidate would be $<f,g> = ∫ f(x) g(x) dx$ but this does not have to be well defined as there is no integrability condition on $f$ and $g$. One traditionally work with L^2 space but again, this is a Banach space and we do not want to work with those. There is even a theorem that the space of all smooth function from R to R do not form a Banach space i.e. there is no way to define a norm on this spaces.

  Thus can we define adjunction without introducing inner product? Yes we can! With the help of distributions. Any function $R ~~> R$ can be understood as a distribution, that is a linear functional on test functions. And linear maps between distributions have a well defined adjoints! So we can take a linear function $F : (R ~~> R) -> (R ~~> R)$ that has the property that for test function ϕ the function $F ϕ$ is again a test function. Then this function defines a function on distributions $F' : D'(R) -> D'(R)$. This function has an adjoint in distributional sense: $F'† : D'(R) -> D'(R)$. Now the question is if there is a function $F† : (R ~~> R) -> (R ~~> R)$ that preserves test functions and when interpreted as a function between distributions $F†' : D'(R) -> D'(R)$, do we have equality $F†' = F'†$?

  If a function $F$ indeed preserves test functions and there is function $F†$ with above properties, we say that $F$ has adjoint and the adjoint is $F†$.

  In practices, this means you have to find a function $F†$ such that
  $$
  ∫ (F Φ)(x) Ψ(x) dx = ∫ Φ(x) (F† Ψ)(x) dx
  $$
  for every test function Φ Ψ.

* Forward mode

  Now we shift from symbolic to automatic differentiation i.e. we are a bit more concerned about the performance of the generated code. The problem with symbolic differentiation is that naive application of chain rule and product rule blows up resulting expression.
  The main problem is that
  $$
  ∂ (f ∘ g) ≠ ∂ f ∘ ∂ g
  $$
  A program is a long chain of function compositions and if we want to preserve the complexity of the program we should transform a composition of n function to composition of n functions.

  There are two ways how to define a forward mode differentiation operator with the above compositionality rule.

  The dual number approach, taking a function $f : X -> Y$ and extending it to its dual number version $𝓓 f : X × X -> Y × Y$. Then
  $$
  𝓓 (f ∘ g) = 𝓓 f ∘ 𝓓 g
  $$
  However, the dual number approach does not generalize to reverse mode, thus we do will not talk about it here. But we would like to point out that another useful code transformation is analytical continuation. I.e. extend a code defined over real numbers to complex numbers. In our applications, we have future plans on using some theorems from complex analysis. Because extension to complex number is very similar to extension to dual number we will deal with it in the future.

  The alternative is to take a function $f : X -> Y$ and produce the function that in addition to the value $f(x)$  also returns a program that computes the differential. $𝓣 f : X -> (Y × (X -> Y))$. In term of differential
  $$
  𝓣 f x = (f x, fun dx => ∂ f x dx)
  $$

  Then we have
  $$
  𝓣 (f ∘ g) = 𝓣 f ⋄ 𝓣 g
  $$
  The composition $⋄$ on the right is not the standard function composition $∘$, but it correctly composes the functions and their differentials. The definition is
  $$ (F : Y -> (Z × (Y -> Z))) (G : X -> (Y × (X -> Y)))
  F ⋄ G = fun x =>
    let (y, dg) := G x
    let (z, df) := F y
    (z, fun dx => df (dg x))
  $$

  
* Reverse mode


* Differentiation with side effects


* Future work

** Inversion

** Solver struct

** Working with quotients

** Testing
