http://localhost/index/
index
/index/
重庆森林
This is the forest of Ilias Garnier. It acts as a core-dump of my brain: usefulness to third-parties would be purely incidental.
Those articles should be readable. Use ctrl-k to search the rest.
http://localhost/000A/
000A
/000A/
TODOs
This is mostly a backlink accumulator for unfinished writing.
Other stuff:
Finish cleaning up the vlayout library, and then the vplot one. Would be cool to make drawing diagrams with arrows easy. Look at circuit layout algos? Could I use MCTS/MCGS for this?
Cleanup beamforming code and write an article out of it
HMM, split old page on wind power prediction using particle filters
Add a term indexing page
Look at discrimination trees
Plenty of other stuff on my old hard drives: exhaustive molecule generation, schreier-sims and so on.
The selection monad x memoizable choice tree monad => greedy algorithms, dynamic programming
A category of weights and rewards
See Liam O'Connor's forester plugins on syntax highlighting via nova
Ah! I've been through too much:-But, sweet Satan, I beg of you, a less blazing eye! and while waiting for the new little cowardly gestures yet to come, since you like an absence of descriptive or didactic skills in a writer, let me rip out these few ghastly pages from my notebook of the damned.
References
Context
Backlinks
Related
http://localhost/000O/
000O
/000O/
重庆森林 › Articles
I'll post new notes and articles and backport old stuff here. This blog also got an Atom feed.
2026
1
12
http://localhost/001K/
001K
/001K/
Compiling with continuations, continued
article
Reading notes and toy implementation for Andrew Kennedy's paper.
The paper introduces a continuation-based intermediate representation (IR) which simplifies and makes efficiently implementable some
classical optimizations (shrinking and inlining) and lends itself well to lowering to e.g. SSA representations.The paper provides an abstract definition of the IR and its properties, defines a transformation from an ML-like language to the IR
and presents the optimizations as rewriting rules. Finally, a graph-based implementation is presented in a high-level fashion.There's a flurry of work on the respective merits of various IRs such as SSA (with or without named blocks), CPS, ANF etc. In this page, we'll focus on fleshing out an implementation of Kennedy's ideas.
#require "fmt";;
2026112Source languageOur source language is lambda-calculus with pairs and integers. The paper goes further (sums, let bindings, exceptions)
but we focus on the core idea.type expr =
| Var of string
| App of expr * expr
| Lam of string * expr
| Int of int
| Pair of expr * expr
| Fst of expr
| Snd of expr
;;Let's see how to encode typical expressions.(* Example: [fun x -> (x, x)] *)
let dup = Lam("x", Pair (Var "x", Var "x"));;val dup : expr = Lam ("x", Pair (Var "x", Var "x"))
The definition of pp_expr is omitted, but we can use it to print stuff:let () = Format.printf "%a@." pp_expr dup;;fun x -> (x, x)We'll use the following expression as a tiny running example.let running_example =
Lam ("f", Lam ("x", App (Var "f", Pair (Int 1, Var "x"))));;Or in readable form:let () = Format.printf "%a@." pp_expr running_example;;fun f -> fun x -> f (1, x)All good. Let's look at the definition of the IR!
2026112Target languageThe target language has explicitly named continuations. We make regular variables (type var) and continuation variables (kvar) distinct at the type level.type var = Var of string
type kvar = KVar of string
;;Like with the CPS transform, the IR makes control flow explicit. All intermediate computations are bound to names.
Note that there is a specific syntactic class for values, which includes in particular lambda-abstractions. Those take
an additional continuation as a parameter (i.e. they never return). The Kapp constructor corresponds to calling
a continuation, while Fapp (f, k, x) calls a function with argument x and continuation k.
(* Expressions *)
type kexpr =
(* let x = v in K *)
| Let_val of { var : var; bound : kvalue; body : kexpr }
(* let x = fst y in K *)
| Let_fst of { var : var; bound : var; body : kexpr }
(* let x = snd y in K *)
| Let_snd of { var : var; bound : var; body : kexpr }
(* let k x = K in L *)
| Let_cont of { kvar : kvar; arg : var; fbody : kexpr; body : kexpr }
(* k x *)
| Kapp of kvar * var
(* f k x *)
| Fapp of var * kvar * var
(* Values *)
and kvalue =
| Kv_int of int
| Kv_pair of var * var
(* \lambda k. \lambda x. K *)
| Kv_lam of kvar * var * kexpr
;;
2026112Naive translationTaking our running example and trying to convert it by hand is a nice exercise (try it without looking!).We proceed in a inside-out fashion. First the easy part: the innermost pair must be bound to a name. After the pair is evaluated, we proceed to the call of f.let f_pair =
Let_val {
var = Var "one";
bound = Kv_int 1;
body = Let_val {
var = Var "pair";
bound = Kv_pair (Var "one", Var "x");
body = Fapp (Var "f", KVar "kinner", Var "pair")
}
}
;;Or in perhaps easier to read form:let () = Format.printf "%a@." pp_kexpr f_pair;;let one = 1 in
let pair = (one, x) in
f kinner pairNote how we call a yet unbound function "f" with a yet unbound continuation "kinner". The variable "x" is free as well. Let's use pair to define the image of fun x -> f (1, x).let inner_lambda =
Kv_lam (KVar "kinner", Var "x", f_pair)
;;This lambda is a value and has to be bound to a name as well.let inner =
Let_val {
var = Var "inner_lambda";
bound = inner_lambda;
body = Kapp (KVar "kouter", Var "inner_lambda")
}
;;let () = Format.printf "%a@." pp_kexpr inner;;let inner_lambda = fun kinner x ->
let one = 1 in
let pair = (one, x) in
f kinner pair in
kouter inner_lambdaIt remains to treat the outermost function. The last continuation is the "identity": it marks the end of the transformation.let outer =
Let_val {
var = Var "outer_lambda";
bound = Kv_lam (KVar "kouter", Var "f", inner);
body = Kapp (KVar "identity", Var "outer_lambda")
}
;;let () = Format.printf "%a@." pp_kexpr outer;;let outer_lambda = fun kouter f ->
let inner_lambda = fun kinner x ->
let one = 1 in
let pair = (one, x) in
f kinner pair in
kouter inner_lambda in
identity outer_lambdaThis was rather cumbersome. We can automate this, but beforehand we need some tiny helpers.(* Fresh variable generation *)
let fresh f =
let c = ref 0 in
fun () ->
let v = !c in
incr c ;
Format.kasprintf f "%d" v
let kfresh = fresh (fun s -> KVar ("k" ^ s))
let vfresh = fresh (fun s -> Var ("x" ^ s))
;;val fresh : (string -> 'a) -> unit -> 'a = <fun>
val kfresh : unit -> kvar = <fun>
val vfresh : unit -> var = <fun>
(* Convenient infix syntax to "bind" transformation-time continuation. *)
let ( let@ ) f m = f m;;val ( let@ ) : ('a -> 'b) -> 'a -> 'b = <fun>
Here comes the "naive" transformation. Look at its type: it takes an expr and returns a kexpr.
But there is an extra "transformation-time" continuation of type var -> expr in the middle. It acts as an "accumulator" which encodes how to produce a kexpr when the input expr
has been transformed and the result of its evaluation bound a var.
let rec naive : expr -> (var -> kexpr) -> kexpr =
fun expr k ->
match expr with
| Var x -> k (Var x)
| App (f, arg) ->
let@ fv = naive f in
let@ argv = naive arg in
let kvar = kfresh () in
let arg = vfresh () in
Let_cont { kvar; arg; fbody = k arg; body = Fapp (fv, kvar, argv) }
| Lam (arg, body) ->
let var = vfresh () in
let kvar = kfresh () in
let body =
let@ xv = naive body in
Kapp (kvar, xv)
in
Let_val { var; bound = Kv_lam (kvar, Var arg, body); body = k var }
| Int i ->
let var = vfresh () in
Let_val { var; bound = Kv_int i; body = k var }
| Pair (l, r) ->
let@ lv = naive l in
let@ rv = naive r in
let var = vfresh () in
Let_val { var; bound = Kv_pair (lv, rv); body = k var }
| Fst p ->
let@ pv = naive p in
let var = vfresh () in
Let_fst { var; bound = pv; body = k var }
| Snd p ->
let@ pv = naive p in
let var = vfresh () in
Let_snd { var; bound = pv; body = k var }
let run_naive : expr -> kexpr =
fun expr -> naive expr (fun exprv -> Kapp (KVar "id", exprv))
;;val naive : expr -> (var -> kexpr) -> kexpr = <fun>
val run_naive : expr -> kexpr = <fun>
Let's apply this to our example!let () = Format.printf "%a@." pp_kexpr (run_naive running_example);;let x0 = fun k0 f ->
let x1 = fun k1 x ->
let x2 = 1 in
let x3 = (x2, x) in
let k2 x4 =
k1 x4
in
f k2 x3 in
k0 x1 in
id x0Up to -conversion, this is almost identical to our hand-rolled transform. Except that the naive transformation introduced a tiny -expanded continuation let k2 x4 = k1 x4!
We don't want that.
2026112http://localhost/kennedy-ccc-optimizing/kennedy-ccc-optimizing/kennedy-ccc-optimizing/Smarter translationThe optimizing translation, as its name suggests, is smarter about transforming expressions in tail position and avoids introducing those dummy -expanded forms.let rec optimizing : expr -> (var -> kexpr) -> kexpr =
fun expr k ->
match expr with
| Var x -> k (Var x)
| App (f, arg) ->
let@ fv = optimizing f in
let@ argv = optimizing arg in
let kvar = kfresh () in
let arg = vfresh () in
Let_cont { kvar; arg; fbody = k arg; body = Fapp (fv, kvar, argv) }
| Lam (arg, body) ->
(* [fun arg -> body]@k = letval v = fun ark k0 -> [body][fun bv -> k0 bv] in k v *)
let var = vfresh () in
let kvar = kfresh () in
let body = optimizing_tailcall body kvar in
Let_val { var; bound = Kv_lam (kvar, Var arg, body); body = k var }
| Int i ->
let var = vfresh () in
Let_val { var; bound = Kv_int i; body = k var }
| Pair (l, r) ->
let@ lv = optimizing l in
let@ rv = optimizing r in
let var = vfresh () in
Let_val { var; bound = Kv_pair (lv, rv); body = k var }
| Fst p ->
let@ pv = optimizing p in
let var = vfresh () in
Let_fst { var; bound = pv; body = k var }
| Snd p ->
let@ pv = optimizing p in
let var = vfresh () in
Let_snd { var; bound = pv; body = k var }
and optimizing_tailcall : expr -> kvar -> kexpr =
fun expr kv ->
match expr with
| Var x -> Kapp (kv, Var x)
| Int i ->
let var = vfresh () in
Let_val { var; bound = Kv_int i; body = Kapp (kv, var) }
| App (f, arg) ->
let@ fv = optimizing f in
let@ argv = optimizing arg in
Fapp (fv, kv, argv)
| Lam (arg, body) ->
let var = vfresh () in
let kvar = kfresh () in
Let_val
{ var;
bound = Kv_lam (kvar, Var arg, optimizing_tailcall body kvar);
body = Kapp (kv, var)
}
| Pair (l, r) ->
let@ lv = optimizing l in
let@ rv = optimizing r in
let var = vfresh () in
Let_val { var; bound = Kv_pair (lv, rv); body = Kapp (kv, var) }
| Fst p ->
let@ pv = optimizing p in
let var = vfresh () in
Let_fst { var; bound = pv; body = Kapp (kv, var) }
| Snd p ->
let@ pv = optimizing p in
let var = vfresh () in
Let_snd { var; bound = pv; body = Kapp (kv, var) }
let run_optimizing : expr -> kexpr =
fun expr -> optimizing expr (fun exprv -> Kapp (KVar "id", exprv))
;;val optimizing : expr -> (var -> kexpr) -> kexpr = <fun>
val optimizing_tailcall : expr -> kvar -> kexpr = <fun>
val run_optimizing : expr -> kexpr = <fun>
Let's apply this again to our example.let () = Format.printf "%a@." pp_kexpr (run_optimizing running_example);;let x5 = fun k3 f ->
let x6 = fun k4 x ->
let x7 = 1 in
let x8 = (x7, x) in
f k4 x8 in
k3 x6 in
id x5Much better!
2026112Optimizing by rewriting on a graph-like IRWe have seen how to efficiently translate to the IR but not how to perform optimizations on IR terms. The author argues that
the purely functional representation is ill-suited for rewriting purposes and suggests using a graph-based, mutable
term representation combined with a union-find data structure to represent binding constructs. Besides some useful
figures, an actual implementation is not shown. We remedy this.The first key ingredient is a representation for variable binders and variable occurrences. Prescriptions from the paper:
All occurrences for a variable should be linked in a doubly-linked (circular) list, for efficient iteration.
Those occurrences should also be organized in a union-find data structure, whose representative element should store a pointer to the binding construct
for that variable. This allows constant-time access to the binder from any variable.
The binder should point to the canonical occurrence, if any.
These might seem like weird prescriptions, but together they allow near-O(1) substitution of a variable for another one! An especially important operation in a IR
where all intermediate computations are named, and the basis of many optimizations. It also allows to know in constant time whether a binder is dead or not.
We start with the representation of binders and variable occurrences. Note that the type of binders and
occurrences is parameterized by 'expr which represents the type of expressions, to be defined later.
This is a way to break a mutually recursive definition.type 'expr cell = {
mutable repr : 'expr uf_node;
mutable next : 'expr cell;
mutable prev : 'expr cell;
}
and 'expr uf_node = Root of 'expr binder | Node of 'expr cell
and 'expr binder = {
var : string;
mutable binding : 'expr option;
mutable repr_cell : 'expr cell option;
};;Let's unpack this definition. We have three types, defined in a mutually recursive fashion.
A 'expr cell encode variable occurrences, organized in a doubly linked circular list for convenient iteration.
'expr uf_node are union-find nodes. All occurrences for a given binder are in the same class, the canonical representative being the binder.
A 'expr binder contain the name of the variable, a pointer to the binding expression (i.e. the actual AST node)
and a backpointer to the representative cell. This backpointer is initialized "late", via a wrapper
make_binder which ensures correctness.(* Initializes a [binder] named [var] and gives it to the callback [f] *)
let make_binder : type a. var:string -> (a binder -> a) -> a =
fun ~var f ->
let uninitialized = { var; binding = None; repr_cell = None } in
let expr = f uninitialized in
uninitialized.binding <- Some expr ;
expr
(* [repr cell] returns the binder to which the occurrence is attached. *)
let repr : type expr. expr cell -> expr binder =
fun cell ->
let rec loop cell acc =
match cell.repr with
| Node cell' -> loop cell' (cell :: acc)
| Root binder as repr ->
List.iter (fun c -> c.repr <- repr) acc ;
binder
in
loop cell []
;;val make_binder : var:string -> ('a binder -> 'a) -> 'a = <fun>
val repr : 'expr cell -> 'expr binder = <fun>
make_occurrence creates a fresh variable occurrence associated to a binder. If there is no occurrence,
the a singleton cell is created and set as being the union-find canonical representative. If not, the fresh
cell is inserted into the doubly linked list of occurrences.let insert_fresh_after cell binder =
(* assert that cell is in same binder eq. class *)
assert (repr cell == binder) ;
let repr_cell = Option.get binder.repr_cell in
let rec fresh = { repr = Node repr_cell; next = fresh; prev = fresh } in
let next = cell.next in
fresh.prev <- cell ;
fresh.next <- next ;
next.prev <- fresh ;
cell.next <- fresh ;
fresh
let make_occurrence : 'expr binder -> 'expr cell =
fun binder ->
match binder.repr_cell with
| None ->
let rec occurrence =
{ repr = Root binder; next = occurrence; prev = occurrence }
in
binder.repr_cell <- Some occurrence ;
occurrence
| Some repr_occ -> insert_fresh_after repr_occ binder
;;We can now define substitution subst x y, which substitutes y by x.
It proceeds as follows:
If y is dead (no occurrences), there's nothing to do.
Otherwise, we merge the lists of cells for x and y and let the representative cell for y
point to that for x.
Note that it is assumed in the implementation that x is not dead: this can be easily relaxed.
Note also that out of laziness, I didn't bother implementing the classical rank-based union computation - this also could be easily
added. As such, the implementation presented here is not near-constant time.
let unsafe_merge_dll (repr1 : _ cell) (repr2 : _ cell) =
let next1 = repr1.next in
let prev2 = repr2.prev in
repr1.next <- repr2 ;
repr2.prev <- repr1 ;
next1.prev <- prev2 ;
prev2.next <- next1
let subst (x : _ binder) (y : _ binder) =
assert (x != y) ;
begin
match (x.repr_cell, y.repr_cell) with
| (_, None) -> ()
| (None, _) -> assert false
| (Some x_repr_cell, Some y_repr_cell) ->
(* Invariant: the binder is what it should *)
assert (repr y_repr_cell == y) ;
assert (repr x_repr_cell == x) ;
unsafe_merge_dll x_repr_cell y_repr_cell ;
(* Reconnect representative cell to binding class of [x] *)
(* TODO: use rank-based UF merge *)
y_repr_cell.repr <- x_repr_cell.repr
end
;;val unsafe_merge_dll : 'a cell -> 'a cell -> unit = <fun>
val subst : 'a binder -> 'a binder -> unit = <fun>
That's it for the core variable handling! We define a pair of extra helpful functions.let get_name : 'expr binder -> string = fun binder -> binder.var
let is_dead binder = Option.is_none binder.repr_cell
let retrieve_binder : 'expr cell -> 'expr =
fun var ->
let binder = repr var in
Option.get binder.binding
;;val get_name : 'expr binder -> string = <fun>
val is_dead : 'a binder -> bool = <fun>
val retrieve_binder : 'expr cell -> 'expr = <fun>
We pack those definitions under a generative module whose signature is:module type Var = sig
type 'expr t
type 'expr binder
val make_binder : var:string -> ('a binder -> 'a) -> 'a
val make_occurrence : 'expr binder -> 'expr t
val get_name : 'a binder -> string
val subst : 'a binder -> 'a binder -> unit
val is_dead : 'a binder -> bool
val repr : 'expr t -> 'expr binder
val retrieve_binder : 'expr t -> 'expr
end
;;(We could have directly implemented that module but the REPL does not allow for piecewise module definitions).module Make_var () : Var =
struct
type 'expr t = 'expr cell
type nonrec 'expr binder = 'expr binder
let make_binder = make_binder
let make_occurrence = make_occurrence
let get_name = get_name
let subst = subst
let is_dead = is_dead
let repr = repr
let retrieve_binder = retrieve_binder
end;;module Make_var : () -> Var
Let's instantiate this module for regular and continuation variables.module Var = Make_var ()
module KVar = Make_var ()
;;We can now define the graph-like IR! The main difference, besides relying on the variable representation above, is that terms are mutable and that each expression contains a pointer to its enclosing expression
for easier rewriting. We have to let go the mutually recursive definition of kexpr and kvalue, whose type constraints will be encoded as dynamic checks.A side note: the reliance of parent pointers makes the sharing of identical sub-expressions impossible.
This is not much of an issue for the purpose of compiling and optimizing but it could be a show-stopper
in the setting of e.g. program synthesis.type kexpr_desc =
| Let_val of { var : var_binder; bound : kexpr; body : kexpr }
| Let_fst of { var : var_binder; bound : var; body : kexpr }
| Let_snd of { var : var_binder; bound : var; body : kexpr }
| Let_cont of { kvar : kvar_binder;
var : var_binder;
fbody : kexpr;
body : kexpr }
| Kapp of kvar * var
| Fapp of var * kvar * var
| Kv_int of int
| Kv_pair of var * var
| Kv_lam of kvar_binder * var_binder * kexpr
and kexpr = { mutable desc : kexpr_desc; mutable parent : kexpr option }
and var_binder = kexpr Var.binder
and kvar_binder = kexpr KVar.binder
and var = kexpr Var.t
and kvar = kexpr KVar.t
;;The presence of forward and backpointers makes constructing expressions not entirely trivial. We
make good use of the Var interface to provide a bunch of smart constructors taking care of
proper initialization.let make_var = Var.make_occurrence
let make_kvar = KVar.make_occurrence
let make desc = { desc; parent = None }
let is_value kexpr =
match kexpr.desc with Kv_int _ | Kv_pair _ | Kv_lam _ -> true | _ -> false;;val make_var : 'a Var.binder -> 'a Var.t = <fun>
val make_kvar : 'a KVar.binder -> 'a KVar.t = <fun>
val make : kexpr_desc -> kexpr = <fun>
val is_value : kexpr -> bool = <fun>
Let's unpack the definition of let_val. We see that the Let_val constructor takes
a binder as argument. We allocate this binder via Var.make_binder. In the scope
of the callback, where the binder is in scope, we evaluate the body of the construct.Something that we must take proper care of is the initialization of parent pointers: the convention is
that this initialization is left to use-time of the sub-expressions.let let_val ~(var : string) (bound : kexpr) (body : var_binder -> kexpr) =
(* Since we had to collapse the mutually recursive [expr] and [value]
types into one, we defer what was a statically checked invariant
to a dynamically checked one. *)
assert (is_value bound) ;
Var.make_binder ~var (fun var ->
let body = body var in
assert (Option.is_none bound.parent) ;
assert (Option.is_none body.parent) ;
let root = make (Let_val { var; bound; body }) in
let some_root = Some root in
bound.parent <- some_root ;
body.parent <- some_root ;
root)
;;val let_val : var:string -> kexpr -> (var_binder -> kexpr) -> kexpr = <fun>
The other constructs follow the same pattern.let let_fst ~(var : string) (bound : var) (body : var_binder -> kexpr) =
Var.make_binder ~var @@ fun var ->
let body = body var in
let root = make @@ Let_fst { var; bound; body } in
assert (Option.is_none body.parent) ;
let some_root = Some root in
body.parent <- some_root ;
root
let let_snd ~(var : string) (bound : var) (body : var_binder -> kexpr) =
Var.make_binder ~var @@ fun var ->
let body = body var in
let root = make @@ Let_snd { var; bound; body } in
assert (Option.is_none body.parent) ;
let some_root = Some root in
body.parent <- some_root ;
root
let let_cont ~(kvar : string) ~(var : string) (cont : var_binder -> kexpr)
(body : kvar_binder -> kexpr) =
KVar.make_binder ~var:kvar @@ fun kvar ->
Var.make_binder ~var @@ fun var ->
let fbody = cont var in
let body = body kvar in
let root = make @@ Let_cont { kvar; var; fbody; body } in
assert (Option.is_none fbody.parent) ;
assert (Option.is_none body.parent) ;
let some_root = Some root in
fbody.parent <- some_root ;
body.parent <- some_root ;
root
let kapp kv v = make (Kapp (kv, v))
let fapp fv k x = make (Fapp (fv, k, x))
let int i = make (Kv_int i)
let pair l r = make (Kv_pair (l, r))
let lam ~kvar ~var body =
KVar.make_binder ~var:kvar @@ fun kvar ->
Var.make_binder ~var @@ fun var -> make (Kv_lam (kvar, var, body kvar var))
;;val let_fst : var:string -> var -> (var_binder -> kexpr) -> kexpr = <fun>
val let_snd : var:string -> var -> (var_binder -> kexpr) -> kexpr = <fun>
val let_cont :
kvar:string ->
var:string -> (var_binder -> kexpr) -> (kvar_binder -> kexpr) -> kexpr =
<fun>
val kapp : kvar -> var -> kexpr = <fun>
val fapp : var -> kvar -> var -> kexpr = <fun>
val int : int -> kexpr = <fun>
val pair : var -> var -> kexpr = <fun>
val lam :
kvar:string -> var:string -> (kvar_binder -> var_binder -> kexpr) -> kexpr =
<fun>
Now that we can construct expressions, time to revisit conversion from the source language to the IR.
We use the optimizing translation. We just need to adjoin an environment mapping regular
variable names to binders and use the combinators defined above rather than straight constructors.module Env = Map.Make (String)
(* Scope representation: a map from names to binders. *)
type env = var_binder Env.t
let get_var (env : env) (v : string) : var_binder = Env.find v env
let kfresh = fresh (fun s -> s)
let vfresh = fresh (fun s -> s)
;;let rec optimizing : env -> expr -> (var_binder -> kexpr) -> kexpr =
fun env expr k ->
match expr with
| Var x -> k (get_var env x)
| App (f, arg) ->
let@ fv = optimizing env f in
let@ argv = optimizing env arg in
let kvar = kfresh () in
let var = vfresh () in
let_cont ~kvar ~var k (fun kv ->
fapp (make_var fv) (make_kvar kv) (make_var argv))
| Lam (arg, body) ->
(* [fun arg -> body]@k = letval v = fun ark k0 -> [body][fun bv -> k0 bv] in k v *)
let var = vfresh () in
let kvar = kfresh () in
let_val
~var
(lam ~kvar ~var:arg (fun kv v ->
let env = Env.add arg v env in
optimizing_tailcall env body kv))
k
| Int i ->
let var = vfresh () in
let_val ~var (int i) k
| Pair (l, r) ->
let@ lv = optimizing env l in
let@ rv = optimizing env r in
let var = vfresh () in
let_val ~var (pair (make_var lv) (make_var rv)) k
| Fst p ->
let@ pv = optimizing env p in
let var = vfresh () in
let_fst ~var (make_var pv) k
| Snd p ->
let@ pv = optimizing env p in
let var = vfresh () in
let_snd ~var (make_var pv) k
and optimizing_tailcall : env -> expr -> kvar_binder -> kexpr =
fun env expr kv ->
match expr with
| Var x -> kapp (make_kvar kv) (make_var (get_var env x))
| Int i ->
let var = vfresh () in
let_val ~var (int i) (fun v -> kapp (make_kvar kv) (make_var v))
| App (f, arg) ->
let@ fv = optimizing env f in
let@ argv = optimizing env arg in
fapp (make_var fv) (make_kvar kv) (make_var argv)
| Lam (arg, body) ->
let var = vfresh () in
let kvar = kfresh () in
let_val
~var
(lam ~kvar ~var:arg (fun kv v ->
let env = Env.add arg v env in
optimizing_tailcall env body kv))
(fun v -> kapp (make_kvar kv) (make_var v))
| Pair (l, r) ->
let@ lv = optimizing env l in
let@ rv = optimizing env r in
let var = vfresh () in
let_val
~var
(pair (make_var lv) (make_var rv))
(fun v -> kapp (make_kvar kv) (make_var v))
| Fst p ->
let@ pv = optimizing env p in
let var = vfresh () in
let_fst ~var (make_var pv) (fun v -> kapp (make_kvar kv) (make_var v))
| Snd p ->
let@ pv = optimizing env p in
let var = vfresh () in
let_snd ~var (make_var pv) (fun v -> kapp (make_kvar kv) (make_var v))
let run_optimizing : expr -> kexpr =
fun expr ->
KVar.make_binder ~var:"id" @@ fun id ->
let env = Env.empty in
optimizing env expr (fun exprv -> kapp (make_kvar id) (make_var exprv))
;;We conclude this page by implementing an optimization example presented in the article.
This example is a "shrinking reduction", corresponding to the elimination rules for pairs.
In essence, if we project out a variable which is statically bound to a pair, we may replace
the projection by the corresponding pair component. The rewriting rule is the following, where
C ranges over contexts:
letval x = (x1, x2) in --> letval x = (x1, x2) in
C{let y = fst x in K} C{K[x1/y]}
Let's implement this. We define a type of shrinking reductions (we restrict to the example above).
A value Beta_proj_fst { letval_pair; letproj } corresponds to the left hand side of the rule
above, where letval_pair is the full expression and letproj is the subexpression being
eliminated.
type shrinking_reduction =
| Beta_proj_fst of { letval_pair : kexpr; letproj : kexpr }
;;Performing the reduction boils down to performing the substitution. This makes the
let y = fst x in ... binding dead, allowing us to erase it.let proj_fst pair =
match pair.desc with Kv_pair (fst, _) -> fst | _ -> assert false
let perform_reduction : shrinking_reduction -> unit = function
| Beta_proj_fst { letval_pair; letproj } ->
begin match (letval_pair.desc, letproj.desc) with
| (Let_val { var = x; bound = pair; body = _ },
Let_fst { var; bound; body = k }) ->
assert (Var.repr bound == x) ;
let x1_binder = Var.repr (proj_fst pair) in
(* Perform substitution *)
Var.subst x1_binder var ;
(* Erase now useless binding *)
(* invariant: k record is uniquely owned and is going to become unreachable *)
k.parent <- None ;
letproj.desc <- k.desc
| _ -> assert false
end
;;val proj_fst : kexpr -> var = <fun>
val perform_reduction : shrinking_reduction -> unit = <fun>
That's it for the optimization itself! It remains to find such redexes in the program. We
use a generic iterator over the AST, to which we pass a function which scans for redexes and accumulates
them in a queue.let rec iter kexpr ~f =
f kexpr ;
begin
match kexpr.desc with
| Let_val { var = _; bound; body } ->
iter bound ~f ;
iter body ~f
| Let_fst { var = _; bound = _; body } -> iter body ~f
| Let_snd { var = _; bound = _; body } -> iter body ~f
| Let_cont { kvar = _; var = _; fbody; body } ->
iter fbody ~f ;
iter body ~f
| Kapp _ -> ()
| Fapp _ -> ()
| Kv_int _ -> ()
| Kv_pair _ -> ()
| Kv_lam (_, _, body) -> iter body ~f
end
let is_kv_pair kexpr = match kexpr.desc with Kv_pair _ -> true | _ -> false
let scan_reduction kexpr queue =
iter kexpr ~f:(fun expr ->
match expr.desc with
| Let_fst { var = _; bound; body = _ } ->
(* [bound] is a variable, find its binder. *)
let binder_expr = Var.retrieve_binder bound in
begin match binder_expr.desc with
| Let_val { var = binder_var; bound = bound'; body = _ }
when is_kv_pair bound' ->
assert (Var.repr bound == binder_var) ;
Queue.add
(Beta_proj_fst { letval_pair = binder_expr;
letproj = expr })
queue
| _ -> ()
end
| _ -> ())
;;val iter : kexpr -> f:(kexpr -> 'a) -> unit = <fun>
val is_kv_pair : kexpr -> bool = <fun>
val scan_reduction : kexpr -> shrinking_reduction Queue.t -> unit = <fun>
Let's apply this to some example!let example =
lam ~kvar:"k" ~var:"x" @@ fun k x ->
let_val ~var:"p" (pair (make_var x) (make_var x)) @@ fun p ->
let_fst ~var:"z" (make_var p) @@ fun z ->
let_val ~var:"w" (pair (make_var p) (make_var z)) @@ fun w ->
let_val ~var:"t" (pair (make_var w) (make_var x)) @@ fun t ->
kapp (make_kvar k) (make_var t)
let () = Format.printf "%a@." pp_kexpr example;;fun k x ->
let p = (x, x) in
let z = fst p in
let w = (p, z) in
let t = (w, x) in
k tScanning for the redex bound to z:let queue = Queue.create ()
let () = scan_reduction example queue
let () = Format.printf "queue length: %d@." (Queue.length queue);;val queue : shrinking_reduction Queue.t = <abstr>
queue length: 1Performing the optimization:Queue.iter perform_reduction queue;;- : unit = ()
Printing the result:let () = Format.printf "%a@." pp_kexpr example;;fun k x ->
let p = (x, x) in
let w = (p, x) in
let t = (w, x) in
k t
2026112ConclusionThat was a lot of work, and we only considered a portion of the ideas covered in the paper.
Some things we skipped:
Algebraic datatypes (we only covered products)
Handling exceptions via double-barrelled CPS
Recursion
Other shrinking reductions
Contification
However I think we covered the heart of the paper. The items above are nontrivial but not too hard once
the core structure is pinned down. Very nice paper!
2025
12
14
http://localhost/0019/
0019
/0019/
Backtracking parser combinators
article
Reasonable languages are LR(1). Parsing them is best done using tools like menhir. Out of this convenient class of grammars, one has to resort to other means: manually writing the parser, or using parser combinators.
20251215http://localhost/001A/001A/001A/Parsing as state monads on streams
We'll assume fixed a type of of symbols.
Parsing means consuming part of a stream of symbols and producing a value of a given type, together with the rest of the stream.
A stream can be seen as a function .
The "parsing functor" is
We recognize a thinly veiled state monad. We may also need to handle failure: to this end, compose the state monad with some error monad.
(In this article we'll take .) Let's implement this, see how it falls short of our goals and the how codensity monad of saves the day.
20251214State-monad based parsing combinators#require "fmt";;For simplicity, we will take symbols to be characters and streams to be sequences of characters.type symbol = char
type stream = char Seq.t
(* Note the call to [memoize]: in order to implement backtracking, we want our streams
to be persistent. We must be able to reprocess a chunk of the stream as many times
as it takes. *)
let of_string : string -> stream = fun str -> String.to_seq str |> Seq.memoize;;
Our basic implementation (hinting that there will be a non-basic one) will be a state monad
over streams composed with an option monad for errors. We start of with basic monadic primitives:
bind, return and run.
(* ['a t] is the type of parsers yielding a value of type ['a] *)
type 'a t = stream -> 'a option * stream
(* Monadic composition of parsers.
Note the straightforward handling of errors. *)
let bind : 'a t -> ('a -> 'b t) -> 'b t =
fun m f stream ->
match m stream with
| (None, _) as fail -> fail
| (Some x, rest) -> f x rest
let return x = fun stream -> (Some x, stream)
(* Running a monadic computation. *)
let run (p : 'a t) (str : string) = fst @@ p (of_string str)
;;We add some basic way of expressing failure. let fail () : 'a t = fun stream -> (None, stream);;Now, we can define some basic parsing combinators!(* A convenient syntax for [bind]. Recall that [let* x = m in b]
is equivalent to [bind m (fun x -> b)]. *)
let (let*) = bind
(* [exact sym] matches exactly the symbol [sym], or fails. *)
let exact sym : unit t =
fun stream ->
match Seq.uncons stream with
| None -> (None, stream)
| Some (x, rest) ->
if Char.equal x sym then
(Some (), rest)
else
(None, rest)
(* [try_bind m f on_error] is used as a building block for backtracking.
When [m] evaluates to a failure, [f] is called with the stream as it was
before computing [m]. *)
let try_bind : 'a t -> ('a -> 'b t) -> (unit -> 'b t) -> 'b t =
fun m f on_error stream ->
match m stream with
| (None, _) -> on_error () stream
| (Some x, rest) -> f x rest
(* [first_of [p_1; ...; pk]] tries to parse the stream with each parser in the given list,
in order. *)
let rec first_of : 'a t list -> 'a t =
fun cases ->
match cases with
| [] -> fail ()
| p :: rest ->
try_bind p return (fun () -> first_of rest)
;;val ( let* ) : 'a t -> ('a -> 'b t) -> 'b t = <fun>
val exact : Char.t -> unit t = <fun>
val try_bind : 'a t -> ('a -> 'b t) -> (unit -> 'b t) -> 'b t = <fun>
val first_of : 'a t list -> 'a t = <fun>
20251214Some simple examples and a failurelet print = function
| None -> print_endline "failure"
| Some _ -> print_endline "success!"
(* Matches "ab" *)
let match_ab =
let* () = exact 'a' in
exact 'b'
let () = print (run match_ab "ab");;success!let () = print (run match_ab "ac");;failureNow let's try a more complicated example, involving backtracking.let match_abxx =
let* () = match_ab in
let* () = exact 'x' in
exact 'x'
let match_ab_or_abxx =
first_of [match_ab; match_abxx]
let parser =
let* () = match_ab_or_abxx in
exact '.'
let () = print (run parser "abxx.");;failureWhat gives?
the parser match_ab successfully matches the prefix of the stream
the next expected symbol is '.', however the rest of the stream contains "xx."
but since match_ab succeeded, it can't act as a backtracking point anymore!
The root of the issue is that we'd like try_bind m f on_error to act as a backtracking point for its whole continuation, rather
than just the computation corresponding to m. Using the codensity monad of a well-chosen functor allows to express
this pattern.
2025716http://localhost/000J/000J/000J/Codensity monadDefinitionThe codensity monad of a functor is formally defined as its right Kan extension along itself, when it exists. In programming terms, assuming a (category-theoretic) functor is represented as a type constructor , its codensity monad is . In code:
module Codensity (F : sig type 'a t val return : 'a -> 'a t end) =
struct
type 'a t = { call : 'b. ('a -> 'b F.t) -> 'b F.t }
let bind : type a b. a t -> (a -> b t) -> b t =
fun m f ->
{ call = (fun k -> m.call (fun xm -> (f xm).call k)) }
let return : type a. a -> a t =
fun x -> { call = (fun k -> k x) }
let run : 'a t -> 'a F.t = fun m -> m.call F.return
end
A particularly useful trick is to use a functor of the form . Typically, would be some kind of state or monoid. Since we're programming in CPS, the state can always stay in 'parameter' position!Some good resources about the codensity monad for programming are:
Kan Extensions for Program Optimisation Or: Art and Dan Explain an Old Trick (Ralf Hinze)
Asymptotic Improvement of Computations over Free Monads (Janis Voigtlander)
20251214Codensity monad-based parsing combinatorsTo which functor should we apply the codensity construction? In order to properly implement backtracking,
backtracking points must be able to modify the continuation of their context when a failure occurs.
Failure will hence be handled by passing an extra continuation (an "error continuation state") along the stream
(here in first position):
Yielding:
It is worth emphasizing again that the error continuation should be seen as part of the state, same as the stream.
This type might be reminiscent of a double-barrelled
CPS transformation but the monad above is more expressive as the 'main' continuation can modify its own error continuation.
type error = unit
module M = Codensity(struct
type 'a t = (stream -> 'a option) -> stream -> 'a option
let return x = fun _kerr _stream -> Some x
end);;type error = unit
module M :
sig
type 'a t = {
call :
'b.
('a -> (stream -> 'b option) -> stream -> 'b option) ->
(stream -> 'b option) -> stream -> 'b option;
}
val bind : 'a t -> ('a -> 'b t) -> 'b t
val return : 'a -> 'a t
val run : 'a t -> (stream -> 'a option) -> stream -> 'a option
end
Let's reimplement our parsing primitive. Bear extra attention to try_bind!open M
let ( let* ) = bind
let try_bind : type a b. a t -> (a -> b t) -> (error -> b t) -> b t =
fun m f ferr ->
{ call =
(fun k kerr stream ->
(* We override [kerr] with [on_error].
Observe how [on_error] resumes with the whole [stream]. *)
let on_error = fun _stream -> (ferr ()).call k kerr stream in
m.call (fun x kerr stream ->
(f x).call k kerr stream
) on_error stream)
}
let fail () = { call = (fun _k kerr stream -> kerr stream) }
let exact tok : _ t =
{ call =
(fun k kerr stream ->
match Seq.uncons stream with
| None -> kerr stream
| Some (tok', rest) -> (
match Char.equal tok tok' with
| false -> kerr stream
| true -> k () kerr rest))
}
let first_of (parsers : _ t list) : _ t =
let rec loop parsers =
match parsers with
| [] -> fail ()
| p :: rest -> try_bind p return (fun _err -> loop rest)
in
loop parsers
;;let print = function
| None -> print_endline "failure"
| Some _ -> print_endline "success!"
let run p (str : string) = run p (fun _ -> None) (String.to_seq str)
(* Matches "ab" *)
let match_ab =
let* () = exact 'a' in
exact 'b'
let () = print (run match_ab "ab");;success!let () = print (run match_ab "ac");;failureLet's retry the example involving backtracking.let match_abxx =
let* () = match_ab in
let* () = exact 'x' in
exact 'x'
let match_ab_or_abxx =
first_of [match_ab; match_abxx]
let parser =
let* () = match_ab_or_abxx in
exact '.'
let () = print (run parser "abxx.");;success!
Some remarks/ideas
The idea above is fairly general and doesn't apply just to parsers. try_bind is a limited form of effect handler (our continuations are not delimited though).
To make this reasonably performant, one may need to defunctionalize the implementation, or perhaps use the flambda compiler with unboxing
annotations? In any case, one would need a better representation of streams.
In the basic example given here, the error continuation never uses its input stream meaningfully but it would be useful e.g. to continue parsing despite errors.
The error type could be made more interesting that unit. In that case, the error continuation could take an extra argument corresponding to the error.
Rather than first_of, we could handle non-determinism by enumerating all valid parses.
Or we may want to optimize for some given cost function. In this setting, can we contrast the codensity-based approach vs selection monads?
2025
12
5
http://localhost/0018/
0018
/0018/
Functional machine calculus
article
Links:
The functional machine calculus I, II, III
The relational machine calculus (Barrett et al.)
The functional machine calculus (FMC): a stack-based model of computation that allows to faithfully represent effectful and higher-order computation.
FMC may serve as a compilation target for effectful -calculi;
FMC are closely related to their denotations as string diagrams in appropriate categories.
2025125TODO: FMC semanticsFMC syntax features variables, sequencing, pushing onto a named stack and popping
off a named stack, binding the result to a variable. Stack names are dubbed "locations",
variables ranging over locations are denoted . The term "location" hints at the use of the stacks to implement effectful constructs such as mutable references.
Capture-avoiding substitution on expressions is denoted .
We postulate a "main" location: when location variables are elided, push and pop refer to that one. Extra locations may be used to encode other state-like effects.
Operational semantics is given by transforming stacks (we only use the main location for simplicity). A stack is just a sequence of expressions, with the
top of the stack written on the right. The one-step evaluation relation is denoted . Stacks will be denoted .
Popping triggers a capture-avoiding substitution.
We elude sequencing and no-op rules.
Question: can we implement a VM for the FMC using an environment rather than substitutions?We extend the FMC with string literals and concatenation to have non-trivial computations.Let's code this.
#require "ppx_deriving.show";;
#require "fmt";;
We directly represent sequential composition as list of instructions.
module Syntax =
struct
type literal = string
[@@deriving show { with_path = false }]
(* The type of instructions *)
type instr =
| Var of string (* Variable *)
| Push of t (* Push code onto stack *)
| Pop of string (* Pop from stack, bind to var *)
| Lit of literal (* Push literal to stack *)
| Concat (* String concatenation *)
[@@deriving show { with_path = false }]
(* An FMC program is a list of instructions *)
and t = instr list
[@@deriving show { with_path = false }]
end;;
The target machine uses De Bruijn indices instead of named variables.
module Code =
struct
type literal = Syntax.literal
[@@deriving show { with_path = false }]
(* Continuation is made explicit as last arg of each instruction. *)
type t =
| Var of int * t (* De Bruijn index and continuation *)
| Push of t * t
| Pop of t
| Lit of literal * t
| Concat of t
| Nil
[@@deriving show { with_path = false }]
type scope = string list
[@@deriving show { with_path = false }]
let rec from_instr : Syntax.instr -> scope -> (scope -> t) -> t =
fun instr scope k ->
match instr with
| Var s ->
begin match List.find_index (fun s' -> String.equal s s') scope with
| None -> failwith ("unbound variable " ^ s)
| Some i -> Var (i, k scope)
end
| Push t ->
Push (from_syn t scope, k scope)
| Pop name ->
Pop (k (name :: scope))
| Lit l ->
Lit (l, k scope)
| Concat ->
Concat (k scope)
and from_syn : Syntax.t -> scope -> t =
fun syn scope ->
match syn with
| [] -> Nil
| instr :: rest ->
from_instr instr scope (fun scope -> from_syn rest scope)
end;;
Evaluation implements subtitution using an explicit environment. The resulting machine is reminiscent of
the SECD machine.
module Eval =
struct
(* The type of environments. *)
type env = value list [@@deriving show]
(* A closure is a piece of code in a given environment *)
and closure = env * Code.t [@@deriving show]
(* A value is either a closure or a literal *)
and value =
| Vclos of closure
| Vlit of Code.literal
[@@deriving show]
(* The [stack] type models the main stack of the FMC. *)
type stack = value Dynarray.t
let empty_stack () = Dynarray.create ()
let push : stack -> value -> unit = Dynarray.add_last
let pop : stack -> value = Dynarray.pop_last
let pp_stack = Fmt.using Dynarray.to_seq (Fmt.Dump.seq pp_value)
(* The [kstack] type corresponds to a stack of continuations. *)
type kstack = closure list
(* Evaluation: execute each instruction, updating the stack and/or
environment, and move to the next continuation.
Main correctness question: correctly handling the environment. *)
let rec eval : Code.t -> env -> stack -> kstack -> stack =
fun code env stack kstack ->
match code with
| Code.Var (i, k) ->
let v = List.nth env i in
begin
match v with
| Vlit lit ->
push stack v ;
eval k env stack kstack
| Vclos (env', code') ->
(* Push current cont. to kstack and install closure. *)
eval code' env' stack ((env, k) :: kstack)
end
| Nil -> begin
match kstack with
| [] ->
(* Execution terminates *)
stack
| (env', code) :: rest -> eval code env' stack rest
end
| Push (Var (i, Nil), k) ->
(* This clause is not strictly required but I conjecture that it is
sufficient to ensure that literals are canonically represented
by values of the form [Vlit] on the stack, rather than trivial
closures over a variable. *)
let v = List.nth env i in
push stack v ;
eval k env stack kstack
| Push (code, k) ->
push stack (Vclos (env, code)) ;
eval k env stack kstack
| Pop k ->
let v = pop stack in
eval k (v :: env) stack kstack
| Lit (lit, k) ->
push stack (Vlit lit) ;
eval k env stack kstack
| Concat k ->
let v2 = pop stack in
let v1 = pop stack in
begin
match (v1, v2) with
| (Vlit s1, Vlit s2) ->
(* Conjecture: in a well-typed program, literals on the stack
are of the form [Vlit]. *)
push stack (Vlit (String.cat s1 s2)) ;
eval k env stack kstack
| _ -> failwith "ill-typed program"
end
end;;
Let's play around with some examples. In The functional machine calculus II, a translation from the call-by-value -calculus to the FMC is given.
In code:
type lc =
| Var of string
| Lam of string * lc
| App of lc * lc
| Lit of Syntax.literal
| Concat of lc * lc
(* fresh variable generator *)
let fresh : unit -> string =
let c = ref 0 in
fun () ->
incr c;
"v" ^ (string_of_int !c)
(* Convenience constructors *)
let lambda (f : lc -> lc) = let v = fresh () in Lam (v, f (Var v))
let app (f : lc) (arg : lc) = App (f, arg)
let lit s : lc = Lit s
let concat l r = Concat (l, r)
(* Conversion from lambda-calculus to FMC *)
let rec to_fmc_aux : lc -> Syntax.t -> Syntax.t =
fun term k ->
let open Syntax in
match term with
| Var s -> Push [Var s] :: k
| Lam (s, body) -> Push (Pop s :: to_fmc_aux body []) :: k
| App (f, arg) ->
let var = fresh () in
to_fmc_aux arg (to_fmc_aux f (Pop var :: Var var :: k))
| Lit l -> Lit l :: k
| Concat (lhs, rhs) ->
to_fmc_aux rhs (to_fmc_aux lhs (Concat :: k))
let to_fmc term = to_fmc_aux term [] ;;
Let's print some basic examples.
let print_fmc fmc = print_endline (Syntax.show fmc);;
let print_lambda lc = print_fmc (to_fmc lc);;
print_lambda @@ lambda (fun x -> x);;[(Push [(Pop "v1"); (Push [(Var "v1")])])]
print_lambda @@ lambda (fun fst -> lambda (fun snd -> snd));;[(Push [(Pop "v2"); (Push [(Pop "v3"); (Push [(Var "v3")])])])]
print_lambda @@ app (app (lambda (fun fst -> lambda (fun snd -> snd))) (lit "fst")) (lit "snd");;[(Lit "snd"); (Lit "fst");
(Push [(Pop "v4"); (Push [(Pop "v5"); (Push [(Var "v5")])])]); (Pop "v7");
(Var "v7"); (Pop "v6"); (Var "v6")]
Now let's evaluate stuff.
(* Evaluate a lambda term and return the FMC stack *)
let eval_lambda lc =
let stack = Eval.empty_stack () in
let code = Code.from_syn (to_fmc lc) [] in
Eval.eval code [] stack []
;;val eval_lambda : lc -> Eval.stack = <fun>
We expect this to evalaute to the string "snd".
let stack = eval_lambda @@ app (app (lambda (fun fst -> lambda (fun snd -> snd))) (lit "fst")) (lit "snd");;
print_endline (Fmt.to_to_string Eval.pp_stack stack);;[(Eval.Vlit "snd")]
Now a lambda which concatenates a string with itself, featuring non-linearity.
let stack = eval_lambda @@ (app (lambda (fun s -> concat s s)) (lit "str_"));;
print_endline (Fmt.to_to_string Eval.pp_stack stack);;[(Eval.Vlit "str_str_")]
Remarks/going further:
We could use GADTs to encode the stack typing discipline described in the paper, in the style of what the Michelson interpreter does
State-like effects can be handled with extra stacks (one stack per 'location'/reference). If my understanding in correct this means that the number of stacks may dynamically grow (e.g. when one allocates new references.)
Control: the extension is described in the third paper and seems straightforward, though the description makes it look like rather inefficent (basically, a switch would be executing pieces of code conditionally for each branch until finding the right one rather than actually jumping).
2025
11
8
http://localhost/000P/
000P
/000P/
A plugin system for Forester
article
Forester is a great tool but I have found writing code snippets frustrating, as those are not evaluated nor type-checked. I just prototyped a plugin feature which allows essentially mapping Forester values through an external program. I got a plugin connected to an OCaml toplevel instance. At the time of writing the code is on my github. I have the following needs:
Intersplice text and pieces of code to be evaluated, following the usual literate programming practice.
Maintain simultaneously distinct states for a given plugin.
Ideally, integrate extra artifacts (images, etc) produced by the plugin. This is not done.
I tried to be minimally invasive but my patch certainly requires a lot of discussion before being mergeable. I plan to dogfood it on my own forest for the time being. A demonstration is in order.
let x = 1+1;;val x : int = 2
We note that the code is not highlighted. This could be a future feature. Let's carry on with our computation.
x + 3;;- : int = 5
The Forester code looks like this:
Plugin initialization happens at The resulting (stateful) plugin value is attached to the key
In this case, initialization is nontrivial: the toplevel needs to be compiled in bytecode, so it leaves in its own server, in a separate executable from Forester. When the plugin is initialized, the server is launched and the Forester-side of the plugin acts as a client, pushing phrases and getting the outcome as well as the catpured stdout. Each plugin instance gets its own server, hence its own toplevel. The OCaml toplevel is very stateful so having one process per plugin instance takes a lot of problems away.
Each call to the plugin takes two arguments. The first one is used to pass plugin options, as text content. The second one takes the code to execute.
By default, stdout and topevel output are captured:
Format.printf "stdout message"; 42 ;;- : int = 42
stdout message
Sometimes however we don't want one or the other. One can pass a comma-separated list of options
as first argument: "no-stdout" disables standard output capture, "no-echo" disables the
toplevel echo.
Format.printf "stdout message"; 42 ;;- : int = 42
Some items that needs improving/future features:
general: the command forester new re-evaluates everything, do we really need to run the plugins or can't we just run in 'dummy' mode? Or cache result of evaluation as artifacts, perhaps? The latter requires tracking depenencies across calls to a given plugin instance, which is nontrivial.
general: come up with a plugin protocol that would allow to exchange artifacts produced by the plugin (encompassing what the latex feature does)
general: allow to pass optional parameters at plugin initialization and use time
ocaml: use native toplevel instead of bytecode one
ocaml: syntax highlighting
ocaml: handle newlines and other special characters in string literals
I discovered this RFC after implementing the plugin. Worth reading in depth.
2025
7
13
http://localhost/000D/
000D
/000D/
Monadic Monte Carlo tree search
article
Monte Carlo tree search (MCTS for short) is a heuristic method to solve markov decision process (MDPs).
2025
7
13
http://localhost/000F/
000F
/000F/
Markov decision processes
Definition
Markov decision processes (MDPs) are models of an agent that moves probabilistically from state to state by taking actions, and gets rewards at each state it visits. MDPs are the canonical setting of reinforcement learning (RL). Solving an MDP means finding a policy, i.e. a function which at every state gives an optimal action to take so that the expected reward is maximized.
A monadic interface to describe (non-probabilistic) MDPs is the following:
'a t
val bind : 'a t -> ('a -> 'b t) -> 'b t
val choices : 'a array -> 'a t
val reward : float -> unit t
end
\end{lstlisting}
]]>
return and bind are the usual monadic constructs, while choices is the set of actions available at that program point and reward attributes a (nonnegative) reward to the trace in which it occurs.
MCTS implements this interface by representing computations as lazily expanded choice trees where local exploration decisions are guided by multi-armed bandits. We will explain each of these techniques before coalescing them into an implementation for MCTS. Before, a toy motivating example is in order.
2025
7
14
http://localhost/000G/
000G
/000G/
Monadic MCTS
Example
Assume given an MCTS implementation of the MDP interface: module Mcts : MDP
Consider the following type for arithmetic expressions featuring holes.
We implement a function that non-deterministically evaluates such expressions by filling holes. eval takes an expression and returns a computation in the type Mcts.t which yields the result of the evaluation and the generated expression.
(int * expr) Mcts.t =
fun expr ->
match expr with
| Add (e1, e2) ->
let* (v1, e1') = eval e1 in
let* (v2, e2') = eval e2 in
return (v1 + v2, Add (e1', e2'))
| Mul (e1, e2) ->
let* (v1, e1') = eval e1 in
let* (v2, e2') = eval e2 in
return (v1 * v2, Mul (e1', e2'))
| Const i -> return (i, expr)
| Hole ->
let* choice =
Mcts.choices [| Add (Hole, Hole);
Mul (Hole, Hole);
Const 1;
Const 2 |]
in
eval choice
\end{lstlisting}
]]>
The function eval describes the simultaneous evaluation and generation of an arithmetic expression. Observe how a hole is filled: we call Mcts.choices to describe possible options and proceed with the evaluation on the outcome. The MDP interface allows us to guide the evalation/generation process to produce expressions satisfying some constraints. For instance, the following query function takes an integer and returns a computation which produces an expression e whose evaluation matches the target.
This is achieved by assigning a reward which penalizes bad answers: the reward is of the form
The core of the implementation is a representation of the choice tree.
2025
7
14
http://localhost/000I/
000I
/000I/
Memoizable choice tree monad
Definition
We're concerned with implementing the following signature, corresponding to a monad for finite non-determinism. (Note however than much of what is written here applies to more general signatures.)
('a -> 'b t) -> 'b t
val return : 'a -> 'a t
val choices : 'a array -> 'a t
end
\end{lstlisting}
]]>
A free monadic implementation as given by Free below just reflects the monadic opereations in a GADT. This is fine for some applications, but it makes memoizing choices nontrivial: the callbacks to be memoized are inside bind nodes while the choices are in choices nodes.
'a t
| Bind : 'a t * ('a -> 'b t) -> 'b t
| Return : 'a -> 'a t
let bind m f = Bind (m, f)
let return x = Return x
let choices arr = Choices arr
end
\end{lstlisting}
]]>
If bind nodes were always associated to the right as in bind (choices ...) (fun x -> bind (choices ...) ...), then we could easily memoize. We can enforce this by constraining the structure of the trees to be flat, as in:
'b flat) -> 'b flat
| Return : 'a -> 'a flat
\end{lstlisting}
]]>
But this is overly constraining for users. We can alleviate this issue by applying the Codensity construction to the type above. Rather than using an explicit Bind node to chain computations, we use continuation-passing style combined with a definition of trees which are always explicitly flattened.
'b tree) -> 'b tree
| Return : 'a -> 'a tree
type 'a codensity = {
call : 'b. ('a -> 'b tree) -> 'b tree
}
type 'a t = 'a codensity
let choices arr =
{ call = (fun k -> Choices (arr, k)) }
let bind m f =
{ call = (fun k ->
m.call (fun xm -> (f xm).call k)) }
let return x =
{ call = (fun k -> k x) }
end
\end{lstlisting}
]]>
One can obtain a tree from a codensity computation by running it:
a tree =
fun m -> m.call (fun x -> Return x)
\end{lstlisting}
]]>
And once a tree is obtained, one may explore it. For instance, random sampling from
a tree is a simple recursive function:
a tree -> a =
fun rng -> function
| Choices (arr, callback) ->
let i = Random.State.int rng (Array.length arr) in
monte_carlo_run rng (callback arr.(i))
| Return x -> x
\end{lstlisting}
]]>However, this will reallocate already explored trees. Memoization can make the structure more computationally efficient. We elide the definitions identical to that of the Codensity module.
'b tree
| Return : 'a -> 'a tree
type 'a codensity = {
call : 'b. ('a -> 'b tree) -> 'b tree
}
type 'a t = 'a codensity
let init_choices arr k =
let subtrees = Array.map (fun c -> Lazy.from_fun (fun () -> k c)) arr in
Choices (arr, subtrees)
let choices arr =
{ call = (fun k -> init_choices arr k) }
let bind m f = ...
let return x = ...
let run_codensity (m : 'a t) = ...
let rec monte_carlo_run : type a. Random.State.t -> a tree -> a =
fun rng -> function
| Choices (_, subtrees) ->
let i = Random.State.int rng (Array.length subtrees) in
monte_carlo_run rng (Lazy.force subtrees.(i))
| Return x -> x
end
\end{lstlisting}
]]>
MCTS uses multi-armed bandits as local decision-making units in nodes of a choice trees.
2025
7
13
http://localhost/000E/
000E
/000E/
Multi-armed bandits
Definition
Multi-armed bandits (MAB) are a foundational model of decision-making under uncertainty. In such models, we have:
actions (so-called arms)
an unknown family of reward distributions (assumed to have finite mean)
trials
At each trial we select an action and get a reward sampled from the distribution associated to , i.e. . An optimal strategy on average is to select the arm corresponding to the distribution with maximal mean. The goal is to devise a strategy to select actions so as to minimize regret, i.e. the difference between the obtained reward and that obtained via the optimal strategy.
Some algorithms to solve the MAB problem are -greedy bandits, UCB1, etc.Algorithms to solve the MAB problem share a common interface. We distinguish two kinds of states:
Ready: when in a ready state, the bandit may select an Action, and transition in a Waiting state.
Waiting: when in a waiting state, the bandit may receive a Reward and update its internal statistics, moving back to a Ready state.
This translates in the following interface:
ready
val select : ready -> action * waiting
val set_reward : waiting -> reward -> ready
end
\end{lstlisting}
]]>
In practice, one may want to decompose the waiting state as ready * action as we typically just need to remember the last action taken.
The MCTS tree is defined by adjoining bandits to choices nodes in a codensity-based memoizable tree monad. An implementation could look as follows:
'b tree
| Reward : {
reward : float;
subtree : unit -> 'a tree
} -> 'a tree
| Return : 'a -> 'a tree
\end{lstlisting}
]]>
However we don
'b tree
| Reward : {
reward : float;
subtree : unit -> 'a tree
} -> 'a tree
| Return : 'a -> 'a tree
type 'a codensity = {
call : 'b. ('a -> 'b tree) -> 'b tree
}
let init_choices choices k =
let subtrees =
Array.map (fun c -> Lazy.from_fun (fun () -> k c)) choices
in
let arity = Array.length choices in
Choices { bandit = Int_bandit.init arity; choices; subtrees }
let choices arr = ...
let reward r =
{ call = (fun k -> Reward { reward = r; subtree = ka })}
let bind m f = ...
let return x = ...
let run_codensity = ...
end
\end{lstlisting}
]]>
MCTS explores such trees by iterating the following three phases:
selection: starting at the root, recursively let bandit select actions until an unexplored subtree is encountered;
rollout: Monte Carlo estimation of the reward associated to the chosen actions
backup: assign reward estimate to bandits
MCTS does not prescribe when to stop iterating these three phases: one may for instance inspect how the internal statistics of the bandits evolve and decide to stop when their reward estimate is good enough.
An example might clarify what these three phases are.
2025
7
14
http://localhost/000H/
000H
/000H/
MCTS exploration
Example
We consider a simple example of MDP and how it solved by MCTS. We assume given an MCTS implementation of the MDP interface: module Mcts : MDP and use let* as a notation for Mcts.bind.
The example is set up so that at depth 0, actions are 0, 1; at depth 1, actions are 2, 3, etc. Initially, the tree only contains a root node with unexplored choices corresponding to actions 0 and 1. The root node contains a bandit which is ready to take an action.
We enter the selection phase. The bandit at the root selects action 0 and enters a waiting state. The subtree corresponding to this action was not explored: we unfolded a new node, containing a new bandit.
Since the subtree was not explored, we switch from the selection to the rollout phase, where we must estimate the reward of the choice performed by the root bandit. In basic MCTS, we randomly explore the tree until a leaf and accumulate the reward on the path. Other value estimation techniques exist, e.g. using a neural network as in AlphaGo.
We got a reward of 13. We enter the backup phase, where we assign the reward to all bandits that took a decision in this round. Since this was the first round, only the bandit at the root gets the reward.
This completes the first round. We move back to a state in which the tree is ready for the next round.
The next animation displays the next two phases.
2021
1
24
http://localhost/000S/
000S
/000S/
Automatic differentiation
article
2021
1
24
Composing jacobians
Given a (sufficiently differentiable, etc) function , its jacobian has the following type:
where denotes the space of linear maps (matrices) from to .
The goal of AD is to compute (exactly) the jacobian matrix of a function explicitly
constructed as a composition of functions with known jacobians.
The chain rule for jacobians gives us a neat formula for sequential composition:
given , , we have
Nothing special so far. Note that in order to evaluate the jacobians along the chain,
we need first to evaluate and . We could stop here
and we would have a correct algorithm for computing jacobians.
The real point of AD however is to take advantage of the sparsity of some
of these jacobians. This sparsity arises when we consider parallel composition
of functions, as obtained by taking eg
to the function with the obvious
definition. The jacobian of this last function has the following shape:
We want to take advantage of this sparsity to perform the global jacobian
computation faster. Another instance where sparsity enters the picture is when
data is duplicated. The jacobian of the diagonal operator ;
defined by , is
with the identity matrix.
Let's be pendantic about it and rewrite a small example in
pointfree format:
Interpreting elements of this decomposition as arrows in a suitable monoidal category
(which category? is another interesting question) yields the following string
diagram.
String diagrammatic depiction of
This corresponds to a mathematically grounded notion of computational graph.
Algorithmic differentiation is typically performed directly on such graphs, where
sparsity is manifest and can be easily taken advantage of.
We're left with two problems:
generate these graphs from the source expressions (this is not hard);
perform the jacobian computation taking advantage of the sparsity of the graph.
The second problem can be stated as follows: given a computational graph, what is the schedule for evaluating
the jacobian that minimizes the amount of work (eg measured in basic arithmetic operations)?
As often with interesting problems, this is NP-complete (see this paper).
Empirically, the functions people want to differentiate map a
high-dimensional domain to a lower-dimensional range (neural networks are
typical in this respect). For these functions, an evaluation strategy that
works well in practice consists in evaluating the product of jacobians
left-to-right. This is sometimes called backpropagation, or reverse-mode AD.
2021
1
24
Backpropagation
We present backpropagation as an algorithm computing the partial derivatives of a chosen output of the graph
(a DAG really) with respect to all inputs of the graph simultaneously, by working inductively on the graph and by case
analysis on the kind of nodes. We distinguish two such kinds:
nodes with fan-out, which map one input to several outputs
(in our syntax this corresponds just to nodes);
nodes with fan-in, which map several inputs to one output.
We assume that the graph has been evaluated forward, ie
each node is associated to an evaluated jacobian matrix at the point of interest.
We will thus abuse notation and let denote , leaving
the evaluation point implicit.
The state of the algorithm at each step consists, for the node under consideration,
of partial derivatives of the chosen "global" output with respect to the each
output of the node (i.e. outgoing edges of the node).
The dotted edges in the illustrative diagrams below hint at the fact that each
considered graph may be embedded in a larger one.
Base step. The initial step is to consider the output edge with respect to which we wish
to differentiate the inputs of the graph. This corresponds to some ``dangling''
edge of some arbitrary node.
A dangling edge
We must assign to this edge the partial derivative of the output with respect to itself, ie the identity matrix.
Nodes with fan-in. We consider the case of a node with fan-in as depicted below
(the dimensions have been indicated on the edges). This diagram represents
a node with two inputs from nodes and . The boxed node represents an
arbitrary subgraph that is connected to . has a terminal output of dimension
.
A node with fan-in
Our induction hypothesis is that has been computed, ie we have
the partial derivatives of with respect to the output of .
We want to compute the partial derivatives of
with respect to the outputs of the direct ancestors of , namely and .
That's straightforward: simply compute . is a matrix with
columns: the first columns correspond to the partial derivatives
of with respect to the output of and the last columns are
the counterpart for .
Nodes with fan-out. It remains to consider the case of the nodes.
Let's take the simple case of a node .
A node with fan-out
By induction hypothesis, we have the partial derivatives of with respect
to each outgoing edge of : concretely, the first outgoing edge
is associated to the first columns of and the second outgoing edge
to the last columns of .
Looking at the definition of , one obtains that the partial
derivative of the the composition of with is obtained
as the sum of the first columns of with the last
columns of .
We notice something interesting: in order to evaluate a node,
we must evaluate all its children! These acts as "rendez-vous" points
during the backpropagation procedure.
That's it!?
The algorithm above is quite simple but contains the essence of backpropagation.
Still, there's a big gap between this description and a practical implementation.
Let's name a few.
In practice, the nodes need not be explicitly represented
in the graph. Nodes have a fan-in and a fan-out. This makes it slightly
more complicated to schedule node evaluation: one needs to proceed in
reverse topological order to ensure that each node is evaluated
only once, and only when all its successors have been evaluated.
Some practical details matter for performance: for nodes with
fan-out 1]]>, one can sum the contributions of their successors into
an accumliator that needs to be allocated only once. This yields an
algorithm that is rather imperative in nature but quite efficient.
There's value in making the computation incremental,
ie reevaluate functions and jacobians as lazily as reasonably possible.
This allows eg recomputing gradients efficiently when only changing
part of a high-dimensional input.
2021
1
24
Toy implementation
A toy implementation of the above is not so long to write,
let's give it a try.
We must start with some fairly boring linear algebra.
(* Matrices in column-major format.
n.b.: [float array array] is not a very
good datatype for matrices. *)
type matrix = { c : int; r : int; data : float array array }
type vector = float array
(* Basic accessors *)
let cols m = m.c
let rows m = m.r
(* Creating matrices. *)
let empty () = { c = 0; r = 0; data = [||] }
let zeros ~c ~r = { c; r; data = Array.make_matrix c r 0.0 }
let identity n =
let id = Array.make_matrix n n 0. in
for i = 0 to n - 1 do
id.(i).(i) <- 1.
done ;
{ c = n; r = n; data = id }
(* pointwise matrix addition *)
let matrix_add m1 m2 =
{ c = m1.c; r = m1.r;
data = Array.map2 (Array.map2 ( +. )) m1.data m2.data }
(* Matrix multiplication *)
let mm m1 m2 =
assert (cols m1 = rows m2) ;
let result = Array.make_matrix (cols m2) (rows m1) 0.0 in
for c = 0 to cols m2 - 1 do
for r = 0 to rows m1 - 1 do
let acc = ref 0.0 in
for k = 0 to cols m1 - 1 do
acc := !acc +. (m2.data.(c).(k) *. m1.data.(k).(r))
done ;
result.(c).(r) <- !acc
done
done ;
let r = rows m1 in
let c = cols m2 in
{ c; r; data = result }
(* Submatrix of [m] containing the columns in the
inclusive range [c;c+dim-1]. Columns are indexed from 0. *)
let sub_cols m c dim =
let data = Array.sub m.data c dim in
{ c = dim; r = m.r; data };;type matrix = { c : int; r : int; data : float array array; }
type vector = float array
val cols : matrix -> int = <fun>
val rows : matrix -> int = <fun>
val empty : unit -> matrix = <fun>
val zeros : c:int -> r:int -> matrix = <fun>
val identity : int -> matrix = <fun>
val matrix_add : matrix -> matrix -> matrix = <fun>
val mm : matrix -> matrix -> matrix = <fun>
val sub_cols : matrix -> int -> int -> matrix = <fun>
Let's define the type of a node. We'll take the route of nodes with
fan-in and fan-out (so no explicit ).
type node =
{ transfer : vector -> vector; (** Forward evaluation *)
jacobian : vector -> matrix; (** Jacobian *)
mutable value : vector; (** Current value of the node *)
mutable accumulator : matrix; (** Accumulator *)
back : node array; (** Nodes upstream *)
uid : int (** unique id *)
};;
Most fields should be obvious, except uid.
We will use it to tag nodes so that the following relationship holds:
n1.uid <= n2.uid there exists a topological
ordering of the graph where n1 is before n2.
We'll use the following to generate strictly increasing uids.
let gen =
let c = ref 0 in
fun () ->
let v = !c in
incr c ;
v
;;val gen : unit -> int = <fun>
Let's define a generic function to compose nodes together.
The back argument corresponds to the nodes on which
the newly created node depends for its inputs.
let create ~back ~transfer ~jacobian =
{ transfer;
jacobian;
value = [||];
accumulator = empty ();
back;
uid = gen ()
}
;;val create :
back:node array ->
transfer:(vector -> vector) -> jacobian:(vector -> matrix) -> node = <fun>
As a first example, here's how to create a node for
the function.
let sin x =
create
~back:[| x |]
~transfer:(fun v -> [| sin v.(0) |])
~jacobian:(fun v -> { c = 1; r = 1; data = [| [| cos v.(0) |] |] })
;;val sin : node -> node = <fun>
A graph would be useless without some way to specify its
inputs. We'll do so using the (quite hackish) input node:
let input () =
let value = ref 0.0 in
let node =
create
~back:[||]
~transfer:(fun _v -> [| !value |])
~jacobian:(fun v -> zeros ~c:0 ~r:(Array.length v))
in
let set x = value := x in
(node, set)
;;val input : unit -> node * (float -> unit) = <fun>
This function returns a node and a "setter" function.
Note that the jacobian is empty: it has columns
(because the input node has no backward nodes)
but we must still specify that it has the right number of
rows, so that the linear algebra routines are happy.
We'll first implement code to evaluate the graph forward.
The following helper function gathers all values of the
direct ancestors (elements of node.back) of
the given node and concatenates them into a vector:
let get_backward_values node =
Array.map (fun node -> node.value) node.back
|> Array.to_list
|> Array.concat
;;val get_backward_values : node -> float array = <fun>
Evaluating the graph forward is triggered by querying the
value of a node using the eval function.
let rec eval_loop node =
Array.iter eval_loop node.back ;
let inputs = get_backward_values node in
node.value <- node.transfer inputs
let eval node = eval_loop node ; node.value
;;val eval_loop : node -> unit = <fun>
val eval : node -> vector = <fun>
Upon calling eval, the auxilliary eval_loop
function will first evaluate all ancestors recursively. After it
returns, all values of direct ancestors have been recomputed and
the node.transfer function is called on the now valid
inputs contained in these ancestors.
Let's try our forward evaluation code out:
(* test forward evaluation for a scalar function *)
let () =
(* Create graph *)
let (x, set_x) = input () in
let graph = sin (sin x) in
(* Dummy inputs *)
let inputs = Array.init 100 (fun i -> float_of_int i *. 0.1) in
(* Assert forward evaluation matches direct-style code *)
assert (
Array.for_all
(fun x ->
set_x x ;
let res = eval graph in
res.(0) = Stdlib.(sin (sin x)))
inputs ) ;
Format.printf "Success!@."
;;Success!
Let's carry on with differentiation.
In order to ensure that we process nodes in the good order, we'll
need a form of priority queues. When processing a node, we'll place
its direct ancestors in this queue. Nodes will be taken from the queue
by decreasing uid. The invariant on uids will thus ensure that
a node cannot be processed before all its (transitive) successors
are processed first.
For nodes with fan-out 1]]>, this ensures that the accumulator
contains all jacobian contributions of its direct children.
We'll use a quick to implement priority queue: a sorted list.
(* List of nodes sorted by decreasing uid *)
let rec insert x l =
match l with
| [] -> [x]
| y :: tl ->
if y.uid = x.uid then l
else if y.uid > x.uid then y :: insert x tl
else x :: l
;;val insert : node -> node list -> node list = <fun>
We can now implement backpropagation. The algorithm takes
as parameter a node corresponding to the output variable,
with respect to which all partial derivatives at the roots will
be computed (these roots corresponding to input nodes).
The first step is to initialize the graph:
ensure that all nodes upstream of output have up-to-date values
by calling eval_loop;
allocate suitably dimensioned accumulators;
initialize the accumulator of output
with the identity matrix.
The accumulator at each node must store the jacobian
corresponding to the downstream part of the graph, which is a function from
the outputs of node to output.
The function allocate_accumulators takes care of
recursively initializing all accumulators of nodes upstream of output.
let rec allocate_accumulators node output_dim =
node.accumulator <- zeros ~c:(Array.length node.value) ~r:output_dim ;
Array.iter (fun node ->
allocate_accumulators node output_dim) node.back
;;val allocate_accumulators : node -> int -> unit = <fun>
We can now move to the core loop of backpropagation. At this point,
most of it has already been explaind: no point in commenting further.
let rec backprop_loop node prioqueue =
let ja = node.jacobian (get_backward_values node) in
let acc = node.accumulator in
let start_col = ref 0 in
let queue = ref prioqueue in
for i = 0 to Array.length node.back - 1 do
let back = node.back.(i) in
let odim = Array.length back.value in
(* Accumulate slice of the product in [back] *)
let prod = mm acc (sub_cols ja !start_col odim) in
back.accumulator <- matrix_add back.accumulator prod ;
queue := insert back !queue ;
start_col := !start_col + odim
done ;
backprop_queue !queue
and backprop_queue prioqueue =
match prioqueue with [] -> () | node :: tl -> backprop_loop node tl
;;val backprop_loop : node -> node list -> unit = <fun>
val backprop_queue : node list -> unit = <fun>
The entrypoint for backpropagation proceeds to the initialization,
as described above, before calling the loop with an empty queue.
let backprop node =
eval_loop node ;
let output_dim = Array.length node.value in
allocate_accumulators node output_dim ;
node.accumulator <- identity output_dim ;
backprop_loop node []
;;val backprop : node -> unit = <fun>
Finally! Let's try it out on a simple example:
(* test backward evaluation for a scalar function *)
let () =
let (var, set_var) = input () in
let graph = sin (sin var) in
let inputs = Array.init 100 (fun i -> float_of_int i *. 0.1) in
assert (
Array.for_all
(fun x ->
set_var x ;
backprop graph ;
let derivative = var.accumulator.data.(0).(0) in
derivative = Stdlib.(cos x *. cos (sin x)))
inputs ) ;
Format.printf "Success!@."
;;Success!
let's try a more
complicated example. We'll need some more node constructors:
let ( + ) x y =
create
~back:[| x; y |]
~transfer:(fun v -> [| v.(0) +. v.(1) |])
~jacobian:(fun _v ->
{ c = 2; r = 1; data = [| [| 1.0 |]; [| 1.0 |] |] })
let ( * ) x y =
create
~back:[| x; y |]
~transfer:(fun v -> [| v.(0) *. v.(1) |])
~jacobian:(fun v ->
{ c = 2; r = 1; data = [| [| v.(1) |]; [| v.(0) |] |] })
let pair x y =
create
~back:[| x; y |]
~transfer:(fun v -> [| v.(0); v.(1) |])
~jacobian:(fun _v ->
{ c = 2; r = 2; data = [| [| 1.; 0.0 |]; [| 0.0; 1.0 |] |] })
;;val ( + ) : node -> node -> node = <fun>
val ( * ) : node -> node -> node = <fun>
val pair : node -> node -> node = <fun>
Now, for the final example: I crafted a function from to
and I compared the outcome of our algorithm with
a symbolic solution painstakingly copied from Wolfram alpha.
let approx_eq (f1 : float) (f2 : float) = abs_float (f1 -. f2) < 0.0001
let actual_jacobian x y =
let open Stdlib in (* needed because we shadowed the original defs *)
[| [| cos ((x +. sin y) *. sin (x *. y))
*. (sin (x *. y) +. (y *. (x +. sin y) *. cos (x *. y)));
(y *. y *. sin (x +. sin x))
+. (x *. y *. y *. (cos x +. 1.) *. cos (x +. sin x))
|];
[| cos ((x +. sin y) *. sin (x *. y))
*. ((x *. (x +. sin y) *. cos (x *. y)) +. (cos y *. sin (x *. y)));
2. *. x *. y *. sin (x +. sin x)
|]
|]
(* test backward evaluation for a vector function *)
let success =
let (v1, set_v1) = input () in
let (v2, set_v2) = input () in
let graph =
let graph1 = sin ((v1 + sin v2) * sin (v1 * v2)) in
let graph2 = v1 * v2 * v2 * sin (v1 + sin v1) in
pair graph1 graph2
in
let inputs1 = List.init 100 (fun i -> float_of_int i *. 0.1) in
let inputs2 = List.init 100 (fun i -> 50. +. (float_of_int i *. 0.1)) in
List.for_all2
(fun x y ->
set_v1 x ;
set_v2 y ;
backprop graph ;
let ja = actual_jacobian x y in
let dv1_dx = v1.accumulator.data.(0).(0) in
let dv1_dy = v1.accumulator.data.(0).(1) in
let dv2_dx = v2.accumulator.data.(0).(0) in
let dv2_dy = v2.accumulator.data.(0).(1) in
(approx_eq ja.(0).(0) dv1_dx) &&
(approx_eq ja.(0).(1) dv1_dy) &&
(approx_eq ja.(1).(0) dv2_dx) &&
(approx_eq ja.(1).(1) dv2_dy))
inputs1
inputs2
;;val approx_eq : float -> float -> bool = <fun>
val actual_jacobian : float -> float -> float array array = <fun>
val success : bool = true
That's it!
Contributions