xad-rs 0.3.6

Exact automatic differentiation for Rust — forward-mode, reverse-mode, first- and second-order, with named variable support for ergonomic gradient readback
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

//! Hessian computation via reverse-mode AD and dense forward-mode Dual2Vec.

use ndarray::Array2;

use crate::reverse::areal::AReal;
use crate::forward::dual2vec::Dual2Vec;
use crate::tape::{Tape, TapeStorage};
use crate::registry::VarRegistry;

/// Compute the Hessian of a scalar-valued function using repeated reverse-mode passes
/// with finite-difference perturbation of the gradient.
///
/// For each input direction i, computes the gradient at `x` and at `x + eps * e_i`,
/// then approximates `H[i][j] = (grad_j(x + eps*e_i) - grad_j(x)) / eps`.
#[allow(clippy::needless_range_loop)]
pub fn compute_hessian<T, F>(inputs: &[T], func: F) -> Vec<Vec<T>>
where
    T: TapeStorage,
    F: Fn(&[AReal<T>]) -> AReal<T>,
{
    let n = inputs.len();
    let eps = T::from(1e-7).unwrap();
    let inv_eps = T::one() / eps;
    let half = T::from(0.5).unwrap();
    let mut hessian = vec![vec![T::zero(); n]; n];

    // Compute base gradient
    let base_grad = compute_gradient(inputs, &func);

    // For each direction, perturb and compute gradient
    for i in 0..n {
        let mut perturbed = inputs.to_vec();
        perturbed[i] += eps;
        let perturbed_grad = compute_gradient(&perturbed, &func);

        for (j, h_ij) in hessian[i].iter_mut().enumerate() {
            *h_ij = (perturbed_grad[j] - base_grad[j]) * inv_eps;
        }
    }

    // Symmetrize off-diagonal entries.
    for i in 0..n {
        for j in (i + 1)..n {
            let avg = (hessian[i][j] + hessian[j][i]) * half;
            hessian[i][j] = avg;
            hessian[j][i] = avg;
        }
    }

    hessian
}

/// Compute the gradient of a scalar-valued function using reverse-mode AD.
fn compute_gradient<T, F>(inputs: &[T], func: &F) -> Vec<T>
where
    T: TapeStorage,
    F: Fn(&[AReal<T>]) -> AReal<T>,
{
    let mut tape = Tape::<T>::new(true);
    tape.activate();

    let mut ad_inputs: Vec<AReal<T>> = inputs.iter().map(|&v| AReal::new(v)).collect();
    AReal::register_input(&mut ad_inputs, &mut tape);

    let mut output = func(&ad_inputs);
    AReal::register_output(std::slice::from_mut(&mut output), &mut tape);

    output.set_adjoint(&mut tape, T::one());
    tape.compute_adjoints();

    let grad: Vec<T> = ad_inputs.iter().map(|x| x.adjoint(&tape)).collect();

    Tape::<T>::deactivate_all();
    grad
}

// ============================================================================
// Labeled Hessian (dense full Hessian via Dual2Vec)
// ============================================================================

/// Labeled dense full Hessian produced by [`compute_full_hessian`].
pub struct NamedHessian {
    pub vars: VarRegistry,
    pub value: f64,
    pub gradient: Vec<f64>,
    pub hessian: Array2<f64>,
}

/// Compute a named dense full Hessian via a single `Dual2Vec` forward pass.
pub fn compute_full_hessian<F>(inputs: &[(String, f64)], f: F) -> NamedHessian
where
    F: Fn(&[Dual2Vec]) -> Dual2Vec,
{
    let n = inputs.len();

    let names: Vec<String> = inputs.iter().map(|(name, _)| name.clone()).collect();
    let vars = VarRegistry::from_names(names);

    let d2v_inputs: Vec<Dual2Vec> = inputs
        .iter()
        .enumerate()
        .map(|(i, (_, value))| Dual2Vec::variable(*value, i, n))
        .collect();

    let output = f(&d2v_inputs);

    let value = output.value();
    let gradient: Vec<f64> = output.gradient().to_vec();
    let hessian: Array2<f64> = output.hessian().clone();

    NamedHessian {
        vars,
        value,
        gradient,
        hessian,
    }
}