-
-
Notifications
You must be signed in to change notification settings - Fork 609
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add a "gradient zoo" page? #2447
base: master
Are you sure you want to change the base?
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #2447 +/- ##
===========================================
+ Coverage 33.20% 71.52% +38.31%
===========================================
Files 31 31
Lines 1843 1921 +78
===========================================
+ Hits 612 1374 +762
+ Misses 1231 547 -684 ☔ 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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Out of curiousity, what models does it not work on atm?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
# TESTS BELOW ARE BROKEN FOR ZYGOTE BUT CORRECT FOR ENZYME! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Haven't used on Flux functions, but it is well tested in our CI
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