Using Julia autodiff code from python with JAX

[gdalle]

I think your best bet combines Reactant.jl with Enzyme.jl. I’m not sure exactly how to combine both to do what you want, but @wsmoses will know.
You may also be interested in past discussion on the topic:

Calling Python function JIT compiled with JAX from Julia without overhead

It could be very useful to call JIT compiled JAX functions from Python libraries from Julia, and in turn, to be able for Julia libraries to call user-provided JAX functions. This is possible today via PythonCall, e.g.: using PythonCall jax = pyimport("jax") numpy = pyimport("numpy") # Define a simple test function func_str = """ def f(x): return jax.numpy.sin(x) + jax.numpy.cos(x) """ # Create namespace and define function namespace = pydic…