Autodiff support in Numba

[adam-urbanczyk]

Feature request

It would be great if Numba supported automatic differentiation. Maybe using Enzyme would be the easiest way as it operates directly on the IR of LLVM. Another possible source of inspiration could be Taichi.

[gmarkall]

I'd been meaning to look into Enzyme, as it also seems to support CUDA-specific techniques. @adam-urbanczyk have you any thoughts or experiments in the direction of AD support for Numba that you could add here?

[adam-urbanczyk]

Nothing too specific. My use case is estimating probability distributions and I'd be already quite happy with jacobians and hessians of the functions from math and being able to differentiate through if statements.

[gmarkall]

I think this is an area none of the Numba maintainers have any significant amount of experience with, so it'll be difficult for us to push forward on this without a significant amount of hand-holding, but maybe we can establish some scope here with a bit of further discussion. To kick this off, I have a few questions:

• What is a small example of a bit of code that you'd like to be able to write?
• What would Numba do / generate under the hood - for example what would some expected transformation of this code look like?
• How was that transformation derived?

[adam-urbanczyk]

* What is a small example of a bit of code that you'd like to be able to write?

Maybe something like this:

@jit
def f(x,a,b):

   return exp(-a*x**2)*erf(b*x)  

df = f.gradient(a=True, b=True)

a = 1
b = 2
x = np.linspace(-1,1)
f_val = f(x,a,b) # return the function value
f_jacobian = df(x,a,b) # return the gradient w.r.t. a and b. Note that f_jacobian is 2D.

* What would Numba do / generate under the hood - for example what would some expected transformation of this code look like?

I'm just a user of AD so cannot say too much here. My understanding is that Enzyme allows to directly differentiate LLVM IR and Numba generates IR at certain point. It should be "easy" to combine those two tools.

Here is an official example on the IR level from their compiler explorer instance:

declare dso_local double @__enzyme_autodiff(i8*, double)

define double @square(double %x) {
entry:
  %mul = fmul double %x, %x
  ret double %mul
}

define double @dsquare(double %x) local_unnamed_addr {
entry:
  %call = tail call double @__enzyme_autodiff(i8* bitcast (double (double)* @square to i8*), double %x)
  ret double %call
}

The way I understand it Numba already generates something like @square and what would be needed is to add the @dsquare part and run the Enzyme transformation on this IR. More examples on IR level are here: https://github.com/EnzymeAD/Enzyme/tree/main/enzyme/test/Enzyme/. Probably ForwardMode and ForwardModeVector are most usefel for the usacase I proposed at the beginning.

* How was that transformation derived?

That is in my understanding already handled by Enzyme via the __enyzme_autodiff and related functions.

[wsmoses]

In contrast Enzyme dev here, with little Numba internal knowledge. Regardless happy to help and this has been something that definitely would be really useful to a lot of people.

In particular what you see by forward, etc is that Enzyme supports several different AD algorithms/modes. Forward mode gets the derivative of all results wrt a single input, reverse mode gets the derivative of all inputs wrt an output (gradient), and vector modes thereof change single in either of the above to one’s choice of N.

I think the first question here is what communication mechanism makes the most sense. The easiest thing to get working is that Enzyme has a LLVM module pass that replaces all AD calls with a call to a newly generated derivative function. Alternatively, there’s also an API that takes an LLVM function and returns a newly generated derivative function. The API level has more flexibility, but requires a bit more setup.

Regardless either way I imagine the overall pipeline to be Numba generated LLVM, and via a choice of above, during compilation Numba calls Enzyme to generate the derivative functions.

[ludgerpaehler]

Fellow Enzyme-dev here who already looked at Numba and is already working on a prototype - would be really happy to link up to get a more complete coverage and move to have what you are envisioning @adam-urbanczyk

What I think we can consciously move towards is support for

• Forward-Mode
• Vector Forward-Mode
• Reverse-Mode
• Vector Reverse-Mode

What kind of target platforms are you thinking of? From an Enzyme perspective Threads, & OpenMP are the easiest. CUDA might be possible? We do not support Intel Threading Building Blocks (TBB) at the current moment, but it would be a matter of software engineering understanding for TBB into Enzyme.

@gmarkall do I understand it correctly that Numba does not use CUDA-Streams, or CUDA-Graphs for its CUDA support?

[ludgerpaehler]

cc @lukasheinrich

[lukasheinrich]

Hi,

This is a great initiative. as @ludgerpaehler says, there has already been some effort in this direction.

in terms of API, I think it makes sense to stay as close as possible to the JAX

i.e.

def f(x):
  ...
  
numba.jvp(f...)
numba.vjp(f,...)
numba.grad(f,..)
numba.value_and_grad(f,..)

etc

[lukasheinrich]

cc @niklasschmitz

[gmarkall]

@gmarkall do I understand it correctly that Numba does not use CUDA-Streams, or CUDA-Graphs for its CUDA support?

Apologies for the delay in my reply here. It doesn't use graphs, but it supports streams: https://numba.readthedocs.io/en/latest/cuda-reference/host.html#stream-management

[adam-urbanczyk]

@gmarkall coming back to the original question, is the description good enough? Do you need more details?

@ludgerpaehler can you share the code you are working on? Do you know if enzyme supports higher order derivatives (specifically calculation of hessians)?

[wsmoses]

You can indeed use Enzyme to take arbitrary higher order derivatives.

[gmarkall]

@gmarkall coming back to the original question, is the description good enough? Do you need more details?

I guess so - as my experience in this area is extremely limited, I feel that the most I can do is to help facilitate the discussion and experimentation by those more qualified than myself, and that seems to be happening so I'd imagine it is good enough...

Perhaps others who have commented in this thread can suggest if any more detail or clarification is required, if necessary.

[adam-urbanczyk]

I see that there was a relevant talk in the Enzyme Conference 23

Numba-Enzyme: A Differentiable JIT-ed Python

Anyone here is/knows the authors and can share slides/video?