Writing AD-safe scientific numerics

A function can return the right number and still be wrong. When a gradient is part of its contract, a silently zero, NaN, or detached derivative can break the science even though the value “worked.” When the operation is deliberately discrete, pretending it has a useful gradient is equally wrong. It ran is not it is correct. Elegant nonsense is still nonsense.

jaxstro therefore requires every public numerical path to name its transform contract. Smooth pathwise gradients receive independent finite-difference checks. Expected-zero, blocked, surrogate, validation-only, and discrete paths state their narrower claim instead of borrowing the language of smooth inference. This page is the thesis: ten principles for making those scientific contracts explicit, testable, and teachable.

If you are starting from a research question rather than a method name, use Science patterns enabled by Jaxstro to route events, equilibria, integrals, tabulated models, limiting distributions, spatial interactions, and provenance to the relevant module and evidence.

Gradient contracts¶

The first question is not “does jax.grad run?” It is “what derivative claim is scientifically valid here?” jaxstro.testing exposes five live contracts, and the audit gate interprets each one differently.

1. Classify the transform contract first¶

Before implementation, classify the transform contract first. A smooth kernel, an expected-zero sensitivity, a stopped gradient, a surrogate, and a discrete index builder require different evidence. For smooth_pathwise and known_zero, the audit computes both AD and an independent finite-difference estimate. Other contracts fail closed for inference and carry only the narrower claim they name. See Validation for the measured audit.

2. Fixed iteration is necessary, not sufficient¶

Fixed scan lengths give JAX a static computation and avoid reverse-mode limits around data-dependent convergence loops. They do not make the update map smooth. A fixed-step solver can still contain branch-selected intervals, clips, or singular derivatives.

For a smooth function with a nonzero derivative and a parameter-independent initial guess, Newton can carry a smooth_pathwise contract after AD–FD verification. In contrast, bisection is a branch-selected forward solve: it can deliver an accurate root value without providing the smooth inverse sensitivity needed for inference. Iteration count and gradient contract are separate facts.

→ Root-finding — the distinct contracts of bisect, newton, and newton_ppf.

3. Guard singularities without killing the gradient — the where-trap¶

The natural way to avoid a division by zero is jnp.where(d == 0, fallback, a / d). The selected forward value may be finite, but both branch expressions are traced. If the unselected expression creates an inf or NaN, its reverse-mode cotangent can meet a zero multiplier and the inactive branch can still poison a derivative. The discipline is to guard the operand, not only the selected result: sanitize the denominator before division, then select the intended value. safe_div and safe_log implement that policy for their documented domains.

4. Saturation is a silent gradient killer¶

clip, min, max, and floor are piecewise or discrete operations. They can zero, route, or make a gradient convention-dependent at their boundaries. Sometimes that is the intended hard-bound contract. Sometimes it pins a parameter on a wall while an optimizer reports convergence. Name which case you intend. When newton_ppf clips iterates to [lo, hi], for example, its smooth pathwise claim applies to an interior solution, not to saturation at the support.

→ newton_ppf — the inverse-CDF solver, and the clip-to-support caveat discusses the clip-to-support trade-off.

5. Floating point is part of the math¶

Catastrophic cancellation, overflow in exp, and underflow in log are not edge cases — they are the common case in likelihood code. Work in the log domain, use log1p/expm1 near zero, and sum in the order that minimizes error. jaxstro provides stable_log1p, stable_expm1, safe_log, safe_exp, and Neumaier compensated summation for reductions where the ordinary sum loses digits. And turn on float64 first (8. Precision discipline).

6. Non-differentiable operations are forbidden in the differentiable graph¶

argmax, argsort, sort, integer casts, and data-dependent shapes have no useful gradient. They are not banned from the package — the spatial module needs them — but they must be isolated from any path you will differentiate. Build the Morton codes and neighbor lists once, as discrete preprocessing; keep the differentiable physics downstream of them.

→ Spatial indexing and neighbor contracts — fixed-capacity cells, candidate recall, exact-pair overflow, and the boundary between discrete identity and downstream differentiable values.

7. Quadrature and sampling differentiate through the values, not the nodes¶

A Gaussian quadrature rule has fixed nodes and weights; an inverse-CDF sampler has a fixed grid. Differentiate through the integrand evaluated at the nodes or the values being interpolated, never through the node positions. This is why the quadrature factory generates nodes once on the host with numpy and freezes them to constants: the gradient flows through f(x_i), and the constant $x_i$ contributes nothing it should not.

→ Newton–Cotes integration — Newton–Cotes integration over a grid of values.

