tenferro-ad 0.8.0

Eager runtime, eager tensors, and traced AD extension traits for tenferro.
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

use std::sync::Arc;

use crate::{EagerRuntime, EagerTensor, Tensor};

#[test]
fn backward_shares_one_vjp_across_leaves_and_separates_active_masks() {
    let ctx = EagerRuntime::new().unwrap();
    let leaf = |values| {
        EagerTensor::requires_grad_in(
            Tensor::from_vec_col_major(vec![2], values).unwrap(),
            Arc::clone(&ctx),
        )
        .unwrap()
    };
    let x = leaf(vec![2.0_f64, 3.0]);
    let y = leaf(vec![5.0_f64, 7.0]);
    let unused = leaf(vec![11.0_f64, 13.0]);
    let loss = ctx
        .with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
            let product = s.mul(&x, &y)?;
            s.reduce_sum(&product, Some(&[0]))
        })
        .unwrap();
    let first = loss
        .backward(&mut tenferro_cpu::CpuEngine::serial())
        .unwrap();
    assert_eq!(first.len(), 2);
    assert_eq!(
        x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[5.0, 7.0]
    );
    assert_eq!(
        y.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[2.0, 3.0]
    );
    assert!(unused.grad().unwrap().is_none());
    let stats = ctx.cache_stats().unwrap().prepared_derivatives;
    assert_eq!(stats.entries, 1, "all leaves must share one prepared VJP");
    assert_eq!(stats.misses, 1);

    // A single-input VJP must not reuse the two-input derivative's output map.
    let gx = ctx
        .grad(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x)
        .unwrap();
    assert_eq!(gx.value().unwrap().as_slice::<f64>().unwrap(), &[5.0, 7.0]);
    assert_eq!(ctx.cache_stats().unwrap().prepared_derivatives.entries, 2);
    let before = ctx.cache_stats().unwrap().prepared_derivatives;
    let second = loss
        .backward(&mut tenferro_cpu::CpuEngine::serial())
        .unwrap();
    assert_eq!(second.len(), 2);
    let after = ctx.cache_stats().unwrap().prepared_derivatives;
    assert_eq!(after.hits, before.hits + 1);
    assert_eq!(after.misses, before.misses);
    assert_eq!(
        x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[10.0, 14.0]
    );
    assert_eq!(
        y.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[4.0, 6.0]
    );
}