reverse_mode_ad.ml

(* Reverse-mode algorithmic differentiation using effect handlers.
   Adapted from https://twitter.com/tiarkrompf/status/963314799521222656.
   See https://openreview.net/forum?id=SJxJtYkPG for more information. *)
   
module F = struct
  type t = { v : float; mutable d : float }

  let mk v = {v; d = 0.0}

  effect Plus : t * t -> t
  effect Mult : t * t -> t

  let grad f x =
    let x = mk x in
    begin match f x with
    | r -> r.d <- 1.0; r
    | effect (Plus(a,b)) k ->
        let x = {v = a.v +. b.v; d = 0.0} in
        ignore (continue k x);
        a.d <- a.d +. x.d;
        b.d <- b.d +. x.d;
        x
    | effect (Mult(a,b)) k ->
        let x = {v = a.v *. b.v; d = 0.0} in
        ignore (continue k x);
        a.d <- a.d +. (b.v *. x.d);
        b.d <- b.d +. (a.v *. x.d);
        x
    end;
    x.d

  let (+.) a b = perform (Plus(a,b))
  let ( *. ) a b = perform (Mult(a,b))
end;;

for x = 0 to 10 do
  let x = float_of_int x in
  assert (F.(grad (fun x -> x +. x *. x *. x) x) =
            1.0 +. 3.0 *. x *. x)
done