aquakin.DifferentiationConfig#

class aquakin.DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False)[source]#

Bases: object

How a gradient / sensitivity flows through a solve.

A single frozen value object replacing the old AD selectors (adjoint=, gradient=, ad_mode=, check_finite=, stable_adjoint_max_steps=). Passed as diff=DifferentiationConfig(...).

Parameters:

  • mode ({"reverse", "forward"}) – Autodiff direction. "reverse" (default) is the calibration-gradient direction (jax.grad / a discrete adjoint); "forward" is the sensitivity-screen direction (jax.jacfwd / a variational solve).

  • method ({"stable", "through_solve"}) –

    How the adjoint / sensitivity is formed for the chosen mode.

    • mode="reverse", method="stable" (the default) – a plant uses the cap-free hand-written discrete adjoint (the old gradient="stable_adjoint" / "auto" routing: a concrete forward solve still takes the fast cached path, a differentiated solve takes the discrete adjoint). A reactor has no stable reverse adjoint, so it always uses RecursiveCheckpointAdjoint regardless of method.

    • mode="reverse", method="through_solve" – differentiate through the diffrax solve (RecursiveCheckpointAdjoint; the old gradient="jax_adjoint"). Needs a dtmax cap for a stiff model.

    • mode="forward", method="stable" – the augmented [y; S] variational solve (cap-free, exact). Routed through solve_sensitivity / augmented_forward_sensitivity.

    • mode="forward", method="through_solve" – jax.jvp / jacfwd via a forward-capable adjoint (DirectAdjoint / ForwardMode).

  • check_finite (bool) – Raise a friendly error if a freshly computed gradient/Jacobian is non-finite, instead of returning silent NaN values.

  • adjoint_max_steps (int) – Buffer length for the discrete-adjoint backward scan (the old stable_adjoint_max_steps); set it to a tight upper bound on the forward step count.

  • adjoint_low_memory (bool) – Low-memory mode for the method="stable" discrete adjoint. False (default) saves each forward step’s dense-output stage increments, so the backward reconstructs the stages cheaply – but the saved buffer is ~n_stages × the trajectory. True stores only the step states and recomputes each step’s stages in the backward pass (a per-step Newton scan), so the dense buffer is never allocated – trading compute for memory when that buffer is the binding constraint on a long, large-state plant solve. The recompute matches the saved-stage gradient (it is the same well-conditioned per-stage solve), so it is a pure memory/compute trade. Reverse method="stable" only – ignored for a through_solve or forward solve.

__init__(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False)#
Parameters:

  • mode (str)

  • method (str)

  • check_finite (bool)

  • adjoint_max_steps (int)

  • adjoint_low_memory (bool)

Return type:

None

Methods

__init__([mode, method, check_finite, ...])

forms_jacfwd()

Whether the AD direction is forward (jax.jacfwd / a variational solve) rather than reverse (jax.grad / a discrete adjoint).

gradient_backend()

The reverse-adjoint backend this config selects.

reactor_adjoint()

The diffrax adjoint object a reactor / plant solve should carry.

validated()

Validate the mode / method vocabulary; return self.

Attributes

adjoint_low_memory

adjoint_max_steps

check_finite

method

mode

validated()[source]#

Validate the mode / method vocabulary; return self.

Raises ValueError for an out-of-range mode or method. This is the vocabulary check only – whether a particular consumer serves a given valid pairing (e.g. Plant.solve rejecting the forward + stable augmented solve) is that consumer’s own guard.

Return type:

DifferentiationConfig

forms_jacfwd()[source]#

Whether the AD direction is forward (jax.jacfwd / a variational solve) rather than reverse (jax.grad / a discrete adjoint).

Return type:

bool

gradient_backend()[source]#

The reverse-adjoint backend this config selects.

"stable_adjoint" for the cap-free hand-written discrete adjoint (method="stable"), or "jax_adjoint" for differentiating through the diffrax solve (method="through_solve"). This is the canonical method -> backend meaning the calibration entry points consume; the plant solve routes method="stable" through its own "auto" dispatch instead (see Plant._resolve_diff_config()).

Return type:

str

reactor_adjoint()[source]#

The diffrax adjoint object a reactor / plant solve should carry.

Reactors and the plant have no hand-written (“stable”) reverse adjoint on the reactor object – reverse mode always uses the default RecursiveCheckpointAdjoint (signalled by None). method matters only for forward mode: through_solve -> a forward-capable DirectAdjoint (forward_adjoint()); stable -> the augmented solve_sensitivity backend, invoked explicitly (not via this adjoint), so None is still returned.