Bringing your own code: JAX functions with gradients inside a PyMC model

Sooner or later the model you want is not made of PyMC building blocks. The mean function is a simulator, an ODE solver, an optimiser, a neural network, or a colleague's library. This notebook shows how to put such a function inside a PyMC model without giving up NUTS, using JAX to supply the gradients. The application is the classic one: the Lotka-Volterra predator-prey equations fitted to the lynx and hare pelt records.

Several names in this corner of the API changed recently. The ones that matter here:

#1 · The question and the data

The Hudson's Bay Company bought pelts from trappers for centuries, and the number of lynx and snowshoe-hare pelts it took in each year is the textbook example of a predator-prey cycle. The question: are these cycles what the Lotka-Volterra equations predict, and what growth, predation and death rates do the data imply?

$$ \frac{du}{dt} = (\alpha - \beta v)\,u, \qquad \frac{dv}{dt} = (-\gamma + \delta u)\,v $$

with $u$ hares and $v$ lynx. Hares grow at rate $\alpha$ and are eaten at rate $\beta v$; lynx die at rate $\gamma$ and grow by $\delta u$ when there are hares to eat. There is no closed-form solution: to get the populations at the observation years you must numerically integrate the equations - and that is the external code PyMC cannot see into.

Two cycles of roughly ten years, with lynx peaking a year or two after the hares. Pelts are a noisy proxy for population, and the noise is plausibly multiplicative (a good trapping year scales everything), which is why the likelihood below is LogNormal.

#2 · Why you cannot just call your function

NUTS needs the gradient of the log-density with respect to every parameter. PyMC gets it by differentiating the symbolic PyTensor graph of the model. A PyMC random variable is a node in that graph, not a number - so code that needs actual numbers fails the moment you hand it one. We start with a function small enough to see through: an exponential decay evaluated on a fixed grid of five times.

Every attempt fails, each in its own words, for the same reason: k_rv is a symbolic TensorVariable. (Plain NumPy arithmetic on a random variable often appears to work, because TensorVariable overloads +, *, np.exp... and quietly builds a graph. That is PyTensor doing the work, not your function. Anything that loops, branches on values or calls compiled code - a solver, in short - cannot be rescued that way.)

The way in is to teach PyTensor a new primitive: an Op. An Op tells PyTensor three things: the types of its outputs (make_node), how to compute them from numbers (perform), and how to build the graph of its gradient (pullback).

#3 · Route 1 - write the Op by hand

You rarely need to do this any more (Route 2 is one line), but doing it once shows exactly what the one-liner automates, and it is still the route for non-JAX code: if you have a Fortran simulator with an adjoint, this is the template.

Reverse-mode differentiation never needs the full Jacobian $J$ of your function. It needs the vector-Jacobian product: given the gradient of the final scalar (the log-density) with respect to your outputs, $\bar{y}$, return the gradient with respect to your inputs, $\bar{x} = \bar{y}^\top J$. JAX hands you exactly that function: jax.vjp.

The Op now accepts symbolic variables. Compiling it produces one warning that you will see in every project that uses a Python perform, so here it is once, on purpose:

PyTensor's default backend is now Numba. Numba cannot compile an arbitrary Python perform, so it calls back into the interpreter for this one node ("object mode") and compiles everything around it. That is harmless: the node is a jitted JAX call anyway. It only means a little per-call overhead, which we will measure in section 8.

Now the gradient. pt.grad walks the graph backwards and, on reaching our node, calls pullback. Printing the gradient graph shows the second Op sitting where the derivative of DecayOp should be:

The gradient with respect to k is output .1 of a DecayVJPOp node with three inputs: the original a and k, and a Mul sub-graph (collapsed here) that computes the cotangent $\bar{y} = 2y$ - the derivative of sum(y**2) with respect to y, a length-5 vector like y itself. That node is what our pullback returned. The contract in full:

• inputs, outputs: the symbolic inputs and outputs of this node.
• cotangents: one symbolic variable per output, with that output's shape - the gradient of the final scalar with respect to it. (With several outputs, the ones the cost does not depend on arrive as DisconnectedType, and your code must cope.)
• return: one symbolic variable per input, with that input's shape.

Never trust a hand-written gradient until you have checked it. verify_grad compares it with finite differences at random projections; we also compare with JAX's own Jacobian.

#The classic pitfalls of Route 1

• float32. Without jax_enable_x64, JAX silently computes in single precision. Our perform casts the result to float64, which hides the problem instead of fixing it:

An error of 1e-8 in a value looks innocent. In the gradient of an ODE solution, accumulated over hundreds of steps, it is enough to make NUTS reject good proposals for no visible reason. Set the flag in the first cell, before any JAX array exists.

