github ↗

raven · beta 1

Scientific computing for OCaml, all the way down to the GPU.

Arrays, gradients, neural networks, numerical methods, probabilistic inference, dataframes and plots. Libraries that compose as ordinary OCaml code.

opam install raven
five-minute quickstart →
new to OCaml? set up opam first →

nx

let () =
  (* Five evenly spaced float32 values from 0 to 10, and their squares. *)
  let x = Nx.linspace Nx.float32 0. 10. 5 in
  let y = Nx.Infix.(x * x) in
  Nx.print x;
  Nx.print y;

  (* Two independent keys from seed 0, one per random matrix. *)
  let keys = Nx.Rng.split (Nx.Rng.key 0) in
  let a = Nx.Rng.uniform ~key:keys.(0) Nx.float32 [| 3; 3 |] in
  let b = Nx.Rng.uniform ~key:keys.(1) Nx.float32 [| 3; 3 |] in
  let c = Nx.matmul a b in
  Format.printf "matmul shape: %a@." Nx.pp_shape (Nx.shape c)
n-dimensional arrays · like NumPynx docs →

rune

let params = Nx.Ptree.(pair tensor tensor)

let inputs = Nx.create Nx.float64 [| 4 |] [| 0.; 1.; 2.; 3. |]
let targets = Nx.create Nx.float64 [| 4 |] [| 1.; 3.; 5.; 7. |]

let mse (w, b) =
  let err = Nx.Infix.((w * inputs) + b - targets) in
  Nx.mean Nx.Infix.(err * err)

let step =
  Rune.jit
    Nx.Ptree.(params @-> returns (pair tensor params))
    (fun ((w, b) as p) ->
      let loss, (gw, gb) = Rune.value_and_grad params mse p in
      (loss, (Nx.Infix.(w - (gw *$ 0.1)), Nx.Infix.(b - (gb *$ 0.1)))))

let rec train n p =
  let loss, p = step p in
  if n = 1 then (loss, p) else train (n - 1) p

let loss, (w, b) =
  train 200 (Nx.scalar Nx.float64 0., Nx.scalar Nx.float64 0.)
automatic differentiation · like JAXrune docs →

kaun

open Kaun

module Mlp = struct
  type ('a, 'd) t = { l1 : ('a, 'd) Linear.t; l2 : ('a, 'd) Linear.t }

  let walk c { l1; l2 } =
    let open Nx.Ptree.Walk in
    let l1 = field c "l1" Linear.walk l1 in
    let l2 = field c "l2" Linear.walk l2 in
    { l1; l2 }

  let apply p x = Linear.apply p.l2 (Nx.tanh (Linear.apply p.l1 x))
end

let mlp = Nx.Ptree.instantiate (module Mlp)

let () =
  Nx.Rng.with_key (Nx.Rng.key 42) @@ fun () ->
  let x =
    Nx.create Nx.float32 [| 4; 2 |] [| 0.; 0.; 0.; 1.; 1.; 0.; 1.; 1. |]
  in
  let y = Nx.create Nx.float32 [| 4; 1 |] [| 0.; 1.; 1.; 0. |] in
  let params =
    {
      Mlp.l1 = Linear.init ~key:(Nx.Rng.next_key ()) ~inputs:2 ~outputs:8;
      l2 = Linear.init ~key:(Nx.Rng.next_key ()) ~inputs:8 ~outputs:1;
    }
  in
  let loss p = Loss.sigmoid_bce (Mlp.apply p x) y in
  let step (params, ostate) =
    let l, grads = Rune.value_and_grad mlp loss params in
    let params, ostate =
      Vega.adam_step mlp ~lr:(Vega.lr 0.05) ostate ~params ~grads
    in
    ((params, ostate), Nx.item [] l)
  in
  let state = ref (params, Vega.adam_init mlp params) in
  for _ = 1 to 500 do
    let s, _ = step !state in
    state := s
  done;
  Printf.printf "%s\n"
    (Nx.to_string (Nx.sigmoid (Mlp.apply (fst !state) x)))
  (* predictions approach [0; 1; 1; 0] *)
neural networks · like Flaxkaun docs →

brot

open Brot

let tokenizer =
  bpe
    ~vocab:
      [ ("h", 0); ("e", 1); ("l", 2); ("o", 3); (" ", 4); ("w", 5);
        ("r", 6); ("d", 7); ("he", 8); ("ll", 9); ("llo", 10);
        ("hello", 11); ("wo", 12); ("rl", 13); ("rld", 14); ("world", 15) ]
    ~merges:
      [ ("h", "e"); ("l", "l"); ("ll", "o"); ("he", "llo");
        ("w", "o"); ("r", "l"); ("rl", "d"); ("wo", "rld") ]
    ()

