| example | N | implementation | device | gradient |
|---|---|---|---|---|
| 01. Particles | 1,000,000 | enzyme | cpu | 4.78 s |
| 01. Particles | 1,000,000 | enzyme | gpu | 2.57 s |
| 01. Particles | 1,000,000 | jax | cpu | 5.98 s |
| 01. Particles | 100,000 | jax | gpu | 10.3 ms |
| 01. Particles | 1,000,000 | reactant | gpu | 196.1 ms |
| 02. Gauss–Seidel | 1,048,576 | enzyme / lexicographic | cpu | 1.32 s |
| 02. Gauss–Seidel | 1,048,576 | enzyme / red-black | cpu | 3.48 s |
| 02. Gauss–Seidel | 1,048,576 | enzyme / red-black | gpu | 21.9 ms |
| 02. Gauss–Seidel | 262,144 | jax / lexicographic | cpu | 19.70 s |
| 02. Gauss–Seidel | 4,096 | jax / lexicographic | gpu | 7.32 s |
| 02. Gauss–Seidel | 1,048,576 | jax / red-black | cpu | 1.81 s |
| 02. Gauss–Seidel | 1,048,576 | jax / red-black | gpu | 20.0 ms |
| 02. Gauss–Seidel | 4,096 | reactant / lexicographic | gpu | 8.34 s |
| 02. Gauss–Seidel | 1,048,576 | reactant / red-black | gpu | 16.8 ms |
| 03. Compaction | 10,000,000 | enzyme | cpu | 3.18 s |
| 03. Compaction | 10,000,000 | enzyme | gpu | 146.9 ms |
| 03. Compaction | 10,000,000 | jax / compacted | cpu | 4.82 s |
| 03. Compaction | 10,000,000 | jax / compacted | gpu | 52.6 ms |
| 03. Compaction | 10,000,000 | jax / masked | cpu | 11.99 s |
| 03. Compaction | 1,000,000 | jax / masked | gpu | 28.6 ms |
| 04. Newton | 1,000,000 | enzyme / implicit | cpu | 69.1 ms |
| 04. Newton | 1,000,000 | enzyme / unrolled | cpu | 130.3 ms |
| 04. Newton | 1,000,000 | jax / fixed 50 steps | cpu | 3.32 s |
| 04. Newton | 1,000,000 | jax / fixed 50 steps | gpu | 63.4 ms |
| 04. Newton | 1,000,000 | jax / implicit | cpu | 185.4 ms |
| 04. Newton | 1,000,000 | jax / implicit | gpu | 3.5 ms |
Why Enzyme and not JAX?
Differentiating programs whose efficient form is irregular
Note
Draft for review. All four examples are in; examples 03 and 04 do not have a Reactant column yet. Plan and status: AGENTS.md.
“Why do you use Julia instead of JAX?” is a fair question, and “JAX can’t sort” is the wrong answer: jax.lax.sort is differentiable. The real difference is the execution model.
JAX differentiates pure functions over fixed-shape arrays with a static computation graph. Anything that fits that model runs well, often better than hand-written code, because XLA fuses and vectorises it. Enzyme differentiates the imperative program as written, including in-place mutation, data-dependent loop bounds and growing buffers. When an algorithm’s efficient form is irregular, JAX makes you regularise it with padding, masks, static trip counts or implicit differentiation. Enzyme lets you keep it.
Each example below implements the same loss three times, with Enzyme.jl, with Reactant.jl, and with JAX, checks that all three give the same gradient, and times them on a CUDA GPU. Reactant is the control: it is Julia compiled through XLA, the same static-graph model as JAX. Where Reactant hits the same wall as JAX, the cause is the execution model, not the language.
Where the gap shows up
- Data-dependent output size or sparsity
-
Compaction and filtering (
x[mask],push!), sparse assembly with data-dependent fill-in, contact detection, tree codes (octrees, Barnes–Hut). JAX needs a static upper bound and masks. → 03. Compaction - Data-dependent control flow
-
Iterative solvers with convergence tests, adaptive time stepping. JAX cannot reverse-differentiate
lax.while_loop, so you fix the iteration count or switch to implicit differentiation. → 04. Newton - Sequential in-place updates
-
Gauss–Seidel, SOR, Thomas algorithm, sweeps where element i depends on the just-updated i−1. In JAX these become
lax.scanloops with one XLA step per element, and for nonlinear updates reverse mode stores every carry unless you rematerialise. → 02. Gauss–Seidel - Irregular neighbour access
- Particle methods with cell lists (sort by cell, loop over a data-dependent number of neighbours), scatter-accumulate with data-dependent indices, histogramming, resampling. → 01. Particles
Results at a glance
Gradient time at the largest problem size each example ran, Float64. Every number comes from the notebooks on the example pages, which also check that the gradients agree. Read each example’s “Where JAX wins” note before quoting a number.
Where JAX wins
Dense, batched, regular work: most of machine learning, many PDE stencils, anything whose naive cost already equals its smart cost. There the regularised JAX version is often faster despite padding, because XLA fuses and the hardware wants regular access. And the advantage on this site is Enzyme’s, not Julia’s: Julia AD tools that do not support mutation (Zygote) cannot do this either.