package neural_nets_lib
Install
Dune Dependency
Authors
Maintainers
Sources
md5=179e3bf2a47d49e71eaffa8229bcf1b8
sha512=5e93c09c065ff34b7bdc32be9e5783a5ee22c7506c08d78201a81ab749708354a487a60bd101a80c0c087ba0937e055e4875295ab4c2c2ea107cc57a4241da44
Description
OCaml Compiles Algorithms for Neural Networks Learning is a compiled Deep Learning framework that puts emphasis on lowlevel backends (like tinygrad), shape inference, concise notation (ab)using PPX.
README
ocannl
NOTE TO POTENTIAL CONTRIBUTORS: reach out so I can adjust my work style  start using branches for refactoring. Otherwise you face frustration as the code might be broken. Tagged versions of the code are guaranteed to work as well as the given stage of the project permitted.
NEWS: the upcoming version 0.4.0 has significant design changes around the backend API and synchronization, with Cuda streams exposed as virtual devices. Version 0.4.1 will come soon after, with tests for (and small improvements to) mixed precision computation; it will require OCaml 5.2.
OCANNL is sponsored by Ahrefs! Visit the Ahrefs website.
OCANNL  OCaml Compiles Algorithms for Neural Networks Learning
A fromscratch, compiled Deep Learning framework.
Implements backpropagation (i.e. firstorder reverse mode autodiff) and shape inference.
The longterm goal is to provide several "lowlevel" backends, aiming to seek inspiration from projects such as TinyGrad, TVM, Luminal.
OCANNL starts with a highlevel representation, but can compile everything down to
for
loops.
The library users can compile any amount of code into a routine (i.e. a compilation unit). The user decides explicitly what the scope of a compilation unit is, by putting together the corresponding code. 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 builtin 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.%op
stands for "operation", it's meant to express tensors:Tensor.t
, and tensor functions.
Plain computations, centered around the
%cd
syntax extension. It integrates thearrayjit
backend library with shape inference.%cd
stands for "code", it's meant to express assignments:Assignments.t
.
The support for mixedprecision computations is upcoming.
E.g. higherprecision 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 pointwisemultiplication, e.g.
2*.m
orm*.2
.Matrixmultiplying 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 outputaxesonly tensor (multiaxisvector) that is the scaled sum over the input axes of the tensorm
.Matrixmultiplying 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 ofm
, and the output shape is 1D (scalar), that is the scaled sum over the output axes of the tensorm
.The matrixmultiply operation behaves pointwise along the batch axes.
Usage
A possible route to learning OCANNL:
Get some basic grasp of the aims and design of the project by reading or skimming files in test/ and bin/.
Read the syntax extensions documentation lib/syntax_extensions.md.
Read the introductory part of the shape inference documentation lib/shape_inference.md.
Improve your understanding by reading or skimming lib/shape.mli, lib/tensor.mli, lib/operation.ml, lib/train.ml, and (since 0.4.1) lib/nn_blocks.ml.
Read the implementation overview:
Shape inference details lib/shape_inference.md.
Backendindependent optimizations arrayjit/lib/lowering_and_inlining.md  lowering means translating (compiling) from the highlevel representation (as assignments) to the lowlevel representation.
More documentation to come.
Upcoming milestones
This is very tentative.
0.4.1
Half precision. Maybe improvements for mixedprecision computations.
Resolve remaining issues with the new scheduler.
Initial version of lib/nn_blocks.ml.
0.5
More of primitive numeric operations.
Useful building blocks for models in lib/nn_blocks.ml.
A language model example.
0.6
Getting more out of GPUs: better CUDA generation.
Releases
For more details, see CHANGES.
v0.4 merge buffers, Csyntax backend builder: a significant refactoring of the API.
v0.3 shape inference, jitted routines: a major rewrite of the whole project.
v0.3.3: continuous integration and opam release.
v0.3.2: new shape inference feature: tracking leftmost axes  complete inference for splicing, ellipsisinthemiddle allowed in einsum notation.
v0.3.1: sanitizing code inclusion (rootness checks).
v0.3.0: declarative shape inference; replaced the session interface with a "jitted code routines" API. Cuda defunct.
v0.2 inching toward GPU:
v0.2.1 naivecuda: a Cuda backend where blocks and threads are exposed via dedicated axis types.
v0.2.0 stackasdevice: treating the C function stack as the "device memory".
v0.1 GCCJIT backend:
v0.1.2: multicore computations using a threadlocal "task id" index.
v0.1.1: inlining scalar constants, improved inlining for virtual nodes.
v0.1.0: a
Gccjit
backend, single and double precision floats, code compiled as a monolithic update step function.
v0.0 untagged: basic design around shape inference, highlevel and lowlevel code representation. Nowabandoned MetaOCaml and OCaml backends.
Why not just use OWL?
OCANNL follows different design choices than OWL. For example:
OCANNL is not functorized.
OCANNL has fewer abstraction layers.
OCANNL has a more powerful shape inference.
OCANNL only supports backpropagation, while OWL supports full forward and backward autodiff.
Some aspects are more centralized in OCANNL than in OWL and form the "infrastructure":
Tensor indexing mechanisms are not extensible, other than changing OCANNL code.
Shape inference is fully handled by OCANNL and not extensible, other than changing OCANNL code.
Tensor
implements "putting pieces together".Train
has the optimization "frontend" and utilities.arrayjit
, which may one day become a standalone library: generates the code, performs backendagnostic 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.
OCANNL provides lowerlevel compilation backends than OWL, it is more selfcontained 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 ocamlcudajit
is optional, so you have to install it first to enable the Cuda backend.
Dependencies (16)

patdiff
>= "v0.15.0"

ppx_minidebug
>= "2.0.0"
 ppx_expect
 ppx_jane
 ppxlib
 num
 stdio

angstrom
>= "0.15"
 ocannl_npy
 printboxtext
 printbox

arrayjit
= "0.4.0"
 core
 base

dune
>= "3.11"

ocaml
>= "5.1.0"
Used by
None
Conflicts
None