Autodiff and refinement API

Chunked autodiff

Memory-chunked Jacobian construction (the jac_chunk_size knob).

The dense Jacobian of f : R^n -> R^m is assembled column by column (forward mode: one JVP against each input basis vector) or row by row (reverse mode: one VJP against each output basis vector). Evaluating all n (or m) directional derivatives in a single jax.vmap() — what jax.jacfwd() / jax.jacrev() do — replicates the intermediate program state that many times at once, so peak memory grows with the full Jacobian width. Chunking trades that peak for a modest slowdown: the basis is split into blocks of chunk_size, each block is vmapped, and the blocks are walked with jax.lax.map(), so peak memory scales with

memory ~ m0 + m1 * chunk_size, time ~ t0 + t1 * (n / chunk_size),

i.e. a knob between the fast/hungry chunk_size = n (plain jax.jacfwd/jacrev) and the slow/lean chunk_size = 1. This is the memory lever that makes otherwise-OOM optimization Jacobians (residual of a large parameter vector) and matrix-free operator materializations fit on a single accelerator; it is the analogue of DESC’s jac_chunk_size argument, factored out here for reuse across kinetic and equilibrium codes.

The chunked Jacobians are numerically identical to their JAX counterparts (the same JVP/VJP is evaluated for every basis vector; only the batching changes), and remain jit/vmap/grad-transparent.

Contract: fun maps one array argument (selected by argnums; arbitrary shape, flattened internally) to one array output (arbitrary shape). The returned Jacobian follows the JAX convention output_shape + input_shape. Pytree inputs/outputs are out of scope — ravel them (e.g. with jax.flatten_util.ravel_pytree()) before calling.

References

solvax.autodiff.chunk_map(fun, xs, *, chunk_size=None)

Map fun over the leading axis of xs in fixed-size chunks.

A thin wrapper choosing between a single wide jax.vmap() and a chunked jax.lax.map():

  • chunk_size is None: one jax.vmap over all len(xs) slices (maximum parallelism, maximum memory).

  • chunk_size = k: jax.lax.map(fun, xs, batch_size=k) — the leading axis is processed k slices at a time (vmapped within a chunk, scanned across chunks), so peak memory scales with k instead of len(xs). The final chunk is padded internally by lax.map and the padding discarded.

Parameters:

  • fun (Callable) – callable applied to a single leading-axis slice of xs.

  • xs (Array) – array (or pytree of arrays) with a common leading axis.

  • chunk_size (int | None) – chunk width, or None for a single vmap.

Returns:

The stacked results, leading axis equal to len(xs).

Return type:

Array

solvax.autodiff.auto_chunk_size(dim, output_size=1, *, max_memory_bytes=None, element_bytes=8, memory_fraction=0.5)

Pick a chunk width for a Jacobian with dim basis vectors.

Two regimes:

  • Memory budget (max_memory_bytes given, or a device limit is queryable): return the largest chunk whose block of directional derivatives fits the budget,

    chunk = floor(memory_fraction * budget / (output_size * element_bytes)),

    clamped to [1, dim] — DESC’s “largest that fits” semantics.

  • Heuristic (no budget available): return the square-root-balanced width ceil(sqrt(dim)), where the peak memory ~ chunk and the chunk count ~ dim / chunk are balanced, clamped to [1, dim].

Parameters:

  • dim (int) – number of basis vectors (input size for forward mode, output size for reverse mode).

  • output_size (int) – element count of one directional-derivative result (the other Jacobian dimension), used only in the budget regime.

  • max_memory_bytes (int | None) – explicit byte budget; if None, the first local device’s bytes_limit is used when available, else the heuristic regime applies.

  • element_bytes (int) – bytes per array element (8 for float64, 4 for float32).

  • memory_fraction (float) – fraction of the budget the Jacobian block may use.

Returns:

A chunk width in [1, dim] (returns 1 for dim <= 1).

Return type:

int

solvax.autodiff.chunked_jacfwd(fun, argnums=0, *, chunk_size='auto')

Forward-mode Jacobian assembled in column chunks.

A drop-in for jax.jacfwd() (single array argument, single array output) whose columns — one JVP per input basis vector — are evaluated chunk_size at a time, bounding peak memory. Forward mode is the efficient choice when the input is smaller than the output (tall Jacobian).

Parameters:

  • fun (Callable) – function of one array argument (position argnums) returning an array; other positional arguments are held fixed.

  • argnums (int) – index of the argument to differentiate.

  • chunk_size (int | str | None) – column-chunk width. "auto" (default) uses auto_chunk_size() on the input dimension; None reproduces jax.jacfwd() exactly (a single vmap); an int fixes the width.

Returns:

A callable jac(*args) -> Jacobian of shape output_shape + input_shape.

Return type:

Callable

