Add a "gradient zoo" page? #2447
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #2447 /- ##
===========================================
Coverage 32.69% 65.55% 32.86%
===========================================
Files 34 35 1
Lines 1976 2015 39
===========================================
Hits 646 1321 675
Misses 1330 694 -636 ☔ View full report in Codecov by Sentry. |
|
|
||
| * Long compilation times, on the first call. | ||
|
|
||
| * Does not at present work on all Flux models, due to missing rules. |
There was a problem hiding this comment.
Out of curiousity, what models does it not work on atm?
There was a problem hiding this comment.
FWIW all the CPU Flux tests I know of presently pass (and please let me know any failing ones so I can fix them), including some which Zygote fails on:
Flux.jl/test/ext_enzyme/enzyme.jl
Line 143 in d39d13b
There was a problem hiding this comment.
That's great!
Sorry this was written a few months back, and I don't recall what issues I ran into. I see many have been fixed.
|
|
||
| New package which works on the LLVM code which Julia compiles down to. | ||
|
|
||
| * Allows mutation of arrays. |
There was a problem hiding this comment.
I might also add high-performance :) [by virtue of pioneering the "run after optimization" trick (https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf)]
|
|
||
| * Does not always handle type instability. | ||
|
|
||
| * Custom rules by its own rules... Generally fewer such rules than Zygote, and at a lower level -- applied to `BLAS.gemm!` not `*`. |
There was a problem hiding this comment.
Enzyme also supports importing chainrules as well (though it is discouraged as they are usually lower-performance for sake of not understanding mutation)
| Zygote.hessian is like this. | ||
|
|
||
| ### Enzyme.jl | ||
|
|
There was a problem hiding this comment.
Haven't used on Flux functions, but it is well tested in our CI
mcabbott
force-pushed
the
grad_zoo
branch
2 times, most recently
from
86c7fd0 to
6707c6f
Compare
A while back I tried to write up a docs page pointing out that Flux is AD-agnostic -- anything which produces an appropriately nested structure should work.
This is a rough draft, help with tidying it up would be welcome.
Perhaps it's too opinionated about other packages. Certainly any opinions it offers should have clear dates attached.
Rendered version: https://github.com/mcabbott/Flux.jl/blob/grad_zoo/docs/src/tutorials/gradient_zoo.md