Support nested forward/reverse mode #156

sritchie · 2023-12-08T03:09:17Z

This PR:

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
Adds support for emmy.operator/Operator-shaped outputs in reverse-mode

In a future PR I want to add to / modify the IPerturbed protocol so that it's as supportive of reverse-mode as forward-mode. I'm leaving that to the future for now.

codecov · 2023-12-08T03:19:43Z

Codecov Report

Attention: Patch coverage is 94.91525% with 6 lines in your changes are missing coverage. Please review.

Project coverage is 87.67%. Comparing base (d354463) to head (2617c36).

❗ Current head 2617c36 differs from pull request most recent head 1e86dc5. Consider uploading reports for the commit 1e86dc5 to get more accurate results

Files	Patch %	Lines
src/emmy/abstract/function.cljc	71.42%	2 Missing and 2 partials ⚠️
src/emmy/tape.cljc	97.72%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #156      +/-   ##
==========================================
- Coverage   87.68%   87.67%   -0.02%     
==========================================
  Files          99       99              
  Lines       15862    15791      -71     
  Branches      853      848       -5     
==========================================
- Hits        13908    13844      -64     
+ Misses       1101     1099       -2     
+ Partials      853      848       -5

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

src/emmy/tape.cljc

sritchie

Some comments to help you on the journey, @littleredcomputer ...

src/emmy/abstract/function.cljc

sritchie · 2024-04-22T18:25:39Z

src/emmy/abstract/function.cljc