• Type declarations are promises. make_node said "float64 vector of length 5". If perform returns something else, the error appears far from its cause.
• jax.jit compiles on the first call for each new input shape/dtype. Jit once at module level (as above), not inside perform, or you pay compilation at every evaluation.
• No second derivatives: DecayVJPOp has no pullback of its own, so pt.grad(grad_k, k) raises NotImplementedError. NUTS does not need them; Laplace approximations and some optimisers do.
• The old names. If you implement grad or L_op as older tutorials say, PyTensor 3.3 still calls it, with a FutureWarning telling you to rename it to pullback.

That is some forty lines for a one-line function, and it has to be repeated for each new function. Hence:

#4 · Route 2 - pytensor.wrap_jax

wrap_jax takes a JAX-jittable function and returns a function of PyTensor variables.

Same numbers as the hand-written version, and the same construction underneath: a JAXOp whose perform calls the jitted function and whose pullback creates a second JAXOp (vjp_decay) that calls jax.vjp. What it did for us:

• Output types by abstract evaluation. It ran jax.eval_shape on dummy inputs to learn that two float64 scalars give a float64 vector of length 5. No declarations to get wrong.
• Gradients of any order, because the VJP Op is itself a JAXOp:

• Pytrees in and out. Dicts, tuples, even Equinox modules holding PyTensor variables are flattened for you; non-array arguments are treated as static.

#The limitation I actually hit: shapes must be known when the graph is built

jax.eval_shape needs concrete input shapes. A PyMC variable declared with dims= has a symbolic length (so that pm.set_data can resize it). wrap_jax then evaluates the shape once to get going, but throws the static output shape away; and if it cannot evaluate the shape at all, it refuses:

An output of shape (?,) still runs, but PyTensor can no longer check shapes while the graph is built, so mistakes surface at run time instead. The fix is one call to pt.specify_shape at the boundary, which is what the model below does.

#5 · Route 3 - stay in JAX end to end

So far PyTensor compiled the graph with Numba and called out to JAX for one node. But PyTensor can also compile the whole graph to JAX (mode="JAX"). This is what happens when you sample with NumPyro or with nutpie's JAX backend. Each Op is then translated by a function registered with jax_funcify. Our hand-written Op has no translation:

For Route 1 you must therefore register one yourself - trivial here, because the implementation is a JAX function. wrap_jax's JAXOp ships with its registration, so Route 2 needs nothing.

In JAX mode there is no callback and no copy between NumPy and JAX at the boundary: your function is inlined into one XLA program together with the priors and the likelihood. Whether that is faster depends on the model; we measure it on the real one below.

Summary of the three routes

#6 · The real function: a Lotka-Volterra solver in JAX

Two design decisions before any code.

Integrate the logarithms. With $x = \log u$, $y = \log v$ the equations become $\dot{x} = \alpha - \beta e^{y}$, $\dot{y} = -\gamma + \delta e^{x}$. Populations are then positive by construction, and the solver returns exactly what a LogNormal likelihood wants as its location. Section 10 shows what happens if you do not do this.

Two solvers. A hand-rolled fixed-step RK4 written with jax.lax.scan (twenty steps per year), and diffrax's adaptive Tsit5. lax.scan is the key to the first: a Python for loop would be unrolled into thousands of operations at trace time, and lax.while_loop cannot be reverse-differentiated at all. diffrax solves the same problem for adaptive step sizes with a bounded, checkpointed loop (RecursiveCheckpointAdjoint, the default) or by integrating the adjoint ODE backwards in time (BacksolveAdjoint).

Check the solvers against a reference (an 8th-order method at tolerance 1e-12), and time what NUTS will actually call: the value and gradient. The point is a plausible one: rates from the Stan case study on these data, starting from the first observation.

All three are far more accurate than the data warrant. The timings (they vary from run to run, and this notebook was built on a busy machine, so read ratios, not digits) show the pattern you should expect. The first call is dominated by compilation: hundreds to thousands of times the cost of a later call. On a problem this small and smooth the fixed-step scan is as cheap to differentiate as anything; diffrax's default checkpointed adjoint costs three to four times more, and its backsolve adjoint closes the gap. (The backsolve gradient is that of the continuous adjoint ODE, not exactly the derivative of the numbers the forward solve produced, which is why diffrax does not make it the default.) An adaptive solver earns its keep when the dynamics are stiff or the right step size varies wildly across the prior.

#7 · The model

