# Using Julia autodiff code from python with JAX

**URL:** <https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062>\
**Category:** New to Julia\
**Tags:** python, autodiff\
**Created:** [September 2, 2025, 6:19pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062 "2025-09-02T18:19:36Z")\
**Posts on this page:** 12\
**Page:** 1

<div class="post-metadata">

**Author:** ![gpxl](https://avatars.discourse-cdn.com/v4/letter/g/f14d63/32.png) [@gpxl](https://discourse.julialang.org/u/gpxl)\
**Post date:** [September 2, 2025, 6:19pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/1 "2025-09-02T18:19:37Z")

</div>

I have a pretty basic question. I am a scientist who is new to Julia (hence the category I am posting this in), but interested in learning more both out of personal curiosity as well as the possible benefits for my own work.

There is one step in a longer analysis pipeline in my field that would benefit from autodiff while not sacrificing speed. This step solves a system of differential equations to compute the theoretical predictions of a model, which is then compared against data in, e.g., MCMC. Other elements of this pipeline have autodiff implementations written by others in the community, but predominantly using python with JAX.

I am interested in plugging this hole in the pipeline by writing a Julia package that computes these theoretical predictions with autodiff. I am interested in Julia above using python with JAX for a few reasons.

1. Performance is important, as a single computation of the theoretical quantities mentioned above, in the modern C and/or Fortran codes widely used, can take up to O(minutes) for execution. Given the number of such computations needed for an MCMC (for example), it is important that the code is fast. My suspicion is that it would be much harder to write python code that is as performant, compared to Julia.
2. Ease of autodiff. It is probably possible to write code with autodiff that accomplishes this task in Python with JAX, but my impression is that it would be tedious to do while balancing clarity and performance. I think this is why no one in the community has really done this yet.
3. Clarity. The module I plan to write, in the best case scenario, would be easily extensible by others looking to use the code, say by adding terms to the system of equations that gets solved. I think, with well structured code, this should be easy even for new Julia users.

Given these motivations, I think Julia is a good option for me. But, since most of the pipeline is done with JAX, I am wondering if my idea is even possible to begin with. If I have a Julia module with autodiff implemented, can I benefit from that autodiff in a python pipeline that uses JAX, say by using the pyjulia package? Are there other considerations here that I am not aware of?

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [September 2, 2025, 6:40pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/2 "2025-09-02T18:40:12Z")

</div>

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](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552):
>
> 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…

---

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [September 2, 2025, 6:40pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/3 "2025-09-02T18:40:49Z")

</div>

Hello, just to make sure, do you want Jax to autodiff the computation made in Julia ? In which case I don’t think it’s possible without generating the MLIR in Julia ( with Reactant.jl for instance), however if you only want the forward pass to be in Julia I’m pretty sure it’s fine but you will want to reduce as much as possible the communication Julia Array → python array → Jax array which may be bad if done too much.  
I guess you should try it and if something is really bothering you ask it here.  
Also, Jax won’t be able to fuze as much as it does in python if you’re going through Julia it won’t even compile it well so you will loose the XPU advantage which you may have a hard time to rebuild

---

<div class="post-metadata">

**Author:** ![gpxl](https://avatars.discourse-cdn.com/v4/letter/g/f14d63/32.png) [@gpxl](https://discourse.julialang.org/u/gpxl)\
**Post date:** [September 2, 2025, 6:44pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/4 "2025-09-02T18:44:47Z")

</div>

Thanks for the pointer to Reactant.jl, definitely seems useful. I was already familiar with Enzyme.jl and currently plan to use it if I go forward with the project.

Thanks as well for the link to the other discussion. I had come across that thread in my searches prior to posting, and it’s definitely useful for me, but my problem is sort of the converse of that (I want to call Julia autodiff code from a predominantly JAX pipeline). Since I’m relatively new to the language and ecosystem it wasn’t clear to me that what I’m after was possible.

---

<div class="post-metadata">

**Author:** ![gpxl](https://avatars.discourse-cdn.com/v4/letter/g/f14d63/32.png) [@gpxl](https://discourse.julialang.org/u/gpxl)\
**Post date:** [September 2, 2025, 6:50pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/5 "2025-09-02T18:50:46Z")

</div>

I don’t think that’s exactly what I’m after, but I may be misunderstanding you. For simplicity you can picture the pipeline I’m talking about as having two steps.

\mathcal{T}(\theta) \to \mathcal{L}(\mathcal{T}(\theta))

where the first step is the theoretical prediction I want to implement in Julia, and the second step takes that prediction and passes it to a likelihood function written with JAX. My question is: is there a way to do this such that the whole pipeline is autodifferentiable (say, by calling the Julia code that gives you \mathcal{T} from python). I would also be interested in a setup where, e.g., I call the \mathcal{L} code from Julia to do the same.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [September 2, 2025, 6:53pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/6 "2025-09-02T18:53:57Z")

</div>

> [@gpxl](#):
>
> Thanks as well for the link to the other discussion. I had come across that thread in my searches prior to posting, and it’s definitely useful for me, but my problem is sort of the converse of that (I want to call Julia autodiff code from a predominantly JAX pipeline).

This link to the Lux.jl docs seems close to what you’re after?

> **[Exporting Lux Models to Jax (via EnzymeJAX & Reactant) | Lux.jl Docs](https://lux.csail.mit.edu/dev/manual/exporting_to_jax#Exporting-Lux-Models-to-Jax-(via-EnzymeJAX-and-Reactant))**
>
> Documentation for LuxDL Repositories

---

<div class="post-metadata">

**Author:** ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)\
**Post date:** [September 2, 2025, 6:59pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/7 "2025-09-02T18:59:40Z")

</div>

Can’t you just define a custom VJP rule in Jax for your Julia wrapper function, and then call some Julia AD package manually within the VJP?

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [September 2, 2025, 7:21pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/8 "2025-09-02T19:21:08Z")

</div>

Yeah Enzyme-JaX will let you do exactly that converse (call Julia from JaX, see the lux docs on it above). the readme for EnzymeJaX has an example of doing c++ from jax (and we don’t have a native julia as a string form, the setup from julia into jax requires the export – but if folks are interested we can add that feature too). Note that even with calling jax.grad on the outside, using EnzymeJaX to import the code will use Enzyme(MLIR) to autodiff it while compatible with JaX’s existing primitives, since jax doesn’t actually support autodiffing non-python code.

also cc @avikpal

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [September 2, 2025, 7:22pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/9 "2025-09-02T19:22:40Z")

</div>

Also in addition to autodiff, EnzymeJaX contains a bunch of linear algebra optimizations that will generically speed up existing jax code too 🙂

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [September 2, 2025, 7:29pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/10 "2025-09-02T19:29:18Z")

</div>

How can we use those 👀 ?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [September 2, 2025, 7:30pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/11 "2025-09-02T19:30:54Z")

</div>

Reactant already enables them

---

<div class="post-metadata">

**Author:** ![gpxl](https://avatars.discourse-cdn.com/v4/letter/g/f14d63/32.png) [@gpxl](https://discourse.julialang.org/u/gpxl)\
**Post date:** [September 2, 2025, 7:34pm UTC](https://discourse.julialang.org/t/using-julia-autodiff-code-from-python-with-jax/132062/12 "2025-09-02T19:34:01Z")

</div>

Thanks all, great to know this is generally possible and I really appreciate the links to concrete ways to get started. Enzyme-JAX seems like it might be the way to go for me.