djrakita / reverse

Zero-dep reverse mode automatic differentiation in Rust.

reverse

Zero-dependency crate for reverse mode automatic differentiation in Rust.

To use this in your crate, add the following to Cargo.toml:

[dependencies]
reverse = "0.2"

Examples

use reverse::*;

fn main() {
  let tape = Tape::new();
  let a = tape.add_var(2.5);
  let b = tape.add_var(14.);
  let c = (a.sin().powi(2) + b.ln() * 3.) - 5.;
  let gradients = c.grad();

  assert_eq!(gradients.wrt(&a), (2. * 2.5).sin());
  assert_eq!(gradients.wrt(&b), 3. / 14.);
}

The main type is Var<'a>, so you can define functions that take this as an input (possibly along with other f64 arguments) and also returns this as an output, and the function will be differentiable. For example:

use reverse::*;

fn main() {
    let tape = Tape::new();
    let params = tape.add_vars(&[5., 2., 0.]);
    let data = [1., 2.];
    let result = diff_fn(&params, &data);
    let gradients = result.grad();
    println!("{:?}", gradients.wrt(&params));
}

fn diff_fn<'a>(params: &[Var<'a>], data: &[f64]) -> Var<'a> {
    params[0].powf(params[1]) + data[0].sin() - params[2].asinh() / data[1]
}