Memory-bounded Jacobians

jax.jacfwd and jax.jacrev batch all basis directions by default. If a single derivative evaluation has large intermediates, full batching can exceed device memory even when the final Jacobian fits.

SOLVAX chunks the derivative basis while preserving JAX’s Jacobian layout.

Forward and reverse mode

For \(f:\mathbb{R}^n\to\mathbb{R}^m\):

• forward mode evaluates Jacobian columns and is usually favored when \(n\) is small relative to \(m\);

• reverse mode evaluates Jacobian rows and is usually favored when \(m\) is small relative to \(n\).

J_fwd = sx.chunked_jacfwd(f, chunk_size=8)(x)
J_rev = sx.chunked_jacrev(f, chunk_size=8)(x)
J_auto = sx.chunked_jacobian(f, mode="auto", chunk_size="auto")(x)

The output shape follows JAX:

\[ \operatorname{shape}(J)= \operatorname{shape}(f(x))+\operatorname{shape}(x). \]

argnums selects the differentiated positional argument.

Chunking model

Let \(q\) be the number of basis directions and \(c\) the chunk size. SOLVAX uses vmap inside each width-\(c\) block and jax.lax.map across blocks. A useful cost model is

\[ M(c)\approx M_0+cM_1, \qquad T(c)\approx T_0+\left\lceil\frac{q}{c}\right\rceil T_1. \]

Larger chunks expose more parallelism and use more memory. chunk_size=None is the unchunked extreme; chunk_size=1 minimizes batched intermediate state.

Automatic chunk size

chunk = sx.auto_chunk_size(
    dim=q,
    output_size=other_dimension,
    max_memory_bytes=None,
    element_bytes=8,
    memory_fraction=0.5,
)

With an explicit or device-reported memory budget, the policy chooses the largest estimated chunk fitting the chosen fraction. Without a usable device budget, it uses \(\lceil\sqrt q\rceil\), balancing chunk width and chunk count. This is a heuristic because JAX intermediate storage depends on the traced program, not only the final Jacobian.

General chunked mapping

ys = sx.chunk_map(expensive_function, xs, chunk_size=16)

xs may be an array or a pytree with a common leading axis. None performs a single vmap; an integer uses batched lax.map, including internal handling of a short final chunk. This helper is useful for parameter scans and batched local physics even when no Jacobian is formed.

Numerical equivalence

Chunking changes batching, not the JVP or VJP being evaluated. Results should match jax.jacfwd/jax.jacrev to normal floating-point tolerance. Different batching may change reduction order inside user code, so bitwise identity is not a portable guarantee.

Comparison with other memory strategies

Strategy	Saves memory by	Trade-off
chunked Jacobian	reducing simultaneous derivative directions	more sequential chunks
matrix-free JVP/VJP	never materializing the Jacobian	only operator actions are available
rematerialization/checkpointing	recomputing primal intermediates	extra primal work
finite differences	avoiding AD trace	truncation error and one solve per direction

If the consumer is a Krylov method, prefer a matrix-free JVP over materializing the full Jacobian. Use chunked Jacobians when the matrix itself is required by a direct factorization, export, or dense diagnostic.

Failure and tuning guidance

• mode="auto" selects by input/output sizes, not by the cost of the traced function. Benchmark both modes for unusual programs.

• A device memory statistic may be unavailable or may not reflect concurrent allocations. Pass an explicit budget for predictable jobs.

• Chunk size is static under jit; changing it produces a different compilation.

• Choose element_bytes consistent with the differentiated dtype.

API summary

• solvax.autodiff.chunk_map()

• solvax.autodiff.auto_chunk_size()

• solvax.autodiff.chunked_jacfwd()

• solvax.autodiff.chunked_jacrev()

• solvax.autodiff.chunked_jacobian()

Runnable counterpart: examples/15_chunked_jacobian.py.