→ Fixed-node quadrature — fixed-node Gaussian, Clenshaw-Curtis, and cumulative Simpson rules differentiate through values rather than node generation.

→ Cubic interpolation — PCHIP-style interpolation differentiates inside stable limiter branches and avoids inventing monotone-table overshoot.

→ Regular-grid interpolation — multilinear interpolation differentiates inside grid cells while making out-of-domain policy explicit.

→ B-splines — B-spline evaluation differentiates cleanly through coefficients and interior coordinates for fixed knots.

8. Precision discipline¶

Float32 carries about 7 decimal digits; one bad subtraction can spend all of them. Enable float64 with jaxconfig.enable_high_precision() before creating any array, and request the highest matmul precision so reductions are not silently downcast on accelerators. This is cheap insurance and the default posture for everything here.

9. Correctness over comfort¶

Every constant cites its source — CODATA 2018, IAU 2015, Oke & Gunn 1983 — so a reader can audit the number, not trust it. Every method is validated against an analytic result or a known answer. “It converged” and “it’s elegant” are not evidence. The radiation constant in this package is $a = 7.565733250\times10^{-15}\, \erg\,\mathrm{cm}^{-3}\,\mathrm{K}^{-4}$ precisely because it is derived as $4\sigma_\mathrm{SB}/c$ from the CODATA values, not rounded independently (see Release notes).

10. Vectorize and compose¶

Prefer vmap over Python loops, pure functions over mutable state, and immutable PyTrees (equinox modules) over in-place updates. Composition is what lets a foundation stay small: a handful of well-behaved primitives, combined, cover the ecosystem’s needs without each package reinventing them.

What we just established¶

These ten principles are not style preferences. They separate gradients that can support inference from expected-zero, blocked, surrogate, validation-only, and discrete paths with narrower claims. The rest of the theory section shows those boundaries in specific methods. Read on:

• Root-finding — fixed-iteration solvers, and the bisect zero-gradient caveat (principles 2, 3, 4).

• Newton–Cotes integration — Newton–Cotes integration and the dx-outside ordering (principles 5, 7).

• Fixed-node quadrature — fixed-node Gaussian and Clenshaw-Curtis quadrature plus cumulative Simpson panel sums (principles 7, 10).

• Cubic interpolation — cubic Hermite and PCHIP-style interpolation for smooth table evaluation without overshoot (principles 3, 4, 7).

• Regular-grid interpolation — static-rank multilinear interpolation for gridded tables with explicit boundary policy (principles 4, 7, 10).

• B-splines — local smooth basis functions for AD-friendly tabulated functions (principles 3, 7, 10).

• Linear algebra helpers — weighted fits, solve wrappers, covariance helpers, and positive-definite jitter for small dense problems (principles 3, 8, 9).

• Autodiff products — JVP, VJP, HVP, Gauss-Newton, and empirical Fisher-style products as named JAX-native helpers (principles 1, 9, 10).

• Geometry helpers — vector normalization, angular distances, rotations, quaternions, rigid transforms, and explicit composition helpers (principles 1, 9, 10).

• Spatial indexing and neighbor contracts — Morton and linear cells, capacity/overflow, approximate candidates, and exact fixed-radius neighbors as discrete preprocessing (principles 1, 6, 9).

• Distribution kernels — logpdf, CDF, and inverse-CDF kernels for normal, lognormal, finite power-law, and truncated-normal families (principles 3, 5, 7).

• Optimization helpers — robust residual losses, objective summaries, fixed-iteration line search, and convergence diagnostics (principles 1, 2, 10).

• Fixed-step ODE integration — fixed-step Euler, midpoint/RK2, RK4, and velocity-Verlet integration with scan-friendly gradient flow (principles 1, 2, 10).

• Linear operators — dense, diagonal, scaled, summed, composed, transposed, and block-diagonal linear operators as PyTrees (principles 1, 9, 10).

• Special functions — stable Planck kernels, normalized log weights, and orthogonal polynomial bases (principles 3, 5, 9).

• Random streams and resampling — explicit key streams, deterministic seed manifests, and systematic/stratified/residual resampling (principles 6, 9, 10).

• Grids and sampling utilities — log grids, conservative binning, and stratified uniforms (principles 7, 9, 10).

• Structured 1D meshes — structured 1D cell/face geometry, finite-volume stencils, and conservative cell-average remapping (principles 7, 9, 10).

• Quantities, units, and dimensional boundaries — exact dimensions, JAX PyTree quantities, parser canonicalization, bases, constants, equivalencies, and the raw-array boundary pattern (principles 1, 9, 10).

Then map principles to call signatures in API reference, and the design choices behind them in Decision log (ADRs).