echidna

A high-performance automatic differentiation library for Rust.

• Forward mode -- dual numbers (Dual<F>, DualVec<F, N>) with tangent propagation and batched JVP
• Reverse mode -- Adept-style two-stack tape with Copy AD variables (12 bytes for f64)
• Bytecode graph-mode AD -- record-once evaluate-many SoA tape with CSE, DCE, constant folding, and algebraic simplification
• Taylor-mode higher-order derivatives -- const-generic and arena-based dynamic implementations
• Sparse Jacobian/Hessian -- automatic sparsity detection and graph coloring
• GPU acceleration -- wgpu (Metal/Vulkan/DX12) and CUDA batch evaluation
• Nonsmooth AD -- branch tracking, kink detection, Clarke generalized Jacobians
• Gradient checkpointing -- binomial, online, disk-backed, and hint-guided strategies
• Cross-country elimination -- Markowitz vertex elimination for optimal Jacobian accumulation
• Stochastic Taylor derivative estimators -- Laplacian, Hessian diagonal, variance reduction
• Composable type nesting -- Dual<BReverse<f64>>, Taylor<BReverse<f64>, K>, Dual<Dual<f64>> for arbitrary-order differentiation

Feature Overview

Quick Start

Add to Cargo.toml:

[dependencies]
echidna = "0.1"

Gradient via reverse mode

use echidna::Scalar;

// Gradient of f(x) = x0^2 + x1^2
let g = echidna::grad(|x| x[0] * x[0] + x[1] * x[1], &[3.0, 4.0]);
assert!((g[0] - 6.0).abs() < 1e-10);
assert!((g[1] - 8.0).abs() < 1e-10);

// Write generic code that works with f64, Dual, and Reverse
fn rosenbrock<T: Scalar>(x: &[T]) -> T {
    let one = T::from_f(1.0);
    let hundred = T::from_f(100.0);
    let t1 = one - x[0];
    let t2 = x[1] - x[0] * x[0];
    t1 * t1 + hundred * t2 * t2
}

let g = echidna::grad(|x| rosenbrock(x), &[1.5, 2.0]);

BytecodeTape: record once, evaluate many

use echidna::{record, record_multi};

// Record a scalar function to tape
let (tape, value) = echidna::record(|x| x[0] * x[0] + x[1] * x[1], &[3.0, 4.0]);

// Re-evaluate at new points without re-recording
let g = tape.gradient(&[1.0, 2.0]);

// Hessian and Hessian-vector product
let (val, grad, hess) = tape.hessian(&[1.0, 2.0]);
let (grad, hvp) = tape.hvp(&[1.0, 2.0], &[1.0, 0.0]);

// Record a multi-output function, compute sparse Jacobian
let (mut tape, vals) = echidna::record_multi(
    |x| vec![x[0] * x[1], x[0] + x[1] * x[1]],
    &[1.0, 2.0],
);
let (vals, pattern, jac_vals) = tape.sparse_jacobian(&[1.0, 2.0]);

Taylor mode: higher-order derivatives

use echidna::{Taylor64, record};

// Record function, then compute Taylor coefficients via taylor_grad
let (tape, _) = echidna::record(|x| x[0].sin() + x[1].exp(), &[0.0, 0.0]);

// K=4 Taylor coefficients along direction v
let (output, adjoints) = tape.taylor_grad::<4>(&[0.0, 0.0], &[1.0, 0.0]);
// output.coeffs() gives [f(x), f'(x)*v, f''(x)*v^2/2!, ...]

Feature Flags

Enable features in Cargo.toml:

echidna = { version = "0.1", features = ["bytecode", "taylor"] }

API

Core (always available)

Bytecode Tape (requires bytecode)

Higher-Order (requires taylor or stde)

GPU (requires gpu-wgpu or gpu-cuda)

CUDA additionally supports _f64 variants of all methods.

Architecture

echidna uses a two-tier AD architecture:

1. Eager mode (Adept-style): Dual<F> and Reverse<F> use operator overloading with a thread-local tape managed by RAII. Reverse<F> stores precomputed partial derivatives (no opcode dispatch); the reverse sweep is a single multiply-accumulate loop. Reverse<f64> is Copy and 12 bytes.

2. Graph mode (BytecodeTape): BReverse<F> records operations to an SoA bytecode tape. The tape is recorded once and evaluated many times at different inputs. Optimization passes (CSE, DCE, algebraic simplification, constant folding) reduce tape size. The bytecode tape enables Hessians, sparse derivatives, GPU acceleration, checkpointing, nonsmooth AD, and cross-country elimination.

Types compose via nesting: Dual<BReverse<f64>> gives forward-over-reverse for Hessian-vector products, Taylor<BReverse<f64>, K> gives Taylor-over-reverse, and Dual<Dual<f64>> gives forward-over-forward.

Comparison

Enzyme operates at the LLVM IR level and is not a Rust-native library; it requires compiler integration. num-dual and the autodiff crate are pure Rust but limited to forward mode.

Performance

Benchmarks use Criterion and cover forward mode, reverse mode, bytecode tape, Taylor mode, STDE, GPU, and comparison with num-dual. Run with:

cargo bench                                         # Forward + reverse
cargo bench --features bytecode --bench bytecode    # Bytecode tape
cargo bench --features stde --bench taylor          # Taylor mode
cargo bench --features gpu-wgpu --bench gpu         # GPU

echidna-optim

The echidna-optim crate provides optimization solvers and implicit differentiation built on echidna:

Solvers: L-BFGS, Newton (Cholesky), trust-region (Steihaug-Toint CG), with Armijo line search.

Implicit differentiation: Differentiate through the fixed point of an optimization problem or nonlinear solve -- implicit_tangent, implicit_adjoint, implicit_jacobian, implicit_hvp, implicit_hessian.

Piggyback differentiation: piggyback_tangent_solve, piggyback_adjoint_solve, piggyback_forward_adjoint_solve.

Sparse implicit (with faer): implicit_tangent_sparse, implicit_adjoint_sparse, implicit_jacobian_sparse.

[dependencies]
echidna-optim = "0.1"

Optional features: parallel (enables rayon parallelism via echidna/parallel), sparse-implicit (sparse implicit differentiation via faer).

Development

cargo test                                    # Core tests
cargo test --features bytecode,taylor,stde    # All features
cargo bench                                   # Run benchmarks
cargo clippy                                  # Lint
cargo fmt                                     # Format

License

Licensed under either of

• Apache License, Version 2.0 (LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0)
• MIT license (LICENSE-MIT or http://opensource.org/licenses/MIT)

at your option.

Contributing

See CONTRIBUTING.md for guidelines.

Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in this project by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions.