OCANNL -- OCaml Compiles Algorithms for Neural Networks Learning
A from-scratch, compiled Deep Learning framework.
Implements backpropagation (i.e. first-order reverse mode autodiff) and shape inference.
The long-term goal is to provide several "low-level" backends, aiming to seek inspiration from projects such as TinyGrad, TVM, Luminal.
OCANNL starts with a high-level representation, but can compile everything down to for loops.
The library users can compile any amount of code into a monolithic routine. Depending on the use case:
the whole training update step can be a single routine,
or the step can be composed of a gradient update routine (a forward pass and a backprop pass) and a params update routine (e.g. SGD with momentum, ADAM, etc.),
or the user can compile parts of a model separately, manually composing the corresponding forward pass code and the backprop code.
Tensor axes are split into kinds: batch, input and output. Tensor dimensions have optional labels.
The labels ensure a more precise semantics for dimension matching.
In the future we might introduce axis labels as an alternative to positional axis selection, it would be a separate naming mechanism.
OCANNL has full support for the einsum notation, integrated with shape inference. Supports static indexing, with a built-in operation to take a slice of the batch axes, integrated with shape inference. Extensible to more static indexing patterns as needs arise.
OCANNL does not have dynamic indexing (using the last axis of one tensor as indices into another tensor). If it's needed, it can be added (we had a prototype once, removed to reduce complexity). Then it would also be integrated with shape inference.
OCANNL has a suite of tutorials doubling as tests with inline expectations.
OCANNL offers two main levels of abstraction.
Differentiable computations, centered around the %op syntax extension.
Plain computations, centered around the %cd syntax extension. It integrates the arrayjit backend library with shape inference.
The support for mixed-precision computations is upcoming.
E.g. higher-precision network components, or gradients at a higher precision than values.
Currently (v0.3), you can select the precision, and individual computation nodes track their precision, but mixing precisions might break things.
Should be easily extensible.
Model surgery should be starightforward (not sure if we are there yet).
It's a feature, not a bug!
To scale a tensor by a number, always use pointwise-multiplication, e.g. 2*.m or m*.2.
Matrix-multiplying a tensor m by a constant number, e.g. m*2, broadcasts the number to the shape of the input axes of the tensor. This results in an output-axes-only tensor (multi-axis-vector) that is the scaled sum over the input axes of the tensor m.
Matrix-multiplying a constant number by a tensor m, e.g. 2*m, broadcasts the number to the shape of the output axes of the tensor. This results in a tensor whose inputs are of the same shape as the inputs of m, and the output shape is 1D (scalar), that is the scaled sum over the output axes of the tensor m.
The matrix-multiply operation behaves pointwise along the batch axes.
Upcoming milestones
On the critical path for the next major release v0.4:
Restore signs of life for the Cuda backend.
Mixed-precision computations: working and convenient.
Train has the optimization "frontend" and utilities.
arrayjit, which may one day become a standalone library: generates the code, performs backend-agnostic optimizations (virtual nodes whose computation is inlined), implements the backends.
Some aspects that are more core to OWL are less encapsulated in OCANNL, so it should be more natural to extend them.
Specifically, Operation and Train are just collections of functions.
OCANNL provides lower-level compilation backends than OWL, it is more self-contained in this sense.
Installation
Although the project is called ocannl, the main package is called neural_nets_lib, to avoid the (opam linter's) complaint that the name can be confused with other packages. This also clarifies that ocannl is composed of arrayjit and neural_nets_lib.
The dependency on ocaml-cudajit is optional, so you have to install it first to enable the Cuda backend.