echidna 0.8.2

A high-performance automatic differentiation library for Rust
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

//! Cross-check tests: verify ad-trait gradient matches echidna.

#![cfg(feature = "bytecode")]

use ad_trait::differentiable_function::{DifferentiableFunctionTrait, ForwardAD, ReverseAD};
use ad_trait::function_engine::FunctionEngine;
use ad_trait::AD;

use echidna::Scalar;

fn rosenbrock<T: Scalar>(x: &[T]) -> T {
    let one = T::from_f(<T::Float as num_traits::FromPrimitive>::from_f64(1.0).unwrap());
    let hundred = T::from_f(<T::Float as num_traits::FromPrimitive>::from_f64(100.0).unwrap());
    let mut sum = T::zero();
    for i in 0..x.len() - 1 {
        let t1 = one - x[i];
        let t2 = x[i + 1] - x[i] * x[i];
        sum = sum + t1 * t1 + hundred * t2 * t2;
    }
    sum
}

fn rosenbrock_ad<T: AD>(x: &[T]) -> T {
    let one = T::constant(1.0);
    let hundred = T::constant(100.0);
    let mut sum = T::constant(0.0);
    for i in 0..x.len() - 1 {
        let t1 = one - x[i];
        let t2 = x[i + 1] - x[i] * x[i];
        sum = sum + t1 * t1 + hundred * t2 * t2;
    }
    sum
}

#[derive(Clone)]
struct RosenbrockAD {
    n: usize,
}

impl<T: AD> DifferentiableFunctionTrait<T> for RosenbrockAD {
    const NAME: &'static str = "Rosenbrock";

    fn call(&self, inputs: &[T], _freeze: bool) -> Vec<T> {
        vec![rosenbrock_ad(inputs)]
    }

    fn num_inputs(&self) -> usize {
        self.n
    }

    fn num_outputs(&self) -> usize {
        1
    }
}

fn make_input(n: usize) -> Vec<f64> {
    (0..n).map(|i| 0.5 + 0.01 * i as f64).collect()
}

#[test]
fn ad_trait_forward_gradient_matches_echidna() {
    let n = 4;
    let x = make_input(n);

    let echidna_grad = echidna::grad(|v| rosenbrock(v), &x);

    let rosen_std = RosenbrockAD { n };
    let rosen_fwd = rosen_std.clone();
    let engine = FunctionEngine::new(rosen_std, rosen_fwd, ForwardAD::new());
    let (_values, jac) = engine.derivative(&x);

    for i in 0..n {
        let ad_trait_gi = jac[(0, i)];
        let echidna_gi = echidna_grad[i];
        assert!(
            (ad_trait_gi - echidna_gi).abs() < 1e-8,
            "forward mismatch at index {}: ad_trait={}, echidna={}",
            i,
            ad_trait_gi,
            echidna_gi
        );
    }
}

#[test]
fn ad_trait_reverse_gradient_matches_echidna() {
    let n = 4;
    let x = make_input(n);

    let echidna_grad = echidna::grad(|v| rosenbrock(v), &x);

    let rosen_std = RosenbrockAD { n };
    let rosen_rev = rosen_std.clone();
    let engine = FunctionEngine::new(rosen_std, rosen_rev, ReverseAD::new());
    let (_values, jac) = engine.derivative(&x);

    for i in 0..n {
        let ad_trait_gi = jac[(0, i)];
        let echidna_gi = echidna_grad[i];
        assert!(
            (ad_trait_gi - echidna_gi).abs() < 1e-8,
            "reverse mismatch at index {}: ad_trait={}, echidna={}",
            i,
            ad_trait_gi,
            echidna_gi
        );
    }
}