@@ -262,15 +277,50 @@
           (g/+ tangent (g/* (literal-apply partial primal-s)
                             dx))))))))

+(defn- reverse-mode-fold


building on the cryptic fold-chain from a previous PR, now we can drop this in and support reverse-mode.

sritchie · 2024-04-22T18:26:04Z

src/emmy/calculus/derivative.cljc

+                                 (rec (d/tangent x))
+                                 (d/tag x))
+
+                                (tape/tape? x)


this feels janky, and is the only place where we deliberately introspect these types.

sritchie · 2024-04-22T18:26:14Z

src/emmy/differential.cljc

@@ -557,36 +556,6 @@
  (boolean
   (some #{tag} *active-tags*)))

-(defn inner-tag


moved to emmy.tape

sritchie · 2024-04-22T18:26:26Z

src/emmy/differential.cljc

@@ -699,280 +668,19 @@

 ;; ## Chain Rule and Lifted Functions
 ;;
-;; Finally, we come to the heart of it! [[lift-1]] and [[lift-2]] "lift", or


this all moves to emmy.tape

sritchie · 2024-04-22T18:26:53Z

src/emmy/tape.cljc

+            d/*active-tags*)
+      (apply max tags)))
+
+(defn tag+perturbation


same as we had in emmy.differential, just using tag-of internally vs emmy.differential/tag

sritchie · 2024-04-22T18:27:04Z

src/emmy/tape.cljc

@@ -548,6 +588,13 @@
        (f/function? output)
        (->partials-fn output tag)

+        (o/operator? output)


operator support, this should be handled with a protocol

sritchie · 2024-04-22T18:29:26Z

src/emmy/tape.cljc

-(defunary g/zero? (by-primal g/zero?))
-(defunary g/one? (by-primal g/one?))
-(defunary g/identity? (by-primal g/identity?))
+(defunary g/zero?


we are careful to only respond true here if the tangent is also zero. should we be just as careful with tapes? by construction we really shouldn't be building duals with empty tangents etc

Do we need g/zero? for tapes? If not, I'd leave it completely undefined: those things are kind of dangerous as you have discovered.

littleredcomputer · 2024-04-23T20:59:33Z

src/emmy/abstract/function.cljc

+     (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))))


it is odd that nothing triggers L300, so in every case, (and (tape/tape? entry) (= tag (tape/tape-tag entry))).

this would happen in the case of nested gradient calls to a literal function, I just don't have those in the tests yet. In that case, the innermost tag wins and any tag with a different tape is treated as a scalar (i.e. partial derivative == 0, so we never add it to the map)

I want to share tests between forward mode and reverse mode... when I do that this should get hit.

littleredcomputer · 2024-04-23T21:04:11Z

src/emmy/calculus/derivative.cljc

+                                 (tape/tape-tag x)
+                                 (tape/tape-id x)
+                                 (rec (tape/tape-primal x))
+                                 (mapv (fn [[node partial]]


looks like the inner thing is mapv too. Does writing it that way feel less janky? [node partial] is almost an instance of Dual, right, except that the tag belongs to the containing tape. Armed with that, you could unify this part with the part above, but (I say sincerely) it would be too clever. Therefore I don't think this is janky: it feels like about what we want to see here

this whole function is too clever, and my cleverness bit me here: #168

cleverness in trying to make this compatible with D, that is...

littleredcomputer · 2024-04-23T21:10:04Z

src/emmy/tape.cljc

+   (let [m (into {} (mapcat
+                     (fn [dx]
+                       (when-let [t (tag-of dx)]
+                         {t dx})))


I wonder if it is cheaper to write [t dx]. However you choose to write it, it's cool (and I do think that {t dx} is a reasonable way to write the lone binding

your way is faster!!

littleredcomputer · 2024-04-23T21:12:29Z

src/emmy/tape.cljc

-  NOTE: `df:dx` has to ALREADY be able to handle [[TapeCell]] instances. The
-  best way to accomplish this is by building `df:dx` out of already-lifted
-  functions, and declaring them by forward reference if you need to."
+  NOTE: `df:dx` has to ALREADY be able to handle [[TapeCell]]


this makes sense; it's so sweet. keep the magic at the most primitive layer

littleredcomputer · 2024-04-23T21:15:51Z

src/emmy/tape.cljc

-(defunary g/zero? (by-primal g/zero?))
-(defunary g/one? (by-primal g/one?))
-(defunary g/identity? (by-primal g/identity?))
+(defunary g/zero?


Do we need g/zero? for tapes? If not, I'd leave it completely undefined: those things are kind of dangerous as you have discovered.

littleredcomputer · 2024-04-23T21:23:49Z

test/emmy/tape_test.cljc

@@ -11,6 +11,7 @@
            [emmy.generators :as sg]
            [emmy.generic :as g]
            [emmy.numerical.derivative :refer [D-numeric]]
+            [emmy.operator :as o]


I just had this random test idea, which I don't think is particularly valuable or anything, but thought it was cool: "the Jacobian of a permutation is the permutation matrix of that permutation", as in :

(defn f [a b c d e] (up d e c b a)) (def M ((D f) 'a 'b 'c 'd 'e)) ;; M = (down (up 0 0 0 0 1) (up 0 0 0 1 0) (up 0 0 1 0 0) (up 1 0 0 0 0) (up 0 1 0 0 0)) ;; and so: (* M [1 2 3 4 5]) ;; => (up 4 5 3 2 1)

Added as a test!

sritchie requested a review from littleredcomputer December 8, 2023 03:09

littleredcomputer approved these changes Dec 8, 2023

View reviewed changes

src/emmy/tape.cljc Outdated Show resolved Hide resolved

Base automatically changed from sritchie/tape to main December 8, 2023 17:10

sritchie force-pushed the sritchie/nested branch 2 times, most recently from 34dea5a to 405160e Compare December 8, 2023 17:19

sritchie force-pushed the sritchie/nested branch 3 times, most recently from 1ba84d7 to e2d3f0e Compare April 9, 2024 00:27

sritchie changed the base branch from main to sritchie/easy April 9, 2024 00:27

Base automatically changed from sritchie/easy to main April 13, 2024 01:15

merge forward and reverse mode

f343c88

sritchie force-pushed the sritchie/nested branch 2 times, most recently from 99a0a08 to c74897e Compare April 18, 2024 18:37

tape works

69ee50a

sritchie force-pushed the sritchie/nested branch from c74897e to 69ee50a Compare April 21, 2024 02:01

add tests

3ece359

sritchie force-pushed the sritchie/nested branch from 227276b to 3ece359 Compare April 22, 2024 18:23

sritchie commented Apr 22, 2024

View reviewed changes

sritchie added 2 commits April 22, 2024 12:49

remove deep-primal

deedf55

new test

2617c36

littleredcomputer approved these changes Apr 23, 2024

View reviewed changes

add another test

1e86dc5

sritchie merged commit bd47efd into main Apr 23, 2024
4 checks passed

sritchie deleted the sritchie/nested branch April 23, 2024 22:22

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support nested forward/reverse mode #156

Support nested forward/reverse mode #156

sritchie commented Dec 8, 2023 •

edited

Loading

codecov bot commented Dec 8, 2023 •

edited

Loading

sritchie left a comment

sritchie Apr 22, 2024

sritchie Apr 22, 2024

sritchie Apr 22, 2024

sritchie Apr 22, 2024

sritchie Apr 22, 2024

sritchie Apr 22, 2024

sritchie Apr 22, 2024

littleredcomputer Apr 23, 2024

littleredcomputer Apr 23, 2024

sritchie Apr 23, 2024

sritchie Apr 23, 2024

littleredcomputer Apr 23, 2024

sritchie Apr 23, 2024

sritchie Apr 23, 2024

littleredcomputer Apr 23, 2024

sritchie Apr 23, 2024

littleredcomputer Apr 23, 2024

littleredcomputer Apr 23, 2024

littleredcomputer Apr 23, 2024

sritchie Apr 23, 2024

Support nested forward/reverse mode #156

Support nested forward/reverse mode #156

Conversation

sritchie commented Dec 8, 2023 • edited Loading

codecov bot commented Dec 8, 2023 • edited Loading

Codecov Report

sritchie left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sritchie commented Dec 8, 2023 •

edited

Loading

codecov bot commented Dec 8, 2023 •

edited

Loading