Module `Kaun`

Neural networks as typed parameter records.

Kaun has no layer or trainer abstraction. A layer is a plain record of tensors with an apply function; a model is a record of layers with hand-written one-line traversals (Nx.Ptree.S plus checkpoint names, see Linear for the pattern); a training step composes Rune.value_and_grad with a Vega optimizer update; the training loop is ordinary Seq iteration over Data minibatches.

let step (params, ostate) (x, y) =
  let loss p = Loss.softmax_cross_entropy_sparse (Model.apply p x) y in
  let l, grads = Rune.value_and_grad (module Model) loss params in
  let params, ostate =
    Vega.adamw_step (module Model) ~lr:1e-3 ostate ~params ~grads
  in
  ((params, ostate), Nx.item [] l)

Random initialization and shuffling draw from the implicit RNG scope; wrap the program in Nx.Rng.run for reproducibility. Randomness inside a jitted training step (Dropout) instead takes an explicit Rune.Rng key threaded through the step's inputs.

Layers

Parameterized building blocks. Each module pairs a parameter record with constructors (Linear.init, Linear.make), an apply function, and the traversals that make it compose into differentiable, checkpointable models.

module Linear : sig ... end

Dense (fully connected) layers, and the model-as-record pattern.

module Embedding : sig ... end

Token-id to dense-vector lookup tables.

module Conv : sig ... end

2-D convolution layers (NCHW).

module Attention : sig ... end

Multi-head self-attention, and the pure Attention.scaled_dot_product_attention core.

module Layer_norm : sig ... end

Layer normalization over the feature axis.

module Batch_norm : sig ... end

Batch normalization, with running statistics as explicit state.

Stateless functions

Pure operations with no parameters: activations, pooling, dropout.

module Fn : sig ... end

Activation functions (relu, gelu, softmax, ...).

module Pool : sig ... end

2-D max and average pooling.

module Dropout : sig ... end

Dropout with an explicit ~training flag and optional Rune.Rng key.

module Init : sig ... end

Weight initializers (Glorot, He, LeCun, variance scaling).

Training

module Loss : sig ... end

Scalar training objectives: regression and classification losses.

module Data : sig ... end

Minibatch Seq.ts over in-memory tensors, with shuffling.

module Metric : sig ... end

Evaluation metrics: accuracy, precision/recall/F1, AUC-ROC.

Persistence

module Checkpoint : sig ... end

Save and load parameter structures as safetensors checkpoints.