solvax.autodiff.chunked_jacrev(fun, argnums=0, *, chunk_size='auto')

Reverse-mode Jacobian assembled in row chunks.

A drop-in for jax.jacrev() (single array argument, single array output) whose rows — one VJP per output basis vector — are evaluated chunk_size at a time, bounding peak memory. Reverse mode is the efficient choice when the output is smaller than the input (wide Jacobian), e.g. a scalar-ish residual of a large parameter vector.

Parameters:

  • fun (Callable) – function of one array argument (position argnums) returning an array; other positional arguments are held fixed.

  • argnums (int) – index of the argument to differentiate.

  • chunk_size (int | str | None) – row-chunk width. "auto" (default) uses auto_chunk_size() on the output dimension; None reproduces jax.jacrev() exactly; an int fixes the width.

Returns:

A callable jac(*args) -> Jacobian of shape output_shape + input_shape.

Return type:

Callable

solvax.autodiff.chunked_jacobian(fun, argnums=0, *, mode='rev', chunk_size='auto')

Memory-chunked Jacobian with forward/reverse/auto mode selection.

Dispatches to chunked_jacfwd() or chunked_jacrev().

Parameters:

  • fun (Callable) – function of one array argument returning an array.

  • argnums (int) – index of the argument to differentiate.

  • mode (str) – "rev" (default, mirrors jax.jacobian()), "fwd", or "auto" — pick forward mode when the input has no more elements than the output (fewer basis vectors), reverse otherwise.

  • chunk_size (int | str | None) – forwarded to the chosen builder.

Returns:

A callable jac(*args) -> Jacobian of shape output_shape + input_shape.

Return type:

Callable

Mixed-precision refinement

Mixed-precision iterative refinement (defect correction).

Given an approximate solver M ~ A^{-1} — typically a factorization or iterative solve carried out in low precision — classic iterative refinement recovers high-precision accuracy by repeated defect correction:

r_i = b - A x_i (residual, computed in high precision) d_i = M r_i (correction, cheap low-precision solve) x_{i+1} = x_i + d_i

Each sweep contracts the error by roughly u_f * kappa(A), where u_f is the unit roundoff of the precision used inside M and kappa(A) the condition number, until the residual stalls at the roundoff floor of the residual precision. Carson & Higham formalised the three-precision variant — factorization precision u_f, working precision u, and residual precision u_r (e.g. float32 / float64 / float64 here, or float16 / float32 / float64 on tensor-core hardware) — and showed convergence to working-precision accuracy for kappa(A) u_f < 1. This is the standard trick for exploiting fast low-precision hardware, and the recommended fallback when the unpivoted block eliminations in solvax.direct are used in weakly diagonally dominant regimes.

References

  • E. Carson & N. J. Higham, Accelerating the Solution of Linear Systems by Iterative Refinement in Three Precisions, SIAM J. Sci. Comput. 40(2), A817–A847 (2018), DOI 10.1137/17M1140819.

  • J. H. Wilkinson, Rounding Errors in Algebraic Processes, Prentice-Hall (1963) — the original fixed-precision analysis.

solvax.refine.iterative_refinement(matvec, b, approx_solve, *, iterations=3, residual_dtype=<class 'jax.numpy.float64'>)

Refine an approximate solve of matvec(x) = b by defect correction.

Runs x_0 = M(b) followed by iterations sweeps of x_{i+1} = x_i + M(b - A x_i), with residuals accumulated in residual_dtype. approx_solve may internally run in float32 (see as_low_precision()); the iterate and residual are kept in residual_dtype so the refined solution reaches high-precision accuracy whenever kappa(A) * u_low < 1.

Parameters:

  • matvec (Callable) – callable computing A @ x in (at least) residual_dtype.

  • b (Array) – right-hand side.

  • approx_solve (Callable) – callable approx_solve(r) -> d applying an approximate inverse of A; may be low precision internally.

  • iterations (int) – number of correction sweeps (static Python int).

  • residual_dtype – precision for iterates and residual accumulation.

Returns:

A pair (x, residual_norms) where x is the refined solution in residual_dtype and residual_norms has shape (iterations + 1,) — the 2-norm of b - A x_i after the initial solve and after each sweep, decreasing until it stalls at the residual_dtype roundoff floor.

Return type:

tuple[Array, Array]

solvax.refine.as_low_precision(solve, dtype=<class 'jax.numpy.float32'>)

Wrap a solve callable to run its inputs in a lower precision.

The returned callable casts every array argument down to dtype, calls solve, and casts the result back up to the original dtype of the first argument — the usual way to build the low-precision inner solve of iterative_refinement().

Parameters:

  • solve (Callable) – callable taking one or more arrays and returning an array (or pytree of arrays).

  • dtype – precision to run solve in.

Returns:

A callable with the same signature operating in dtype internally.

Return type:

Callable