use std::cell::RefCell;
use std::ops::{Add, Div, Mul, Neg, Sub};
#[derive(Clone, Debug)]
struct TapeOp {
output_idx: usize,
inputs: [usize; 2],
n_inputs: usize,
local_grads: [f64; 2],
}
pub struct Tape {
ops: RefCell<Vec<TapeOp>>,
values: RefCell<Vec<f64>>,
grads: RefCell<Vec<f64>>,
}
impl Default for Tape {
fn default() -> Self {
Self::new()
}
}
impl Tape {
pub fn new() -> Self {
Self {
ops: RefCell::new(Vec::new()),
values: RefCell::new(Vec::new()),
grads: RefCell::new(Vec::new()),
}
}
fn push_leaf(&self, value: f64) -> usize {
let mut vals = self.values.borrow_mut();
let mut grds = self.grads.borrow_mut();
let idx = vals.len();
vals.push(value);
grds.push(0.0);
idx
}
fn push_op(
&self,
value: f64,
inputs: [usize; 2],
n_inputs: usize,
local_grads: [f64; 2],
) -> usize {
let idx = self.push_leaf(value);
self.ops.borrow_mut().push(TapeOp {
output_idx: idx,
inputs,
n_inputs,
local_grads,
});
idx
}
pub fn backward(&self, loss: Var<'_>) {
{
let mut grds = self.grads.borrow_mut();
grds[loss.idx] = 1.0;
}
let ops = self.ops.borrow();
let mut grds = self.grads.borrow_mut();
for op in ops.iter().rev() {
let out_grad = grds[op.output_idx];
for i in 0..op.n_inputs {
grds[op.inputs[i]] += op.local_grads[i] * out_grad;
}
}
}
pub fn reset(&self) {
self.ops.borrow_mut().clear();
self.values.borrow_mut().clear();
self.grads.borrow_mut().clear();
}
pub fn values_len(&self) -> usize {
self.values.borrow().len()
}
pub fn ops_len(&self) -> usize {
self.ops.borrow().len()
}
fn value_at(&self, idx: usize) -> f64 {
self.values.borrow()[idx]
}
fn grad_at(&self, idx: usize) -> f64 {
self.grads.borrow()[idx]
}
}
#[derive(Copy, Clone)]
pub struct Var<'t> {
tape: &'t Tape,
pub(crate) idx: usize,
}
impl<'t> Var<'t> {
pub fn leaf(tape: &'t Tape, value: f64) -> Self {
let idx = tape.push_leaf(value);
Self { tape, idx }
}
pub fn value(&self) -> f64 {
self.tape.value_at(self.idx)
}
pub fn grad(&self) -> f64 {
self.tape.grad_at(self.idx)
}
pub fn sin(self) -> Var<'t> {
let v = self.value();
let out = self.tape.push_op(v.sin(), [self.idx, 0], 1, [v.cos(), 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn cos(self) -> Var<'t> {
let v = self.value();
let out = self
.tape
.push_op(v.cos(), [self.idx, 0], 1, [-v.sin(), 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn exp(self) -> Var<'t> {
let v = self.value();
let ev = v.exp();
let out = self.tape.push_op(ev, [self.idx, 0], 1, [ev, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn ln(self) -> Var<'t> {
let v = self.value();
let out = self.tape.push_op(v.ln(), [self.idx, 0], 1, [1.0 / v, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn sqrt(self) -> Var<'t> {
let v = self.value();
let sv = v.sqrt();
let grad = if sv.abs() < f64::EPSILON {
0.0
} else {
0.5 / sv
};
let out = self.tape.push_op(sv, [self.idx, 0], 1, [grad, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn tanh(self) -> Var<'t> {
let v = self.value();
let tv = v.tanh();
let grad = 1.0 - tv * tv;
let out = self.tape.push_op(tv, [self.idx, 0], 1, [grad, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn powi(self, n: i32) -> Var<'t> {
let v = self.value();
let grad = (n as f64) * v.powi(n - 1);
let out = self.tape.push_op(v.powi(n), [self.idx, 0], 1, [grad, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn powf(self, p: f64) -> Var<'t> {
let v = self.value();
let grad = p * v.powf(p - 1.0);
let out = self.tape.push_op(v.powf(p), [self.idx, 0], 1, [grad, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn abs(self) -> Var<'t> {
let v = self.value();
let grad = if v > 0.0 {
1.0
} else if v < 0.0 {
-1.0
} else {
0.0
};
let out = self.tape.push_op(v.abs(), [self.idx, 0], 1, [grad, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
pub fn recip(self) -> Var<'t> {
let v = self.value();
let grad = -1.0 / (v * v);
let out = self.tape.push_op(1.0 / v, [self.idx, 0], 1, [grad, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Add<Var<'t>> for Var<'t> {
type Output = Var<'t>;
fn add(self, rhs: Var<'t>) -> Var<'t> {
let val = self.value() + rhs.value();
let out = self.tape.push_op(val, [self.idx, rhs.idx], 2, [1.0, 1.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Add<f64> for Var<'t> {
type Output = Var<'t>;
fn add(self, rhs: f64) -> Var<'t> {
let val = self.value() + rhs;
let out = self.tape.push_op(val, [self.idx, 0], 1, [1.0, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Add<Var<'t>> for f64 {
type Output = Var<'t>;
fn add(self, rhs: Var<'t>) -> Var<'t> {
rhs + self
}
}
impl<'t> Sub<Var<'t>> for Var<'t> {
type Output = Var<'t>;
fn sub(self, rhs: Var<'t>) -> Var<'t> {
let val = self.value() - rhs.value();
let out = self.tape.push_op(val, [self.idx, rhs.idx], 2, [1.0, -1.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Sub<f64> for Var<'t> {
type Output = Var<'t>;
fn sub(self, rhs: f64) -> Var<'t> {
let val = self.value() - rhs;
let out = self.tape.push_op(val, [self.idx, 0], 1, [1.0, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Sub<Var<'t>> for f64 {
type Output = Var<'t>;
fn sub(self, rhs: Var<'t>) -> Var<'t> {
let val = self - rhs.value();
let out = rhs.tape.push_op(val, [rhs.idx, 0], 1, [-1.0, 0.0]);
Var {
tape: rhs.tape,
idx: out,
}
}
}
impl<'t> Mul<Var<'t>> for Var<'t> {
type Output = Var<'t>;
fn mul(self, rhs: Var<'t>) -> Var<'t> {
let (lv, rv) = (self.value(), rhs.value());
let out = self.tape.push_op(lv * rv, [self.idx, rhs.idx], 2, [rv, lv]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Mul<f64> for Var<'t> {
type Output = Var<'t>;
fn mul(self, rhs: f64) -> Var<'t> {
let val = self.value() * rhs;
let out = self.tape.push_op(val, [self.idx, 0], 1, [rhs, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Mul<Var<'t>> for f64 {
type Output = Var<'t>;
fn mul(self, rhs: Var<'t>) -> Var<'t> {
rhs * self
}
}
impl<'t> Div<Var<'t>> for Var<'t> {
type Output = Var<'t>;
fn div(self, rhs: Var<'t>) -> Var<'t> {
let (lv, rv) = (self.value(), rhs.value());
let g_rhs = -lv / (rv * rv);
let out = self
.tape
.push_op(lv / rv, [self.idx, rhs.idx], 2, [1.0 / rv, g_rhs]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Div<f64> for Var<'t> {
type Output = Var<'t>;
fn div(self, rhs: f64) -> Var<'t> {
let val = self.value() / rhs;
let out = self.tape.push_op(val, [self.idx, 0], 1, [1.0 / rhs, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
impl<'t> Div<Var<'t>> for f64 {
type Output = Var<'t>;
fn div(self, rhs: Var<'t>) -> Var<'t> {
let rv = rhs.value();
let val = self / rv;
let grad = -self / (rv * rv);
let out = rhs.tape.push_op(val, [rhs.idx, 0], 1, [grad, 0.0]);
Var {
tape: rhs.tape,
idx: out,
}
}
}
impl<'t> Neg for Var<'t> {
type Output = Var<'t>;
fn neg(self) -> Var<'t> {
let val = -self.value();
let out = self.tape.push_op(val, [self.idx, 0], 1, [-1.0, 0.0]);
Var {
tape: self.tape,
idx: out,
}
}
}
pub fn finite_diff_grad<F: Fn(f64) -> f64>(f: F, x: f64, h: f64) -> f64 {
(f(x + h) - f(x - h)) / (2.0 * h)
}
pub fn check_gradient<F>(f: F, x: f64, tol: f64) -> bool
where
F: for<'a> Fn(&'a Tape, Var<'a>) -> Var<'a>,
{
let tape = Tape::new();
let xv = Var::leaf(&tape, x);
let out = f(&tape, xv);
tape.backward(out);
let ad_grad = xv.grad();
let fd_grad = finite_diff_grad(
|xi| {
let t2 = Tape::new();
let xv2 = Var::leaf(&t2, xi);
f(&t2, xv2).value()
},
x,
1e-6,
);
(ad_grad - fd_grad).abs() < tol
}
#[cfg(test)]
mod tests {
use std::f64::consts::PI;
use super::*;
fn run_unary<F: Fn(Var<'_>) -> Var<'_>>(x_val: f64, f: F) -> (f64, f64) {
let tape = Tape::new();
let x = Var::leaf(&tape, x_val);
let out = f(x);
tape.backward(out);
(out.value(), x.grad())
}
#[test]
fn test_leaf_node() {
let tape = Tape::new();
let x = Var::leaf(&tape, 7.0);
assert!((x.value() - 7.0).abs() < 1e-15);
assert!((x.grad() - 0.0).abs() < 1e-15);
}
#[test]
fn test_add_grads() {
let tape = Tape::new();
let x = Var::leaf(&tape, 3.0);
let y = Var::leaf(&tape, 5.0);
let z = x + y;
tape.backward(z);
assert!((z.value() - 8.0).abs() < 1e-15);
assert!((x.grad() - 1.0).abs() < 1e-15);
assert!((y.grad() - 1.0).abs() < 1e-15);
}
#[test]
fn test_mul_grads() {
let tape = Tape::new();
let x = Var::leaf(&tape, 3.0);
let y = Var::leaf(&tape, 4.0);
let z = x * y;
tape.backward(z);
assert!((z.value() - 12.0).abs() < 1e-15);
assert!((x.grad() - 4.0).abs() < 1e-15); assert!((y.grad() - 3.0).abs() < 1e-15); }
#[test]
fn test_square_grad() {
let tape = Tape::new();
let x = Var::leaf(&tape, 5.0);
let z = x * x;
tape.backward(z);
assert!((z.value() - 25.0).abs() < 1e-15);
assert!((x.grad() - 10.0).abs() < 1e-14); }
#[test]
fn test_sin_grad() {
let x_val = PI / 4.0;
let (val, grad) = run_unary(x_val, |x| x.sin());
assert!((val - x_val.sin()).abs() < 1e-15);
assert!((grad - x_val.cos()).abs() < 1e-15);
}
#[test]
fn test_cos_grad() {
let x_val = PI / 3.0;
let (val, grad) = run_unary(x_val, |x| x.cos());
assert!((val - x_val.cos()).abs() < 1e-15);
assert!((grad - (-x_val.sin())).abs() < 1e-15);
}
#[test]
fn test_exp_grad() {
let x_val = 1.5;
let (val, grad) = run_unary(x_val, |x| x.exp());
assert!((val - x_val.exp()).abs() < 1e-14);
assert!((grad - x_val.exp()).abs() < 1e-14);
}
#[test]
fn test_ln_grad() {
let x_val = 2.0;
let (val, grad) = run_unary(x_val, |x| x.ln());
assert!((val - x_val.ln()).abs() < 1e-15);
assert!((grad - 1.0 / x_val).abs() < 1e-15);
}
#[test]
fn test_sqrt_grad() {
let x_val = 9.0;
let (val, grad) = run_unary(x_val, |x| x.sqrt());
assert!((val - 3.0).abs() < 1e-14);
assert!((grad - 1.0 / 6.0).abs() < 1e-14);
}
#[test]
fn test_tanh_grad() {
let x_val = 0.5;
let (val, grad) = run_unary(x_val, |x| x.tanh());
let tv = x_val.tanh();
assert!((val - tv).abs() < 1e-15);
assert!((grad - (1.0 - tv * tv)).abs() < 1e-15);
}
#[test]
fn test_chain_rule_sin_square() {
let x_val = 1.2;
let (val, grad) = run_unary(x_val, |x| (x * x).sin());
let expected_val = (x_val * x_val).sin();
let expected_grad = 2.0 * x_val * (x_val * x_val).cos();
assert!((val - expected_val).abs() < 1e-14);
assert!((grad - expected_grad).abs() < 1e-13);
}
#[test]
fn test_triple_product() {
let tape = Tape::new();
let x = Var::leaf(&tape, 2.0);
let y = Var::leaf(&tape, 3.0);
let w = Var::leaf(&tape, 4.0);
let z = (x * y) * w;
tape.backward(z);
assert!((z.value() - 24.0).abs() < 1e-14);
assert!((x.grad() - 12.0).abs() < 1e-13);
assert!((y.grad() - 8.0).abs() < 1e-13);
assert!((w.grad() - 6.0).abs() < 1e-13);
}
#[test]
fn test_powi_grad() {
let x_val = 2.0;
let (val, grad) = run_unary(x_val, |x| x.powi(3));
assert!((val - 8.0).abs() < 1e-14);
assert!((grad - 12.0).abs() < 1e-14); }
#[test]
fn test_gradient_check_complex() {
let x_val = 1.5;
let ok = check_gradient(
|_tape, x| {
let x_sq = x * x;
let s = x_sq.sin();
let e = (-x).exp();
let l = x.ln();
s * e + l
},
x_val,
1e-8,
);
assert!(ok, "AD gradient did not match finite difference");
}
#[test]
fn test_scalar_mul_add() {
let tape = Tape::new();
let x = Var::leaf(&tape, 3.0);
let z = x * 2.0 + 5.0; tape.backward(z);
assert!((z.value() - 11.0).abs() < 1e-14);
assert!((x.grad() - 2.0).abs() < 1e-14);
}
#[test]
fn test_sub_grads() {
let tape = Tape::new();
let x = Var::leaf(&tape, 7.0);
let y = Var::leaf(&tape, 3.0);
let z = x - y;
tape.backward(z);
assert!((x.grad() - 1.0).abs() < 1e-15);
assert!((y.grad() - (-1.0)).abs() < 1e-15);
}
#[test]
fn test_div_grads() {
let tape = Tape::new();
let x = Var::leaf(&tape, 6.0);
let y = Var::leaf(&tape, 3.0);
let z = x / y;
tape.backward(z);
assert!((z.value() - 2.0).abs() < 1e-15);
assert!((x.grad() - 1.0 / 3.0).abs() < 1e-14); assert!((y.grad() - (-6.0 / 9.0)).abs() < 1e-14); }
#[test]
fn test_neg_grad() {
let (_, grad) = run_unary(4.0, |x| -x);
assert!((grad - (-1.0)).abs() < 1e-15);
}
#[test]
fn test_tape_lengths() {
let tape = Tape::new();
let x = Var::leaf(&tape, 1.0);
let y = Var::leaf(&tape, 2.0);
assert_eq!(tape.values_len(), 2);
assert_eq!(tape.ops_len(), 0);
let _z = x + y;
assert_eq!(tape.values_len(), 3);
assert_eq!(tape.ops_len(), 1);
}
#[test]
fn test_tape_reset() {
let tape = Tape::new();
let x = Var::leaf(&tape, 1.0);
let y = Var::leaf(&tape, 2.0);
let _z = x + y;
tape.reset();
assert_eq!(tape.values_len(), 0);
assert_eq!(tape.ops_len(), 0);
}
}