Enzyme
Enzyme copied to clipboard
EnzymeAD
Feasibility of combining Numba and Enzyme
Because they're both based on LLVM, I was wondering what would be involved in integrating Numba with Enzyme, so that I could decorate a python function and get (optimized) gradients. I poked around the documentation a little but didn't see anything relevant.
You could take a look of how the Julia frontend does things, https://github.com/wsmoses/Enzyme.jl GPUCompiler works similar to Numba, and there is a C-API call for creating the enzyme pass.
Happy to chat about the detailed strategy, but I am not familiar enough with Numba to make any direct recommendations.
It definitely should be possible and I agree the Julia front end is likely a good place to start (also happy to video call next week and do some planning).
My bet as to the big things that will need to be done for Numba as a front end:
- Load Enzyme as a plugin and call it as an optimization at the appropriate time (Julia has good example of)
- Expose API for gradients in numba
- Ensure Numba passes down all of the relevant bitcode
- Ensure Numba passes down Type Information that can be understood (as TBAA)
- Re-entry into Numba JIT if AD of non fully statically analzable code is desired
Hi, would just like to add a +1 that this would be of great interest to me. Have any steps been made or anyone else taken a look in the last couple years?
