Sharp bits & FAQ¶

The footguns that catch people, each with the fix, followed by a short FAQ. (Named after JAX's own Sharp Bits.)

Sharp bits¶

Enable 64-bit precision¶

Probabilistic ODE solvers propagate covariances and are sensitive to round-off. In float32 you may see covariances lose positive-definiteness or the log-likelihood go NaN. Turn on double precision before any array is created:

import jax
jax.config.update("jax_enable_x64", True)

`tspan` must be a tuple¶

tspan is a jax.jit static argument, so it must be hashable:

tspan = (0.0, 10.0)   # ✅ tuple
tspan = [0.0, 10.0]   # ❌ list -> unhashable, breaks under jit

Covariances are square-root factors¶

Anything named *_sqr is an upper-triangular factor, not a covariance. Reconstruct with P = P_sqr.T @ P_sqr; for standard deviations use np.sqrt(np.diag(P_sqr.T @ P_sqr)). See Notation.

Build models outside `jit`; pass `prior`/`measure`/`N`/`tspan` as static¶

The prior, measurement model, integer N, and tspan are static configuration — construct them in Python and mark them static (the loops already do this for you via static_argnums). Only arrays (mu_0, Sigma_0_sqr, parameters) are traced.

A solve is differentiable on both stepping modes¶

Both gaussian_filter (fixed grid) and gaussian_filter_adaptive (save-at-grid) are jit/grad/vmap-compatible, so any jax.grad over a solve works either way.

Dynamic calibration absorbs model–data misfit¶

Diffusion calibration sizes the uncertainty, but a dynamic calibration absorbs model–data misfit into sigma^2. When you want the misfit to stay visible (e.g. in a likelihood), use calibration="none" and set the prior scale matrix Xi deliberately. See Diffusion calibration.

FAQ¶

My covariance blows up / the log-likelihood is `NaN`¶

In order: (1) enable float64; (2) check q is not too high for the step size (very high order on a coarse grid is ill-conditioned); (3) make sure the initial covariance is sensible (taylor_mode_initialization pairs the mean with a zero/Dirac covariance); (4) for stiff problems, take smaller steps or use adaptive stepping.

The uncertainty band looks far too tight (or too wide)¶

That is a calibration symptom, not a bug. Plot the whitened residuals (see Diffusion calibration → What calibration cannot fix): a well-calibrated filter has ||z_n||^2 / d ≈ 1. Too tight ≫ 1, too wide ≪ 1.

How does this compare to `scipy` / Diffrax / probdiffeq / ProbNum?¶

scipy/Diffrax are classical solvers — fast point trajectories, no uncertainty. ode_filters, probdiffeq, and ProbNum are probabilistic solvers (Gaussian posterior + likelihood). ode_filters focuses on a transparent square-root EKF/RTS substrate with first-class custom measurement models. See What is a probabilistic ODE solver?.

Which function do I call?¶

See How to choose. Short version: gaussian_filter for a fixed-step solve and gaussian_filter_adaptive for tolerance-based stepping; both support jit/grad/vmap.

What do the fields of the filter result mean?¶

See the filter API reference and the annotated Quickstart step 4.