MiniDiff: A minimal reference implementation of automatic differentiation in Rust

Companion repo for Automatic Differentiation: From Forward to Reverse in Small Steps

MiniDiff implements both forward and reverse mode automatic differentiation, and so enables differentiable programming in Rust.

It is meant to explain how automatic differentiation works, not as a crate to be used - though someone sufficiently motivated could develop it.

The development is largely based on https://simon.peytonjones.org/provably-correct/

Short guide to the repo

Read the article linked above for full details, the tl;dr is:

• playground.rs: introduces the concept of AD. Standalone, also at this Rust playground: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=3575c2c0ce68498e6f1528855c364e7c
• ad.rs: the implementations, and the intermediate steps from forward to reverse mode AD. If you're only interested in the "finished products" (i.e. not the intermediate steps, certainly don't mean they're production ready whatsoever):
  • The code marked as Step 1 is forward mode AD
  • The code marked as Step 3 and Step 4 together is reverse mode AD.
• example1.rs and example2.rs: Examples of usage of the various steps. Again usage of the "finished products":

  • Usage of forward mode AD

    fn f<T: VectorSpace>(t: &Dual<T>) -> Dual<T> {
        t.powi(2) + t + Dual::constant(1.0)
    }
    
    let res = f(&Dual::new(10.0, 1.0));
    println!("{res:?}");
    
    // prints Dual { primal: 111.0, tangent: 21.0 }
    

  • Usage of reverse mode AD and here in example2.rs

    fn f_sharing_fixed<'a>(x: &DualTrace<'a>, y: &DualTrace<'a>) -> DualTrace<'a> {
        let ref s = x * y;
        s + s
    }
    
    fn df_sharing_fixed(x: f64, y: f64) -> Vec<f64> {
        let trace = Trace::new();
        let x = &trace.var(x);
        let y = &trace.var(y);
        let dual_trace = f_sharing_fixed(x, y);
        eval(2, &dual_trace)
    }
    
    let res = df_sharing_fixed(3.0, 2.0);
    println!("{res:?}");
    
    // prints [4.0, 6.0]