Numba-Enzyme: A Differentiable JIT-ed Python
2023-02-22, 15:15–15:30 (US/Mountain), ECNT 312

In this short technical talk we will present a first prototype, as well as a forward-looking roadmap for Numba-Enzyme, a gradient-providing Just-in-time (JIT) compiler for Python relying on Numba to JIT Python, and Enzyme to provide the gradients for respective kernels. Including recent advances in Enzyme for forward-mode, we will prevent the JIT-pipeline for forward-mode differentiation, and reverse-mode differentiation as well as presenting first performance comparisons to C++ differentiated code with Enzyme, as well as gradient computation in JAX. We will conclude by providing an outlook on future extensions to expose Enzyme's vectorization, as well as kernel composability between Numba-Enzyme, and JAX to enable users to leverage the strengths of both ecosystems.