-
Notifications
You must be signed in to change notification settings - Fork 54
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
Switch from Zygote to Enzyme #182
Conversation
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #182 +/- ##
==========================================
+ Coverage 72.69% 74.05% +1.36%
==========================================
Files 37 36 -1
Lines 5560 4942 -618
==========================================
- Hits 4042 3660 -382
+ Misses 1518 1282 -236 ☔ View full report in Codecov by Sentry. |
@jgreener64 if any of these issues are still outstanding, I can try to quickly fix if you have MWE's |
Thanks for asking. It's mostly the GPU reduction stuff (JuliaGPU/CUDA.jl#2471) and any further issues that uncovers. Also EnzymeAD/Enzyme.jl#1672 would be useful but is not blocking. |
Merging this to prevent development stalling elsewhere. There are a few things that need fixing before release related to Enzyme, GPU support and docs updates. |
Switching from using Zygote+Enzyme to just using Enzyme for AD allows us to rewrite much of the main computation code to mutate memory. This improves performance on CPU and GPU.
Opening this PR to test with CI.