$$ \begin{aligned} \log \text{pelts}_{t,s} &\sim \text{Normal}\big(\log z_s(t;\,\theta, z_0),\ \sigma_s\big) \qquad s \in \{\text{hare}, \text{lynx}\} \\ z(t) &= \text{solution of the Lotka-Volterra equations from } z_0 \text{ with rates } \theta \end{aligned} $$

Priors, on the scale of the data (thousands of pelts, years):

• $\alpha, \gamma \sim \text{LogNormal}(\log 1, 0.5)$: rates "per year" of order one. Hares can plausibly multiply by $e$ in a good year with no predators; lynx without prey starve within a year or so. The prior puts 90% of its mass between about 0.4 and 2.3.
• $\beta, \delta \sim \text{LogNormal}(\log 0.05, 1)$: the equilibrium of the system is $(u^\*, v^\*) = (\gamma/\delta,\ \alpha/\beta)$. With pelts in the tens of thousands, that equilibrium should sit somewhere between a few and a few hundred, so the interaction rates are of order $1/20$, give or take a factor of several.
• $z_0 \sim \text{LogNormal}(\log 10, 1)$: initial populations of order ten thousand pelts.
• $\sigma \sim \text{LogNormal}(-1, 1)$: multiplicative noise of roughly 10% to 100%.

The initial state is a parameter: the 1900 observation is as noisy as any other.

#Prior predictive: trajectories, not parameters

Priors on ODE parameters are impossible to judge one at a time - what matters is the trajectories they imply. Nothing needs gradients here, so this also confirms that the wrapped function works in plain forward sampling.

The prior covers the data with room to spare. It includes oscillations of every period and amplitude; a few percent of the trajectories peak above a million pelts, and about half dip below a thousand at some point, some of them far below the bottom of the axis (which is clipped at a single pelt). That is wider than anyone believes, but it rules out nothing plausible, and no draw broke the solver. Good enough to proceed.

#8 · Fit - and the first thing that goes wrong

Default settings first. pm.sample() uses nutpie with the Numba backend here, so the solver and its VJP run as object-mode callbacks (Route 2).

That is a failed fit: divergences, $\hat{R}$ up to 1.14, and a bulk ESS in the twenties for sigma. The per-chain table and the traces show what happened. Three chains agree. Chain 0 was still somewhere else when warm-up ended - gamma between 1 and 2, sigma two to three times larger, i.e. a poor fit explained away as noise. It dropped into the main mode after about 90 draws, but it never mixes properly afterwards: long flat stretches where it does not move at all, and all 101 divergences are its own. Its warm-up was spent adapting to a different region, so its step size and mass matrix are presumably wrong for this one. While preparing this notebook I ran this default fit with eight different seeds: five of them had at least one chain like this, some stuck in the wrong mode for the entire run.

The cause is not the JAX plumbing - it is the likelihood of an oscillator. If the model's period is wrong, the simulated cycles drift in and out of phase with the data, and the fit has many local optima: parameter values where at least some peaks line up. A one-dimensional slice shows it. Set $\alpha = \gamma = \omega$ (so the period of small oscillations is $2\pi/\omega$), fix the equilibrium at the observed mean populations, start from the first observation, and scan $\omega$:

The global minimum is at $\omega \approx 0.70$ (a period of about nine years). But there is a second basin at $\omega \approx 1.15$, where the model cycles faster than the data and only some of the peaks line up, and a string of shallow ones beyond. And PyMC's default starting value for a LogNormal(0, 0.5) rate is 1.13 - at the bottom of the wrong basin. This is only a one-dimensional slice of an eight-dimensional surface, but it is the right picture: a gradient-based sampler started there has no reason to leave quickly.

#The fix: start where the physics says, and tell the sampler about the correlations

Lotka-Volterra hands us a rough answer for free. Over a full cycle the time-averages of the populations equal the equilibrium values, $\bar{u} = \gamma/\delta$ and $\bar{v} = \alpha/\beta$, and small oscillations have period $2\pi/\sqrt{\alpha\gamma}$. The plot in section 1 shows a period of about ten years. So:

Two more settings, both specific to nutpie:

• nutpie adds uniform(-1, 1) jitter to the starting point on the unconstrained scale. For a log-transformed rate that is a factor of up to $e$ either way - enough to throw a chain into a neighbouring mode. compile_kwargs={"jitter_rvs": set()} switches it off. The price: all chains start from the same point, so $\hat{R}$ is a weaker check than usual. The rank plots and a second sampler (NumPyro, below) back it up.
• nuts={"adaptation": "low_rank"} adapts a low-rank-plus-diagonal mass matrix instead of a diagonal one. Section 9 shows why that matters here: the rates are strongly correlated.

