Raven Documentation

Fehu is a reinforcement learning environment toolkit for OCaml. It is designed to provide type-safe environments that follow the Gymnasium interface pattern, where environments expose reset and step methods with typed observation and action spaces.

Key capabilities include:

• Type-safe environments: Observation and action spaces are encoded in the OCaml type system.
• Composable wrappers: Transform observations, actions, or rewards (e.g., map_observation, map_action, map_reward, clip_action, clip_observation, time_limit).
• Data handling: Trajectory collection (structure-of-arrays), replay buffers (fixed-capacity circular buffers), and GAE (generalized advantage estimation) computation.
• Execution utilities: Policy evaluation and vectorized environments for batched execution and auto-reset.

Norn is an MCMC sampling library for OCaml with the following core capabilities:

• Sampling Methods: High-level Norn.hmc and Norn.nuts with automatic adaptation.
• Customization: Use Norn.sample with make_kernel for full control over integrators, metrics, and adaptation.
• Automatic Gradients: Log-density gradients are computed automatically using Rune.
• Symplectic Integrators: Supports leapfrog, mclachlan, and yoshida.
• Mass Matrix Metrics: Supports unit_metric, diagonal_metric, and dense_metric.
• Adaptation: Implements Stan-style adaptation using dual averaging for step size and Welford estimation for the mass matrix.
• Diagnostics: Provides tools for ess (effective sample size) and rhat (split R-hat).

Kaun is a neural network library for OCaml built on rune autodiff. Unlike object-oriented frameworks, Kaun provides building blocks—layers, activations, initializers, losses, data batching, metrics, and checkpoints—as plain records and pure functions.

Key architectural principles:

• No Layer Objects or Trainers: A model is a typed record you define, and a training step is a set of functions you own end-to-end.
• Composition via Nx: Kaun uses Nx.Ptree.S (the traversal interface from nx) to glue components together. Transformations (rune) and optimizers (vega) compose through the single record type you define for your model.
• Statelessness: Most components are pure functions or parameter records with explicit apply and init methods.

Raven is a collection of nine specialized libraries that all share a single core data type: Nx.t (the n-dimensional array). This design allows the libraries to compose seamlessly through tensors.

Core Architecture:

• Nx is the foundation; every library operates on Nx.t tensors.
• Rune provides functional transformations (like gradients) on top of Nx.
• Kaun builds on Rune to provide neural network components.
• Other libraries like Sowilo (Computer Vision) and Fehu (Reinforcement Learning) use Nx directly, making their operations compatible with Rune's grad and vmap because they are plain Nx computations under the hood.

The Nx OxCaml backend is an experimental, high-performance implementation of the Nx backend interface. It is designed to improve performance by leveraging OxCaml's unboxed types, which reduces garbage collection (GC) pressure and avoids the overhead of boxing/unboxing.

Key performance features include:

• Unboxed arithmetic: Utilizes float#, int32#, and int64# for zero-allocation numeric operations.
• Parallel execution: Includes built-in support for parallel operations (currently sequential, with Domain support planned).
• Memory efficiency: Minimizes memory overhead through unboxed types.

Raven is a machine learning stack for OCaml designed to bridge the gap between rapid prototyping and production-scale deployment. It aims to replace the Python-centric workflow by providing a single-language solution that offers strong typing, JIT compilation, and efficient deployment (e.g., as MirageOS unikernels or static binaries).

The Raven ecosystem consists of several core components:

• Nx: NumPy equivalent.
• Rune: JAX equivalent with effects-based autodiff.
• Kaun: Flax equivalent.
• Brot: Tokenization.
• Hugin: Matplotlib equivalent.
• Quill: Notebook environment.

Tolk is a minimal compiler for GPU tensor computation, ported from tinygrad to OCaml. It is designed to take tensor-level computation graphs, optimize them through a multi-stage pipeline, and emit efficient kernels for various hardware backends.

Supported backends include:

• CPU (via Clang)
• Metal (macOS GPU)
• CUDA
• OpenCL