Support nested forward/reverse mode (#156)

Author: sritchie

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## [unreleased]
 
+- #156:
+
+  - Makes forward- and reverse-mode automatic differentiation compatible with
+    each other, allowing for proper mixed-mode AD
+
+  - Adds support for derivatives of literal functions in reverse-mode
+
 - #165:
 
   - Fixes Alexey's Amazing Bug for our tape implementation

src/emmy/abstract/function.cljc

@@ -19,6 +19,7 @@
             [emmy.numsymb :as sym]
             [emmy.polynomial]
             [emmy.structure :as s]
+            [emmy.tape :as tape]
             [emmy.util :as u]
             [emmy.value :as v])
   #?(:clj
@@ -250,7 +251,21 @@
     (->Function
      fexp (f/arity f) (domain-types f) (range-type f))))
 
-(defn- forward-mode-fold [f primal-s tag]
+(defn- forward-mode-fold
+  "Takes
+
+    - a literal function `f`
+    - a structure `primal-s` of the primal components of the args to `f` (with
+      respect to `tag`)
+    - the `tag` of the innermost active derivative call
+
+  And returns a folding function (designed for use
+  with [[emmy.structure/fold-chain]]) that
+
+  generates a new [[emmy.differential/Dual]] by applying the chain rule and
+  summing the partial derivatives for each perturbed argument in the input
+  structure."
+  [f primal-s tag]
   (fn
     ([] 0)
     ([tangent] (d/bundle-element (apply f primal-s) tangent tag))
@@ -262,15 +277,50 @@
            (g/+ tangent (g/* (literal-apply partial primal-s)
                              dx))))))))
 
+(defn- reverse-mode-fold
+  "Takes
+
+    - a literal function `f`
+    - a structure `primal-s` of the primal components of the args to `f` (with
+      respect to `tag`)
+    - the `tag` of the innermost active derivative call
+
+  And returns a folding function (designed for use
+  with [[emmy.structure/fold-chain]]) that assembles all partial derivatives of
+  `f` into a new [[emmy.tape/TapeCell]]."
+  [f primal-s tag]
+  (fn
+    ([] [])
+    ([partials]
+     (tape/make tag (apply f primal-s) partials))
+    ([partials [entry path _]]
+     (if (and (tape/tape? entry) (= tag (tape/tape-tag entry)))
+       (let [partial (literal-partial f path)]
+         (conj partials [entry (literal-apply partial primal-s)]))
+       partials))))
+
 (defn- literal-derivative
-  "Takes a literal function `f` and a sequence of arguments `xs`, and generates
-  an expanded `((D f) xs)` by applying the chain rule and summing the partial
-  derivatives for each [[emmy.differential/Dual]] argument in the input
-  structure."
+  "Takes
+
+    - a literal function `f`
+    - a structure `s` of arguments
+    - the `tag` of the innermost active derivative call
+    - an instance of a perturbation `dx` associated with `tag`
+
+  and generates the proper return value for `((D f) xs)`.
+
+  In forward-mode AD this is a new [[emmy.differential/Dual]] generated by
+  applying the chain rule and summing the partial derivatives for each perturbed
+  argument in the input structure.
+
+  In reverse-mode, this is a new [[emmy.tape/TapeCell]] containing a sequence of
+  pairs of each input paired with the partial derivative of `f` with respect to
+  that input."
   [f s tag dx]
-  (let [fold-fn  (cond (d/dual? dx) forward-mode-fold
-                       :else         (u/illegal "No tape or differential inputs."))
-        primal-s (s/mapr (fn [x] (d/primal x tag)) s)]
+  (let [fold-fn  (cond (tape/tape? dx) reverse-mode-fold
+                       (d/dual? dx)    forward-mode-fold
+                       :else           (u/illegal "No tape or differential inputs."))
+        primal-s (s/mapr (fn [x] (tape/primal-of x tag)) s)]
     (s/fold-chain (fold-fn f primal-s tag) s)))
 
 (defn- check-argument-type
@@ -305,7 +355,7 @@
     (if-let [[tag dx] (s/fold-chain
                        (fn
                          ([] [])
-                         ([acc]     (apply d/tag+perturbation acc))
+                         ([acc]     (apply tape/tag+perturbation acc))
                          ([acc [d]] (conj acc d)))
                        s)]
       (literal-derivative f s tag dx)

src/emmy/calculus/derivative.cljc

@@ -15,6 +15,7 @@
             [emmy.operator :as o]
             [emmy.series :as series]
             [emmy.structure :as s]
+            [emmy.tape :as tape]
             [emmy.util :as u]
             [emmy.value :as v])
   #?(:clj
@@ -531,12 +532,23 @@
      (letfn [(process-term [term]
                (g/simplify
                 (s/mapr (fn rec [x]
-                          (if (d/dual? x)
-                            (d/bundle-element
-                             (rec (d/primal x))
-                             (rec (d/tangent x))
-                             (d/tag x))
-                            (-> (g/simplify x)
-                                (x/substitute replace-m))))
+                          (cond (d/dual? x)
+                                (d/bundle-element
+                                 (rec (d/primal x))
+                                 (rec (d/tangent x))
+                                 (d/tag x))
+
+                                (tape/tape? x)
+                                (tape/->TapeCell
+                                 (tape/tape-tag x)
+                                 (tape/tape-id x)
+                                 (rec (tape/tape-primal x))
+                                 (mapv (fn [[node partial]]
+                                         [(rec node)
+                                          (rec partial)])
+                                       (tape/tape-partials x)))
+
+                                :else (-> (g/simplify x)
+                                          (x/substitute replace-m))))
                         term)))]
        (series/fmap process-term series)))))