We also time every fit from here on, including compilation, for the comparison below.

No divergences, $\hat{R} \le 1.002$, bulk ESS above 3500 for every parameter from 4000 draws, all four chains with similar step sizes, and the rank plots stay inside the envelope expected for chains that sample the same distribution. The posterior means ($\alpha \approx 0.55$, $\beta \approx 0.028$, $\gamma \approx 0.79$, $\delta \approx 0.024$, $\sigma \approx 0.25$) agree with the Stan case study on the same data.

#The same model, three ways of running it

Now the comparison promised in section 5. All runs use the same model object, the same starting point, no jitter, and the default 4 x (1000 + 1000) draws.

• nutpie, Numba backend (above): PyTensor differentiates the graph; the solver and its VJP are JAXOp callbacks.
• nutpie, Numba backend, diagonal mass matrix: as above with the default adaptation, to show what low_rank bought.
• nutpie, JAX backend: backend="jax" compiles the whole log-density to one XLA program; gradient_backend="jax" lets jax.grad differentiate it, solver included. (nuts_sampler_kwargs={"backend": "jax"} is the old spelling and now warns.)
• NumPyro: the sampler itself is JAX too. dense_mass=True is its counterpart of a non-diagonal mass matrix. With one CPU device it runs the four chains one after another (the warning below says so), whereas nutpie runs them in parallel threads.
• nutpie, JAX backend, diffrax: swap the solver for adaptive Tsit5.

Every run is clean and the posterior means agree to two or three significant figures, so the routes differ only in cost. Timings change from run to run (this page was built while other jobs were running), so read the ratios:

• The two nutpie backends are about equally fast on this model: the same sampling time, with the JAX backend roughly a second ahead in total because it compiles faster. The object-mode callback of Route 2 is not the bottleneck you might fear.
• NumPyro takes two to three times longer, almost entirely because it runs the four chains one after another on a single CPU device, while nutpie runs them in parallel. Its ESS is as good or better; per chain it is competitive.
• The mass matrix mattered more than the backend. With the default diagonal adaptation the same model takes about twice as long and yields less than a quarter of the effective draws: roughly a tenth of the ESS per second.
• diffrax costs about three times more than the hand-rolled RK4, for the same posterior. That is the price of adaptive step-size control on a problem that does not need it.

One more number, because sampling times mix the cost of a gradient with how many gradients the sampler chose to take. Below: the cost of one log-density-and-gradient evaluation of the compiled model in each mode - and, for contrast, the same model written with PyMC's built-in pm.ode.DifferentialEquation, which integrates the ODE and its forward sensitivities with SciPy's odeint.

Per evaluation, going through the Numba-mode callback costs a few hundredths of a millisecond more than the all-JAX program: visible for the cheap RK4 solver (roughly a third on top), lost in the noise for diffrax. pm.ode.DifferentialEquation works in PyMC 6.3, but each gradient costs over twenty times more than the wrapped RK4 (it calls back into Python for every right-hand-side evaluation, and integrates 2 x 6 extra sensitivity equations). I would not pick it for anything but a quick prototype.

#9 · What the fit says

Back to the science, using the first (nutpie/Numba, RK4) fit. First the posterior predictive check. For smooth curves we want the solution on a finer grid than the observation years - and because the solver is just a JAX function, we can call it directly on the posterior draws, vectorised with jax.vmap, without touching PyMC.

The two-parameter-per-species mechanistic model reproduces both cycles - timing, amplitude and the lag of lynx behind hares - and 41 of the 42 observations fall inside the 94% predictive band (about 39 expected). The inner band is uncertainty about the solution; the outer one adds the multiplicative observation noise, $\sigma \approx 0.25$. Note what the model cannot do: every cycle is identical, so the lower second lynx peak is attributed to noise.

In the phase plane the Lotka-Volterra solution is a closed orbit, traversed counter-clockwise: hares rise, then lynx, then hares crash, then lynx. The data circle the same region in the same direction, noisily. Derived quantities come with full uncertainty:

The posterior equilibrium sits on the observed mean populations, as the theory says it should, and the period is pinned down to a few months. That is also the explanation of the correlations, which reach -0.94: the data determine the period (roughly the product $\alpha\gamma$, hence the negative correlation between $\alpha$ and $\gamma$) and the equilibria (the ratios $\alpha/\beta$ and $\gamma/\delta$, hence the positive ones) much better than any single rate. A diagonal mass matrix cannot represent such a posterior; the low-rank adaptation can, which is where the four-fold gain in ESS came from.

#10 · What can go wrong

Everything in this section happened while this notebook was being written.

