Differentiable programming patterns

The flow solver is written once, as pure array functions, and JAX differentiates that code. Nothing is differentiated by hand: Newton’s Jacobian products, the adjoint’s transposed products and the sensitivities with respect to shape and flow conditions all come from the kernel that computes the residual, so they cannot fall out of step with it. The price is compilation time — tens of seconds for a new mesh size or polynomial degree — which is why sessions and a persistent cache exist. The patterns below are what a contributor has to keep intact for this to stay true.

Gyre’s numerical kernels are ordinary array code. What makes them a differentiable solver is a small set of structural patterns.

Pure kernels with explicit data#

Every kernel is a pure function of arrays. Operator tables (derivative matrices, gather indices, masks) live in one pytree, Discretization.ops, and are passed to jitted functions as runtime arguments:

python

f = jax.jit(lambda ops, U, geom, th: _residual(disc, ops, U, geom, th, None, True))

Arrays captured by a jitted closure are baked into the compiled program as constants. For index tables with millions of entries that bloats compile time and memory and defeats the persistent cache (the program text changes with every mesh). Passing them as arguments keeps the compiled program a function of shapes only.

Flow parameters and geometry are pytrees#

th is a flat dict of scalars; Geometry is a named tuple of arrays computed inside the traced function from node coordinates. Because both are pytrees, jax.grad, jvp and vjp differentiate with respect to them with no extra code:

python

def lagrangian(d, th, ops, U, psi):
    return J(U, d, th, ops) - jnp.sum(psi * R(U, d, th, ops))

dJ_dd, dJ_dth = jax.grad(lagrangian, argnums=(0, 1))(d, th, ops, U, psi)

One reverse sweep returns the sensitivity with respect to every surface node and every flow parameter.

Linearise inside the solve#

python

def solve(ops, U, geom, th, rhs, ...):                 # one jitted function
    _, lin = jax.linearize(lambda V: Rw(V), U)         # primal once
    A = lambda v: idm * v + lin(v)                     # tangent per iteration
    return fgmres(A, M, rhs, ...)                      # lax.while_loop

jax.linearize returns a function closing over the primal intermediates. Calling it from Python each Newton step would re-trace; jitting it would re-compile. Tracing the whole Krylov solve in one function avoids both: XLA computes the primal part once (it is loop-invariant) and the loop body is the tangent map only. The adjoint solve swaps in jax.vjp.

Split arguments for structure#

AD produces products, not matrices. To obtain the element-diagonal blocks, the residual takes the state twice (own / neighbour roles) and the blocks are jax.lax.map over unit tangents of the own-role argument, batched to bound memory. The same idea — route a dependency through a separate argument to select part of a Jacobian — is general and cheap.

Safe branches#

jnp.where(c, a, b) evaluates and differentiates both branches; a nan derivative in the unselected branch still poisons the result. The logarithmic mean therefore feeds its unselected branch safe substitute arguments (the double-where idiom in physics.ln_mean), and a unit test asserts finite forward- and reverse-mode derivatives on its diagonal. The other switches — the entropy fix of the Roe flux, the pressure-outlet and boundary-face selections, the shock-capturing ramp — are plain wheres whose branches all have finite derivatives; the total-pressure inlet clamps the static pressure before taking a root.

Compilation as a first-class cost#

Mechanism	Effect
persistent XLA cache (`~/.cache/gyre/jit`, `GYRE_JIT_CACHE`)	compiled kernels reused across processes; tracing remains
sessions	one process keeps its compiled functions for a whole design study
per-process device (`GYRE_DEVICE`)	CPU for small problems (faster compile), GPU for large
lazily built, keyed jit caches on `Discretization`, `Hierarchy`, `FlowDiscipline`	each entry point compiled once per object

For small cases cold compilation of the residual, the preconditioner set-up and the linear solve takes tens of seconds in total — noticeably longer on the GPU than on the CPU — while the compiled residual itself runs in milliseconds (measured kernel timings). With a warm persistent cache only tracing remains. The practical rules that follow: keep unit tests small and on the CPU, share solved fixtures across tests, and route repeated evaluations through sessions.

import gyre enables 64-bit floats (implicit solvers and adjoints need them), selects the device from GYRE_DEVICE, and disables GPU memory pre-allocation so several solver processes can share one GPU. Host CPUs without AVX2 must use jax < 0.12; scripts/bootstrap.sh detects this.