--- title: Differentiable programming patterns summary: How the solver is structured so that automatic differentiation yields exact, efficient linearisations — the handful of JAX patterns the code relies on, and the pitfalls each one avoids. section: engineering order: 2 level: 2 keywords: [JAX, automatic differentiation, linearize, vjp, jit, pytrees, compilation] --- **[L1 Essentials] In brief** 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)) ``` **[L3 Mechanics] Why not close over the tables?** 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 **[L3 Mechanics] The pattern** ```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 **[L3 Mechanics] Getting sparsity structure from AD** 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 **[L3 Mechanics] `where` differentiates both sides** `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 `where`s 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 | **[L4 Theory] What dominates** 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](performance.md)). 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. **[L2 Practice] Double precision and devices** `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.