1. Multimodality from a bad start - section 8. The defaults are not safe for oscillatory models: start from a physically motivated point.

2. NaNs and infinities at extreme parameter values. During warm-up the sampler visits places the posterior never would. To see what the solvers do there, draw rates from a deliberately sloppy prior (half-normal with scales 2, 0.5, 2, 0.5) and solve on the natural scale and on the log scale:

On the natural scale the fixed-step RK4 overflows on 14 of the 2000 draws. The adaptive solver fails on 3 and returns negative populations on 377 more. A population that crashes towards zero is only accurate to atol, so it can step just below zero; then the sign of its growth term flips and it runs away (the most negative value here is -63). The logarithm in the likelihood turns all of those into NaN. On the log scale neither happens: a crash is just a large negative number.

That does not make the fixed-step solver right in those corners: there it differs from the adaptive solution by a factor of $e^{6}$ at the 99th percentile. A fixed step never fails; it is just silently wrong where the dynamics are fast. Here those regions carry no posterior mass - the RK4 and diffrax posteriors in section 8 agree - but that comparison is the check to make before trusting a fixed-step solver on a new problem.

3. max_steps and throw. An adaptive solver in a stiff corner can need more steps than its budget. diffrax then either raises, which kills the sampler from inside a callback, or - with throw=False - returns inf, which NUTS treats as a rejected proposal and moves on. For sampling you want the second behaviour:

4. Tracers, Python control flow and dynamic shapes. wrap_jax traces your function when the graph is built, so the usual JAX rules apply, and they bite at model-definition time. Branch on a value with jnp.where/lax.cond, not if; shapes must not depend on values; and lax.while_loop cannot be reverse-differentiated - use lax.scan, a bounded loop, or a library that implements an adjoint (diffrax, optimistix, lineax):

5. Silent float32 - section 3. If the first cell had not enabled x64, every number in this notebook would have been computed in single precision without a single error.

6. PyMC's own NUTS and fork. With nuts_sampler="pymc" the chains run in separate processes. On this machine (macOS, Apple silicon) PyMC starts them with fork, and a forked copy of an initialised, multithreaded JAX runtime deadlocks: the sampler sat at 0% CPU for nine minutes until I killed it. PyMC switches to a safe start method by itself when the backend is JAX, but it cannot know that a Numba-mode graph hides a JAX call. pm.sample(nuts_sampler="pymc", mp_ctx="spawn") or cores=1 both worked. nutpie and NumPyro run their chains in threads of one process, so they never meet the problem.

#11 · Summary

• NUTS needs gradients; an external function is a black box until you wrap it in an Op with a pullback. If the function is written in JAX, jax.vjp is the pullback.
• pytensor.wrap_jax writes both Ops and the JAX-backend registration for you. Give it inputs with static shapes (pt.specify_shape) and enable x64.
• The same wrapped model runs under the default Numba backend (JAX via callback), under nutpie's JAX backend and under NumPyro. Measure before assuming which one is fastest.
• For ODE models: integrate on a scale where the state cannot leave its domain, start the sampler from a physically motivated point, and expect correlated rate parameters.

#Try it yourself

1. An implicit function. The Lotka-Volterra orbit conserves $H(u, v) = \delta u - \gamma \log u + \beta v - \alpha \log v$. The peak hare population $u_{\max}$ is the larger root of $H(u, \alpha/\beta) = H(u_0, v_0)$. Write a JAX function that finds it by Newton iteration (a fixed number of steps in a lax.scan, or a lax.while_loop), give it a jax.custom_vjp from the implicit function theorem, $\partial u_{\max}/\partial\theta = -(\partial H/\partial u)^{-1}\,\partial H/\partial\theta$, wrap it with wrap_jax, and add u_max to the model as a pm.Deterministic. Check the gradient with verify_grad against differentiating through the unrolled iterations.
2. A better ecological model. Hares do not grow without bound in the absence of lynx. Add a carrying capacity, $\dot{u} = \alpha u (1 - u/K) - \beta u v$, with a prior on $K$ justified on the scale of the data. Compare the two models with az.compare (you need pm.compute_log_likelihood first; it works with a wrapped function). Expect a Pareto-k warning or two with only 42 observations - what does that tell you?
3. Swap the solver, measure the effect. Re-run the timing table with dfx.Dopri5(), with tol=1e-3 and tol=1e-9, and with BacksolveAdjoint. How far can you loosen the tolerance before the posterior moves or divergences appear? Then go the other way: how few RK4 steps per year can you get away with?

Next: E10 confronts another physical model with real detector data - LIGO's GW150914.