let encoding = encode tokenizer "hello world"
let ids = Encoding.ids encoding         (* [| 11; 4; 15 |] *)
let tokens = Encoding.tokens encoding   (* [| "hello"; " "; "world" |] *)
let decoded = decode tokenizer ids      (* "hello world" *)
tokenization · like HF Tokenizersbrot docs →

talon

open Talon

let ( let* ) = Result.bind
let delay = Col.int "dep_delay"

let late_by_carrier () =
  let* flights = Talon_csv.file ~nulls:[ "NA" ] "flights.csv" in
  let* carriers = Talon_parquet.file "carriers.parquet" in
  Query.(
    of_source flights
    |> filter Expr.(delay > int 15)
    |> aggregate ~by:[ "carrier" ]
         Expr.[ "mean_delay" := mean delay; "flights" := rows ]
    |> join ~on:(Join.keys [ "carrier" ]) ~each_left:One
         (of_source carriers)
    |> sort [ Order.desc "mean_delay" ]
    |> run)
dataframes · like pandas / Polarstalon docs →

jera

open Jera

let () =
  (* The cube root of each c by bracketing, and its derivative in c. *)
  let cube_root c =
    Root.bracket ~tol:(Tol.ulps 4.)
      (fun x -> Nx.sub (Nx.mul x (Nx.square x)) c)
      ~lo:(Nx.zeros_like c) ~hi:(Nx.full_like c 10.)
    |> Solution.get
  in
  let c = Nx.create Nx.float64 [| 3 |] [| 2.; 8.; 27. |] in
  Nx.print (cube_root c);
  Nx.print (Rune.grad' (fun c -> Nx.sum (cube_root c)) c);

  (* A pendulum q'' = -sin q, solved to t = 10 within a tolerance. *)
  let pendulum _t (q, p) = (p, Nx.neg (Nx.sin q)) in
  let s x = Nx.scalar Nx.float64 x in
  let q, _ =
    Ode.solve Nx.Ptree.(pair tensor tensor) Ode.tsit5
      ~tol:(Tol.v ~rel:1e-10 ~abs:1e-12) ~budget:1000 pendulum
      ~t0:(s 0.) ~t1:(s 10.) (s 1., s 0.)
    |> Solution.get
  in
  Nx.print q
numerical methods · like SciPyjera docs →

hugin

open Hugin

let steps = Nx.linspace Nx.float64 0. 1. 200

let losses =
  let rates = Nx.create Nx.float64 [| 5; 1 |] [| 2.; 2.5; 3.; 3.5; 4. |] in
  let noise = Nx.Rng.normal ~key:(Nx.Rng.key 0) Nx.float64 [| 5; 200 |] in
  Nx.add (Nx.exp (Nx.neg (Nx.mul rates steps))) (Nx.mul_s noise 0.01)

let () =
  line
    ~x:(num ~title:"progress" steps)
    ~y:(num ~title:"loss" losses)
    ~stroke:(dim ~title:"run" 0)
    ()
  |> title "Five runs"
  |> save "runs.png"
plotting · like Matplotlibhugin docs →
raven's counterparts toNumPyJAXPyTorchSciPyNumPyropandasMatplotlibJupyterW&Bcoming from python →

01 — one program, every layer

gpt-oss-20b runs on NVIDIA and AMD GPUs, with every layer in Raven.

Each layer of the model, attention and a mixture of experts, compiles into GPU kernels under Rune.jit. Six libraries carry a prompt to its answer without leaving the language.

  1. brottokenizes the prompt
  2. kaundefines the model
  3. nxholds the weights and the cache
  4. runecompiles each layer
  5. tolkturns it into GPU kernels
  6. rigruns them on the GPU
kaun/examples/06-gpt-oss/gpt_oss.mlfull example ↗
let block cfg layer b cache index x =
  let eps = cfg.norm_eps in
  let a, cache =
    attention cfg layer b cache index (Rms_norm.apply ~eps b.attn_norm x)
  in
  let x = close (Nx.add x a) in
  let h = Rms_norm.apply ~eps b.ffn_norm x in
  let routing = Moe.route ~k:cfg.experts_per_token (Linear.apply b.router h) in
  let experts = Moe.apply ~limit:cfg.swiglu_limit b.moe routing h in
  (close (Nx.add x experts), cache)

98 tok/s

decoding gpt-oss-20b on an RTX 5000 Ada, jitted, one token per 10.2 ms step. llama.cpp: 146 tok/s. The goal post on the road to 1.0: at least as fast as four baselines.

2026-11the goal post →

7.6 MiB

to install rig on NVIDIA through its own driver (115 MiB through CUDA). PyTorch: 5,458 MiB. Compiling kernels loads NVRTC.

2026-10-11source ↗

70 ms

from starting a process to its first GPU result, through rig's own driver (134 ms through CUDA). PyTorch: 720 ms. JAX: 538 ms.

2026-10-11source ↗

02 — the ecosystem

Thirteen libraries. Each does one thing well.

Install the whole stack with opam install raven, or only what you need: opam install nx for arrays alone.

foundation

like CUDA driver API
rigDevices such as GPUs, and the order of the work they run.
like NumPy
nxN-dimensional arrays on the host and on devices, with linear algebra and FFT.
like tinygrad
tolkA minimal ML compiler that fuses tensor programs into GPU kernels.
like JAX
runeAutomatic differentiation, vectorizing maps and compilation.

machine learning

like Flax
kaunNeural networks and training, on typed parameter trees.
like Optax
vegaComposable gradient-based optimizers.
like HF Tokenizers
brotTokenization that reads HuggingFace tokenizer.json.

scientific computing

like SciPy
jeraSystems, minima, integrals, interpolation, ODEs and SDEs.
like NumPyro / Stan
nornProbabilistic models, samplers and diagnostics.

data and visualization

like pandas / Polars
talonDataframes with typed queries, CSV and Parquet.
like Matplotlib
huginPublication-quality plotting.

tools

like Jupyter
quillAn interactive REPL and notebooks as markdown files.
like W&B / MLflow
muninLocal experiment tracking with a live terminal dashboard.

contrib

Built against every change to the core; released on their own schedule.

fehuReinforcement learning for OCamlsowiloDifferentiable computer vision for OCamlymirAstronomy for OCaml

03 — why raven

Prototype like Python. Ship like OCaml.

i.

Dtype and device mistakes stop at the compiler.

An array's type carries its dtype and the device set it lies on, so values from two device sets never meet by mistake. Shapes are checked when an operation runs. Why Raven →

ii.

JAX's ergonomics, in a real functional language.

Every Nx function issues a small set of primitive operations. grad, vmap and jit are OCaml 5 effect handlers that interpret them, so they nest like any other function.

iii.

GPUs with no CUDA, no ROCm, not even a driver.

rig drives NVIDIA and AMD GPUs through their kernel drivers, or, on Linux, with no driver at all: it boots the GPU itself from the vendor's firmware. Precompiled kernels run with nothing else installed; compiling them at run time loads NVRTC or comgr. On the road to 1.0: your whole application, model included, as one static binary. rig →

What people who use Raven say

in use

We picked OCaml for the language and put up with the numerical ecosystem for years. Raven finally gives us both: type-safe models, seamless GPU acceleration, and an effects-based design that’s a pleasure to use. Porting and unifying our previously disparate libraries was effortless. We have never been more productive.”

Guillaume Hennequin ↗Professor of Computational Neuroscience, University of Cambridge

built on raven

  • sofoSecond-order optimization with no backpropagation · NeurIPS 2024
  • symoCurvature from weight-space symmetries · ICML 2026
  • ilqr-vaeNeural dynamics inferred by optimal control · ICLR 2022
  • ModelKitClassical machine learning, like scikit-learn
  • geomGeometric algebra for robotics
  • SoundMLAudio processing
all community projects →

04 — learn

Start from where you are.

all documentation →
new to raven · 5 minQuickstartFrom an empty directory to arrays, gradients and a trained model.know numpy or pytorchComing from PythonLibrary by library mappings, then the differences: explicit types, array literals, slicing, functional transformations.choosing librariesEcosystem overviewHow the libraries layer on each other, and which to reach for.learn by reading codeExamplesXOR and MNIST up to GPT-2, Llama with checkpoint import, gpt-oss and neural ODEs.explore interactivelyNotebooks with QuillA REPL and notebooks that are plain markdown files, diffable and reviewable.look something upIndex, A to ZEvery page, guide, example and public function, filtered as you type.

05 — roadmap

The road to replacing a Python stack in production.

full roadmap →

  1. alpha · 2025-11

    The libraries, eager on the host.

    FFT and a full linear algebra suite in Nx, dataframes in Talon, and the decision to port tinygrad to OCaml.

  2. beta 1 · done

    Compiled, on GPUs.

    A whole training step compiles into one GPU program, with every layer in Raven.

  3. toward 1.0 · now

    gpt-oss, against the fastest.

    JIT-compiled gpt-oss-20b at least as fast as tinygrad, JAX, llama.cpp and vLLM, on NVIDIA and AMD.

06 — from the blog

Notes from building Raven.

all posts →rss

support

Raven is built in public, and funded by its sponsors.

sponsored by

Ahrefs

with thanks to our past sponsor

Tarides

sponsor ravengood first issues