Bagl compiler
I wrote an OCaml compiler for Bagl, a small typed functional language that checks tensor shapes at compile time and derives the gradient of a function inside the compiler.
What it is
Bagl is a small functional language with tensors as a built-in type, where a tensor is an array of floats with one or two dimensions. I wrote its compiler in OCaml, together with a bytecode virtual machine, a REPL (an interactive prompt), a language server for editors, and a browser playground.
- Tensor shapes are part of the type, so multiplying a 2x3 matrix by a 2x2 matrix is a compile error.
gradturns a function into its derivative. For a tensor loss, the compiler writes the gradient code itself, so you can train a model in plain Bagl.
The whole compiler also runs in the browser through js_of_ocaml, as about 160 KB of JavaScript (50 KB compressed). You can try the playground or read the language reference, where you can edit and run the examples.
How it works
A hand-written lexer and a recursive descent parser build a syntax tree. Every node keeps its place in the source, so errors point at the exact line.
Next, the compiler rewrites grad calls into ordinary Bagl code. Scalar functions get symbolic derivatives. Functions of a tensor get reverse mode code, which computes the gradient for all inputs in one backward pass.
Type inference uses Hindley-Milner, a method that works out every type without annotations. Type variables live in a union-find structure, so unifying two types links them in place. Tensor dimensions have their own union-find variables, and that is where shape errors are caught.
The program is then lowered to an intermediate representation made of basic blocks, which are runs of code with no jumps inside. Four optimization passes run on it. Code generation then emits stack bytecode for the virtual machine, and a .baglc file can store it.
Decisions
- I chose OCaml because algebraic data types and pattern matching keep the definitions of the tree, the intermediate representation and the bytecode short and safe. In Rust, the shared union-find graph would conflict with the borrow checker.
gradruns before type inference, so later stages need no special cases, and derivatives are type checked like other code. The cost is a required annotation on tensor parameters. Anything the rewrite can't handle is a compile error, never a wrong gradient.- Bagl has no type classes, so
+picks int or float from its operands. The checker stays single pass, butfn x -> x + xis typedint -> int. - The virtual machine is a stack machine, because code generation for it is close to mechanical and the bytecode is easy to save. A register machine would need a register allocator.
What was hard
The virtual machine has separate int and float instructions, and lowering picks one by running inference again on each subexpression. When a binding was missing from that environment, lowering fell back to int, and correct float programs crashed the virtual machine. I fixed three cases of the bug, once for lambda parameters and twice for letrec. An early grad test found the first one, because the derivative of x * x is x + x, a shape no hand-written test had used.
Other fixes followed.
(x + 1.0) + xtype checked, but(x + x) + 1.0failed. Inference now waits to default to int until a binding is generalized.- A top level
ifcrashed code generation, because local variables were allocated lazily. The whole test suite had passed, since no test had that shape. 0 - 5came back from a.baglcfile as 4294967291. The reader now extends the sign, and round trip tests cover it.
Results
- 104 test cases in
test/test_bagl.mlcover the lexer, the parser, types, lowering, tensors, gradients, and programs that must be rejected. - Tests compare gradients with finite differences, i.e., slopes measured from tiny changes to the input, and they agree within 1e-4.
examples/train_xor.baglruns 200 steps of gradient descent, and a test checks that the weights reach the exact solution [1, 1, -2] within 1e-3.- CI runs the tests, the examples, and all 21 examples from the reference page through the browser bundle.
What I would do next
- Lowering should read types from inference instead of inferring them again. My notes call a typed syntax tree the next refactor.
- The intermediate representation isn't in SSA form, where each variable is assigned once, so common subexpression elimination stays inside one block. SSA would let it work across blocks.
- A
Numconstraint would make the numeric operators fully polymorphic.