This document describes the mathematical foundations and algorithms implemented in echidna. Each section covers the theory, key references, and the corresponding echidna module.
For usage examples, see the crate-level documentation.
Forward-mode automatic differentiation is built on the algebra of dual numbers.
A dual number has the form a + b*epsilon where epsilon^2 = 0. Extending
every elemental operation f to dual numbers yields f(a + b*epsilon) = f(a) + f'(a)*b*epsilon, so the derivative is carried alongside the value
through every operation. This is called tangent propagation: at each step the
pair (f(x), f'(x) * x_dot) is propagated forward through the computation.
In echidna, Dual<F> stores the value in re and the tangent in eps.
For a function f: R^n -> R^m, one forward pass with a tangent seed vector
v computes the Jacobian-vector product (JVP) J * v. Computing the full
Jacobian requires n forward passes (one per input), so forward mode is most
efficient when n is small relative to m.
The batched variant DualVec<F, N> packs N tangent directions into a single
forward pass, amortizing the cost of value computation across multiple seed
vectors. This is particularly effective when combined with SIMD, as tangent
components map naturally to SIMD lanes.
References: Griewank & Walther, Evaluating Derivatives, Chapters 2--3.
Module: echidna::dual, echidna::dual_vec, echidna::api::jvp
Reverse-mode AD computes derivatives by recording a forward trace and then
propagating adjoint values backward. During the forward sweep, each operation
and its intermediate values are stored on a tape. The reverse sweep then
walks the tape in LIFO order, accumulating adjoint contributions via the chain
rule. For a scalar-valued function f: R^n -> R, one reverse sweep computes
the full gradient nabla f — the vector-Jacobian product (VJP) w^T * J
with w = 1. The Baur-Strassen result guarantees that the cost of one reverse
sweep is at most 5 times the cost of the forward evaluation, regardless of n.
echidna implements two reverse-mode backends. The Adept-style tape
(Reverse<F>, managed by echidna::tape) precomputes partial derivatives
during the forward sweep and stores them as scalar multipliers on a two-stack
tape — one stack for multipliers, one for operand indices. The reverse sweep
then consists entirely of multiply-accumulate operations with no opcode
dispatch, yielding tight inner loops. A constant sentinel value (u32::MAX)
marks operands that are literals or constants, avoiding tape entries for
non-differentiable quantities and preventing tape bloat. The bytecode reverse
backend (BReverse<F>) records opcodes for deferred evaluation through the
BytecodeTape. Both backends skip zero adjoints during the reverse sweep to
avoid unnecessary arithmetic in sparse dependency structures.
References: Griewank & Walther, Evaluating Derivatives, Chapter 4; Hogan, R.J. (2014), "Fast reverse-mode automatic differentiation using expression templates in C++", ACM TOMS.
Module: echidna::reverse, echidna::tape, echidna::breverse,
echidna::api::grad
The bytecode tape provides a reusable graph representation of a computation. Rather than re-tracing the function each time (as operator-overloading modes do), the tape is recorded once and can be evaluated repeatedly with different inputs. The tape uses a Structure-of-Arrays (SoA) layout: separate contiguous arrays for opcodes, operand indices, and constant values. This layout improves cache utilization — the opcode stream is compact enough to fit in L1 cache for moderate-sized tapes, and the separation allows independent vectorization of forward and reverse sweeps. The tape currently defines 44 opcodes covering all standard elemental operations (arithmetic, transcendental, power, comparison).
A key advantage of the graph representation is that it enables optimization
passes before evaluation. echidna implements four tape transformations: Common
Subexpression Elimination (CSE) merges duplicate operations sharing the same
opcode and operands; Dead Code Elimination (DCE) removes operations that do
not contribute to any output; algebraic simplification folds identity
patterns at recording time — x * 1 -> x and x / 1 -> x (exact for all
inputs) plus the sign-safe zero identities x + (-0.0) -> x and
x - (+0.0) -> x, while absorbing and self patterns (x * 0, x - x,
x / x) are deliberately NOT folded: their record-time value does not
determine the value at other inputs (x / x at x = 0 is 0/0 = NaN;
x * 0 and x - x at ±Inf are NaN), so folding would silently mask a
genuine singularity; and constant folding evaluates operations
on known constants at recording time. These optimizations are composable and
can be applied in sequence. The tape also enables second-order derivative
computation (Hessians), sparsity pattern detection, cross-country elimination,
and GPU offloading — all of which require a static graph structure.
References: Griewank & Walther, Evaluating Derivatives, Chapter 5 (computational graph representation).
Module: echidna::bytecode_tape, echidna::opcode, echidna::breverse
Higher-order derivatives in echidna are computed by nesting AD modes. The
canonical construction is forward-over-reverse: wrapping a reverse-mode type
inside a forward-mode dual number, Dual<BReverse<f64>>, produces a
Hessian-vector product (HVP) H * v in one combined forward-plus-reverse
pass. The forward seed v is set on the dual component; the forward sweep
records operations on the tape; the reverse sweep then produces adjoints whose
dual parts contain the HVP. To assemble the full n x n Hessian matrix, one
performs n such HVP computations, one per coordinate direction. The Hessian
is symmetric, so only the upper triangle needs to be computed, reducing the
work by nearly half.
The batched variant DualVec<BReverse<f64>, N> processes N tangent
directions simultaneously through a single tape recording, yielding N
columns of the Hessian per forward-plus-reverse pass. This is more efficient
than N separate Dual<BReverse<f64>> passes because the tape recording
overhead is paid only once. For sparse Hessians, graph coloring (see Section 6)
further reduces the number of required directions.
References: Griewank & Walther, Evaluating Derivatives, Chapter 8.
Module: echidna::api::hessian, echidna::api::hvp,
echidna::api::hessian_vec
Taylor mode propagates truncated Taylor series through a computation. Given a
function f and an input expanded as x(t) = x_0 + x_1*t + ... + x_K*t^K,
the output y(t) = f(x(t)) is also a Taylor series whose coefficients y_k
can be computed incrementally. For multiplication, the coefficients follow the
Cauchy product rule: (a*b)_k = sum_{j=0}^{k} a_j * b_{k-j}. For
transcendental functions (exp, log, sin, cos), echidna uses the logarithmic
derivative technique from Griewank Chapter 13, which expresses higher-order
coefficients through recurrence relations involving lower-order ones. This
avoids the need to differentiate the elemental functions symbolically.
echidna provides two Taylor types. Taylor<F, K> is const-generic and
stack-allocated: the degree K is a compile-time constant, enabling
monomorphization and eliminating heap allocation. TaylorDyn<F> is
arena-based (a thread-local, Vec-backed TaylorArena — bumpalo is not
involved), implements Copy, and allows the degree to be
chosen at runtime. Both types implement the full num_traits::Float trait and
echidna's Scalar trait, so they flow through BytecodeTape::forward_tangent
without any tape modifications. Shared propagation logic lives in
taylor_ops.rs.
Two important applications build on Taylor mode. ODE Taylor integration
(ode_taylor_step) computes the Taylor expansion of the solution y(t) to an
ODE y' = f(y) via coefficient bootstrapping: y_{k+1} = coeff_k(f(y(t))) / (k+1). This requires K-1 forward passes through the tape, building up the
Taylor inputs incrementally, and yields both the solution and error estimates
from the trailing coefficients. Reverse-over-Taylor (taylor_grad) combines
a forward Taylor sweep with a reverse adjoint sweep: the zeroth-order adjoint
gives the gradient, the first-order adjoint gives the HVP, and higher-order
adjoints give k-th order directional derivatives — all in a single pass.
References: Griewank & Walther, Evaluating Derivatives, Chapter 13.
Module: echidna::taylor, echidna::taylor_dyn, echidna::taylor_ops
When the Jacobian or Hessian has a known sparsity structure, computing only
the nonzero entries can be vastly cheaper than computing the full dense matrix.
echidna's sparse derivative pipeline has three stages: sparsity detection,
graph coloring, and compressed evaluation. Sparsity detection propagates
bitsets through the bytecode tape to determine which outputs depend on which
inputs, producing a JacobianSparsityPattern or SparsityPattern. Graph
coloring then assigns colors to columns (for Jacobians) or rows/columns (for
Hessians) such that columns sharing a nonzero in the same row receive
different colors. For Jacobians, this is a greedy distance-2 coloring; for
symmetric Hessians, it is a star bicoloring that exploits symmetry to halve
the required colors.
Compressed evaluation performs one forward or reverse sweep per color group.
Each sweep uses a seed vector that is the sum of coordinate basis vectors for
all columns in that color group. Because the columns in a group share no
nonzero row, their contributions can be separated by inspecting which row
positions are nonzero. The result is stored in CSR (Compressed Sparse Row)
format for efficient downstream use. The number of sweeps equals the chromatic
number of the coloring, which for structured problems (banded, block-diagonal)
is far smaller than n.
References: Gebremedhin, Manne & Pothen (2005), "What color is your Jacobian? Graph coloring for computing derivatives"; Griewank & Walther, Evaluating Derivatives, Chapter 7.
Module: echidna::sparse, echidna::api::sparse_jacobian,
echidna::api::sparse_hessian
Cross-country elimination computes Jacobians by performing vertex elimination
on the linearized computational graph — a DAG where edges carry the local
partial derivatives from the forward evaluation. The idea is to eliminate
intermediate vertices one at a time: when vertex v is eliminated, every pair
of incoming and outgoing edges is replaced by a single fill-in edge whose
weight is the product of the two original edge weights (an application of the
chain rule). After all intermediate vertices are eliminated, the remaining
edges directly connect input vertices to output vertices, and their accumulated
weights are the Jacobian entries.
The order in which vertices are eliminated affects the total operation count
due to fill-in. echidna uses Markowitz ordering: at each step, eliminate the
vertex with the smallest product |predecessors| x |successors|. This greedy
heuristic minimizes local fill-in and works well in practice. Cross-country
elimination can be cheaper than either pure forward mode (cost proportional to
n) or pure reverse mode (cost proportional to m) when the function has a
balanced input/output ratio. For extreme aspect ratios, it degenerates to one
of the pure modes.
References: Griewank & Walther, Evaluating Derivatives, Chapter 10.
Module: echidna::cross_country,
BytecodeTape::jacobian_cross_country
Reverse-mode AD requires storing the entire forward trace before the backward sweep can begin. For computations with many steps (long time integrations, deep recurrent networks), this memory requirement can be prohibitive. Checkpointing trades memory for recomputation: instead of storing every intermediate value, only selected checkpoints are stored, and the forward computation between checkpoints is re-executed during the backward sweep.
echidna implements binomial checkpointing via the Revolve algorithm. Given s
checkpoint slots and n time steps, Revolve produces an optimal schedule that
minimizes total recomputation. The schedule is provably optimal for the
binomial model (where each segment costs the same to recompute). Three
extensions address practical limitations. Online checkpointing
(grad_checkpointed_online) handles the case where the total step count is
unknown in advance: it uses periodic thinning, maintaining a fixed-size
checkpoint buffer and discarding every other entry when full, doubling the
spacing. The overhead is O(log N) recomputations for N steps. Disk-backed
checkpointing (grad_checkpointed_disk) stores checkpoint states as raw
binary files when the total checkpoint memory exceeds available RAM, using
unsafe byte transmutation (safe for all Float types since they are Copy + Sized with no pointers) and a DiskCheckpointGuard Drop guard for
panic-safe cleanup. Hint-based placement (grad_checkpointed_with_hints)
accepts user-specified required checkpoint positions and distributes the
remaining slots proportionally across the resulting sub-intervals using the
largest-remainder method, then runs Revolve on each sub-interval independently.
References: Griewank & Walther, Evaluating Derivatives, Chapter 12; Griewank & Walther (2000), "Algorithm 799: Revolve".
Module: echidna::checkpoint
Full second-order derivative tensors are expensive to compute and store. For
many applications — Laplacians in PDE solvers, trace estimation in
optimization, Hessian diagonal preconditioning — only aggregate quantities
are needed. Stochastic Taylor Derivative Estimators use random jet propagation
to estimate these quantities without materializing the full Hessian. The core
idea is the Hutchinson trace estimator: for a random vector v drawn from a
distribution with identity covariance, E[v^T H v] = tr(H), where H is the
Hessian. By pushing Taylor<F, 3> inputs with a random tangent direction
through the bytecode tape, the second-order Taylor coefficient c2 at the
output yields v^T H v / 2 for that direction (the estimators multiply by
two to recover v^T H v). Averaging over S random directions
gives the Laplacian estimate tr(H) ~ (1/S) * sum_s v_s^T H v_s.
echidna provides several estimators built on this primitive. laplacian
computes the Hutchinson trace estimate using user-provided direction vectors
(Rademacher or Gaussian — no rand dependency). hessian_diagonal computes
exact diagonal entries by using coordinate basis vectors as tangent directions.
directional_derivatives returns batch second-order directional derivatives.
Variance reduction is available via laplacian_with_control, which uses the
Hessian diagonal as a control variate: since the diagonal is computed exactly,
it can be subtracted from the stochastic estimate and added back, reducing
variance to match Rademacher performance even when using Gaussian vectors.
laplacian_with_stats uses Welford's online algorithm to track sample
variance and standard error alongside the running estimate.
Higher-order diagonal estimation. The coordinate-basis approach
generalises beyond k=2. diagonal_kth_order pushes basis vectors through
order-(k+1) TaylorDyn jets to extract exact [∂^k u/∂x_j^k] for all
coordinates. The output coefficient at index k stores ∂^k u/∂x_j^k / k!,
so the derivative is k! * coeffs[k]. The stochastic variant subsamples
coordinate indices for an unbiased estimate of the diagonal sum.
Parabolic PDE σ-transform. For parabolic PDEs with diffusion matrix σ,
the operator ½ tr(σσ^T · H) can be rewritten as ½ Σ_i (σ·e_i)^T H (σ·e_i)
where each term is a directional second derivative along σ's columns. This
avoids forming off-diagonal Hessian entries and reduces to d second-order
Taylor pushforwards (d = number of columns of σ).
Sparse STDE. For general differential operators L = Σ C_α D^α, the
coefficient tensor C(L) defines a discrete distribution over sparse k-jets.
Each term α is sampled with probability |C_α| / Z where Z = Σ|C_α|, and
the per-sample estimator sign(C_α) · Z · D^α u is unbiased for Lu. The
derivative D^α u is extracted from a single forward pushforward via the
jet extraction formula from Section 10. This implements the core contribution
of Shi et al. (2024): it reduces the cost of estimating arbitrary k-th order
operators from O(n^k) (exact computation) to O(S) (S random samples), where
each sample costs one forward pushforward.
References: Shi et al. (2024), "Stochastic Taylor Derivative Estimators", NeurIPS 2024 Best Paper; Hutchinson (1990), "A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines".
Module: echidna::stde
The STDE module (Section 9) uses random tangent directions through low-order
Taylor jets to estimate aggregate second-order quantities like the Laplacian.
The diffop module generalises this idea to extract exact arbitrary mixed
partial derivatives from a single jet pushforward per group of active
variables.
The mathematical foundation is the multivariate Faà di Bruno formula applied
to a parameterised curve. Given a function u: R^n -> R and a target mixed
partial ∂^Q u / (∂x_{i₁}^{q₁} ... ∂x_{iT}^{qT}), assign each active
variable i_t a distinct integer "slot" j_t. Construct Taylor inputs where
variable i_t has coefficient 1/j_t! at index j_t (all other higher-order
coefficients zero), then push through the tape via forward_tangent. The
output jet coefficient at index k = Σ j_t · q_t equals the target derivative
divided by the prefactor Π_t (q_t! · (j_t!)^{q_t}). This prefactor absorbs
both the Taylor scaling (1/k! at the output) and the input scaling (1/j_t!
per variable).
The slot assignment must satisfy a uniqueness condition: the target partition
of k into parts drawn from the assigned slots must be the only such
partition. If another partition exists, it contaminates the output coefficient
and the extraction is ambiguous. echidna uses a prime window sliding strategy:
assign slots from consecutive primes PRIMES[offset..offset+T] and verify
uniqueness by exhaustive partition enumeration. If a collision is detected,
increment the offset and retry. For practical PDE operators (total order ≤ 6,
≤ 4 active variables), convergence is fast.
Multi-indices with different sets of active variables receive separate pushforwards ("pushforward groups"). This is essential: if slots from non-active variables are present in the Taylor input, their contributions create partition collisions that cannot be resolved by slot reassignment alone. Multi-indices sharing the same active variable set are batched into a single group, since their slot assignments are compatible and multiple output coefficients can be read from one pushforward.
The JetPlan type precomputes all slot assignments, jet orders, and extraction
prefactors. Once built, a plan is reused across evaluation points — only the
Taylor input values change, not the combinatorial structure. This
plan-once-evaluate-many design mirrors the BytecodeTape philosophy and is
well suited to PDE solvers that evaluate the same differential operator at
many grid points.
The DiffOp type represents a linear differential operator L = Σ C_α D^α
as an explicit list of (coefficient, multi-index) pairs. It provides exact
evaluation via DiffOp::eval (delegating to JetPlan internally) and can
build a SparseSamplingDistribution for stochastic estimation. Convenience
constructors cover common operators: laplacian(n), biharmonic(n),
diagonal(n, k). Inhomogeneous operators (e.g., u_t + ∇²u) can be
decomposed into homogeneous groups via split_by_order(), since the sparse
sampling distribution requires all terms to share the same total order.
References: Shi et al. (2024), "Stochastic Taylor Derivative Estimators", NeurIPS (Section 3, jet extraction formula); Griewank & Walther, Evaluating Derivatives, Chapter 13 (Taylor propagation).
Module: echidna::diffop
Standard AD assumes that all elemental operations are smooth (continuously
differentiable). In practice, many functions contain piecewise-linear
operations — abs, min, max — and step functions — signum, floor,
ceil, round, trunc — that introduce kinks or discontinuities where the
classical derivative does not exist. echidna's nonsmooth AD module detects
these kinks during forward evaluation and provides access to the Clarke
generalized derivative, which is the appropriate generalized derivative concept
for Lipschitz-continuous functions.
BytecodeTape::forward_nonsmooth performs a standard forward evaluation while
recording a KinkEntry for each of the 9 nonsmooth operations. Each entry
stores the tape index, opcode, switching value, and the branch that was taken.
For abs/min/max/signum, the switching value is the operand value
(kink at zero). For floor/ceil/trunc/fract, the switching value is
the distance to the nearest integer (a - a.round()), which is zero exactly
at the discontinuity points; round kinks at half-integers instead, so its
switching value is shifted by 0.5. The resulting NonsmoothInfo provides
active_kinks(tol) to query kinks within a given tolerance,
is_smooth(tol) to check whether the evaluation point lies in a smooth
region, and signature() to characterize the combinatorial branch structure.
echidna uses a two-tier classification for nonsmooth operations.
is_nonsmooth() returns true for all 9 ops and is used for proximity
detection — knowing whether the evaluation point is near any kink.
has_nontrivial_subdifferential() returns true only for abs/min/max,
where forced branch choices produce distinct partial derivatives (e.g.,
abs has slope +1 vs -1). Step functions (signum/floor/ceil/round/
trunc) have zero derivative on both sides of the kink, and fract has an
identical unit derivative on both sides (a value discontinuity rather than a
derivative kink); in both cases enumerating forced branches would add 2^k
cost with no new information.
When kinks with nontrivial subdifferentials are present,
BytecodeTape::clarke_jacobian enumerates all 2^k sign combinations for
k active nontrivial kinks, computing a limiting Jacobian for each
combination by forcing the sign choices in the nonsmooth operations and
running a standard reverse sweep. Step-function kinks are filtered before
enumeration. The convex hull of these limiting Jacobians is the Clarke
generalized Jacobian. A configurable kink limit (default 20) prevents
combinatorial explosion by returning ClarkeError::TooManyKinks when
exceeded.
References: Griewank & Walther, Evaluating Derivatives, Chapter 14; Clarke, F.H. (1983), Optimization and Nonsmooth Analysis.
Module: echidna::nonsmooth,
BytecodeTape::forward_nonsmooth,
BytecodeTape::clarke_jacobian
Laurent series extend Taylor series to handle singularities. A Laurent series
has the form sum_{k=p}^{K} a_k * t^k where the pole order p may be
negative, allowing representation of functions with poles. Laurent<F, K> is
a const-generic, stack-allocated, Copy type that stores K coefficients
along with a tracked pole order. Arithmetic operations reuse the Cauchy
product and division rules from Taylor propagation (taylor_ops), with
additional bookkeeping for pole order: multiplication adds pole orders,
division subtracts them, and the result is always normalized so the leading
coefficient is nonzero.
The type handles degenerate cases deliberately: essential singularities (e.g.,
exp or sin applied to a value with a pole) produce NaN, since the Laurent
expansion does not converge in such cases. Division by zero also produces NaN.
The value() method returns positive or negative infinity for genuine poles,
matching the expected mathematical behavior. Laurent<F, K> implements the
full num_traits::Float trait and echidna's Scalar trait, so it flows
through BytecodeTape::forward_tangent for automated singularity detection
in recorded computations.
References: Standard complex analysis; the implementation follows Taylor propagation conventions from Griewank & Walther, Chapter 13, extended with pole tracking.
Module: echidna::laurent
Many problems involve quantities z* defined implicitly by an equation
F(z*, x) = 0, where z* is the solution and x is the parameter. The
Implicit Function Theorem (IFT) states that if F_z (the Jacobian of F
with respect to z) is nonsingular at (z*, x), then the derivative of z*
with respect to x exists and can be computed without differentiating through
the solver. The tangent-mode formula is dz*/dx = -(F_z)^{-1} * F_x:
compute the partial Jacobians F_z and F_x via AD, then solve a linear
system. The adjoint-mode formula is dL/dx = -lambda^T * F_x where
F_z^T * lambda = dL/dz, which requires only one linear solve regardless
of the number of parameters x. The full Jacobian dz*/dx can be assembled
by iterating either the tangent formula (column by column) or the adjoint
formula (row by row).
Second-order implicit derivatives are computed via nested dual numbers
Dual<Dual<F>>. A single forward pass through forward_tangent with doubly-
nested dual inputs extracts the second-order correction from the .eps.eps
component, which encodes the Hessian-direction-direction product. The function
implicit_hvp solves the resulting linear system to produce the implicit HVP,
and implicit_hessian builds the full second-order tensor by iterating over
n(n+1)/2 direction pairs with LU factor reuse.
Piggyback differentiation handles the case where z* is found by iterating
a fixed-point map z_{k+1} = G(z_k, x) rather than by a direct solve. The
tangent piggyback propagates z_dot_{k+1} = G_z * z_dot_k + G_x * x_dot
alongside the primal iteration using a dual-number forward pass at each step.
The adjoint piggyback iterates lambda_{k+1} = G_z^T * lambda_k + z_bar at
the converged z* via reverse sweeps. The interleaved variant runs both the
primal and adjoint iterations in a single loop, reducing the total iteration
count from K_primal + K_adjoint to max(K_primal, K_adjoint).
For large systems where F_z is sparse, the sparse-implicit feature
exploits structural sparsity. SparseImplicitContext precomputes the full
sparsity pattern and graph coloring, then uses echidna's sparse Jacobian
infrastructure for compressed evaluation and faer's sparse LU for the solve.
References: Griewank & Walther, Evaluating Derivatives, Chapter 15; Christianson, B. (1994), "Reverse accumulation and attractive fixed points".
Module: echidna_optim::implicit, echidna_optim::piggyback,
echidna_optim::sparse_implicit
Batched tape evaluation is naturally parallel: the same sequence of operations is applied independently to many input vectors. echidna exploits this by offloading the bytecode tape evaluation to the GPU. The CPU-side code handles sparsity detection and graph coloring (which are sequential, irregular algorithms), then dispatches batched tangent sweeps to the GPU for the numerically intensive phase. Custom operations are rejected at tape upload time since they cannot be compiled to GPU code.
The wgpu backend (feature gpu-wgpu) provides cross-platform GPU access via
Metal, Vulkan, and DX12. It ships four hand-written WGSL compute shaders:
forward (batched forward evaluation), reverse (batched adjoint sweep),
tangent_forward (forward tangent for JVP and sparse Jacobian), and
tangent_reverse (forward-over-reverse for HVP and sparse Hessian); a fifth
kernel, taylor_forward_2nd (batched second-order Taylor forward propagation
for STDE), is generated at runtime by gpu/taylor_codegen.rs when the stde
feature is enabled. Of the 44 opcodes, the 41 computational ones appear as
dispatch cases in each shader; Input and Const are structural leaf loads,
and Custom is rejected at tape upload — 43 of 44 opcodes are GPU-supported.
The backend is limited to f32 due to
WGSL's lack of native f64 support. The CUDA backend (feature gpu-cuda)
targets NVIDIA GPUs and supports both f32 and f64. It uses a single templated
CUDA kernel file compiled via NVRTC at runtime, instantiated for both float
and double. The WgpuContext and CudaContext types provide matching API
surfaces: forward_batch, gradient_batch, sparse_jacobian, hvp_batch,
and sparse_hessian.
When the stde feature is enabled alongside a GPU backend, the generated
taylor_forward_2nd kernel (introduced above) provides batched second-order
Taylor forward propagation. Each thread processes one batch element,
propagating a (c0, c1, c2) triple through the tape where c0 is the primal
value, c1 is the directional first derivative, and c2 = v^T H v / 2. The
memory layout is interleaved per variable per thread: [c0_i, c1_i, c2_i]
contiguous per variable, optimizing per-thread locality. Every dispatched
computational opcode implements full
second-order Taylor propagation rules: arithmetic via Cauchy products,
transcendentals via logarithmic derivative recurrences, and coupled
recurrences for sin/cos, sinh/cosh, and tan/tanh.
The high-level functions in stde_gpu flatten directions into batch format,
dispatch the Taylor kernel, and perform CPU-side Welford aggregation:
laplacian_gpu— Hutchinson trace estimator: pushes S random directions through the kernel, computesmean(2 * c2)as the Laplacian estimate.hessian_diagonal_gpu— exact diagonal: uses n standard basis vectors as directions, extractsdiag(H)[j] = 2 * c2[j].laplacian_with_control_gpu— subtracts the exact diagonal contribution (control variate) to reduce estimator variance.
Buffer constraint: the working buffer requires 3 * num_variables * batch_size * 4 bytes. For WebGPU's default 128 MB storage buffer limit,
max num_variables * batch_size ≤ ~10M.
References: General GPU computing literature; the tape evaluation model
follows the graph-mode execution paradigm common in ML frameworks, adapted
for AD-specific operations. The Taylor propagation rules follow the same
recurrences as the CPU implementation in taylor_ops.rs.
Module: echidna::gpu (echidna::gpu::wgpu_backend,
echidna::gpu::cuda_backend, echidna::gpu::stde_gpu)
echidna's AD modes are designed to compose via type-level nesting. The key
insight is that Dual<F>, Taylor<F, K>, and DualVec<F, N> are all generic
over F: Float, and Reverse<F> and BReverse<F> both implement Float.
This enables compositions like Dual<BReverse<f64>> (forward-over-reverse =
HVP), Taylor<BReverse<f64>, K> (Taylor-over-reverse = higher-order
adjoints), DualVec<BReverse<f64>, N> (batched tangent-over-reverse =
multiple HVP directions per tape pass), and Dual<Dual<BReverse<f64>>>
(triple nesting for second-order implicit derivatives). The Float and
IsAllZero trait implementations for Reverse<F> and BReverse<F> are what
make these compositions possible: they allow reverse-mode types to serve as
the scalar type inside forward-mode wrappers.
The convenience function composed_hvp provides a one-shot interface for
forward-over-reverse HVP without requiring the user to manually set up nested
types and tapes. Internally, it constructs Dual<BReverse<f64>> inputs,
evaluates the function, and extracts the HVP from the adjoint dual components.
Reverse-wrapping-forward (BReverse<Dual<f64>>) is also supported via
BtapeThreadLocal implementations for Dual<f32> and Dual<f64>, enabling
reverse-over-forward HVP and column-by-column Hessian computation.
References: Griewank & Walther, Evaluating Derivatives, Chapter 8 (higher-order techniques via mode composition).
Module: echidna::api::composed_hvp, echidna::float
The echidna-optim crate provides derivative-based optimization algorithms
that use echidna's AD capabilities for gradient and Hessian computation. All
solvers accept functions via the Objective or TapeObjective traits, which
abstract over the function evaluation and derivative computation interface.
L-BFGS is a limited-memory quasi-Newton method that approximates the
inverse Hessian using the last m gradient-step pairs (typically m = 5..20).
The two-loop recursion computes the search direction H_k * g_k in O(m*n)
time without forming the matrix explicitly. Combined with Armijo backtracking
line search (sufficient decrease condition), L-BFGS requires only gradient
information and is the workhorse method for large-scale unconstrained
optimization.
Newton's method computes the exact Hessian at each iteration via
forward-over-reverse AD and solves the Newton equation H * d = -g using
LU factorization with partial pivoting, falling back to steepest descent
when the Hessian is indefinite. This yields quadratic convergence near the
solution but requires O(n^2) storage and O(n^3) factorization per step,
making it practical only for moderate dimensions.
Trust-region uses the Steihaug-Toint truncated conjugate gradient method
to approximately solve the trust-region subproblem min_d m(d) s.t. ||d|| <= Delta, where m(d) is the quadratic model. The trust radius Delta is
adapted based on the ratio of actual to predicted reduction: expanded when the
model is accurate, contracted when it overpredicts. This method is more robust
than line search for non-convex problems and handles indefinite Hessians
gracefully (CG detects negative curvature and stops early).
References: Nocedal & Wright, Numerical Optimization, Chapters 6--7 (L-BFGS, trust region); standard optimization textbooks.
Module: echidna_optim::solvers::lbfgs,
echidna_optim::solvers::newton,
echidna_optim::solvers::trust_region