Toro

Neural Networks

Toro.NN provides layers, modules, and composition tools for neural networks. Open the namespace to get started:

open Toro
open Toro.NN

Module Interfaces

Toro defines a generic interface for composable layers:

  • IModule<'In, 'Out> -- forward: 'In -> 'Out. The general-purpose module interface, composable via the pipeline { } CE.
  • IModule -- inherits IModule<Tensor, Tensor>. A shortcut for the most common case. Most layers implement this interface.

Layers with richer signatures (e.g. Dropout, BatchNorm) expose typed member methods like forwardT instead of implementing IModule. Use partial application (drop.forwardT train) to get a Tensor -> Tensor function for pipeline composition.

Layers

Linear

A fully connected layer. Implements IModule:

let l = Linear.init 784 256 F32 Cpu
let out = l.forward input

Linear.initNoBias creates a linear layer without a bias term.

Conv1d

A 1D convolution layer. Implements IModule:

let conv = Conv1d.initDefault 3 16 5 F32 Cpu  // inCh=3, outCh=16, kernel=5
let out = conv.forward input

Configure with Conv1dConfig:

let conv = Conv1d.init 3 16 5 { Conv1dConfig.defaultConfig with Padding = 2; Stride = 2 } F32 Cpu

Conv1d.initNoBias creates a layer without a bias term.

Conv2d

A 2D convolution layer. Implements IModule:

let conv = Conv2d.initDefault 1 32 3 F32 Cpu  // inCh=1, outCh=32, kernel=3
let out = conv.forward input

Configure with Conv2dConfig (same fields as Conv1dConfig: Padding, Stride, Dilation, Groups):

let conv = Conv2d.init 1 32 3 { Conv2dConfig.defaultConfig with Padding = 1 } F32 Cpu

Pooling

Pooling layers (MaxPool1d, MaxPool2d, AvgPool2d) reduce spatial dimensions. All implement IModule. Use createDefault with a kernel size, or create for full control over stride and padding:

let pool = MaxPool2d.createDefault 2
let pool = MaxPool2d.create 3 2 1  // kernelSize=3, stride=2, padding=1

See the Pooling API reference for all variants.

Embedding

A lookup table that maps integer indices to dense vectors. Implements IModule:

let emb = Embedding.init 10000 128 F32 Cpu
let out = emb.forward indices  // indices: int64 tensor

Normalization

LayerNorm

Layer normalization. Implements IModule:

let ln = LayerNorm.initDefault 128 F32 Cpu
let out = ln.forward input

Configure with LayerNormConfig (fields: Eps, RemoveMean, Affine).

RmsNorm

Root mean square normalization. Implements IModule:

let rn = RmsNorm.init 128 1e-5 F32 Cpu
let out = rn.forward input

BatchNorm

Batch normalization. Behavior differs between training and inference:

let bn = BatchNorm.initDefault 64 F32 Cpu
let out = bn.forwardT true input   // training
let out = bn.forwardT false input  // inference

Configure with BatchNormConfig (fields: Eps, Momentum, Affine).

GroupNorm

Group normalization. Implements IModule:

let gn = GroupNorm.initDefault 8 64 F32 Cpu  // 8 groups, 64 channels
let out = gn.forward input

Dropout

Randomly zeroes elements during training:

let drop = Dropout.create 0.5

Activation

Activation is a discriminated union (Relu, Gelu, Silu, Tanh, Sigmoid, LeakyRelu, Elu, Mish) that implements IModule. Use them in a sequential model or call forward directly:

let out = Relu.forward input

See the Activation API reference for all cases.

Func and Identity

Wrap any function as a module:

let flatten = Func.create _.flatten(1, -1)

Func.Identity is a module value that returns its input unchanged:

let skip = Func.Identity

Recurrent Layers

Recurrent layers expose two primitive members:

  • zeroState: int -> 'State -- creates the initial state for a given batch size.
  • step: Tensor -> 'State -> 'State -- applies one time step.

Use RNN.scan to fold step over a full sequence:

let states = RNN.scan lstm.zeroState lstm.step inputSequence
let output = Tensor.stack (states |> List.map _.H, 1)

LSTM

Long Short-Term Memory. State type: LSTMState = { H: Tensor; C: Tensor }:

let lstm = LSTM.initDefault 128 256 F32 Cpu
let state0 = lstm.zeroState batchSize
let state1 = lstm.step input state0                      // one step
let states = RNN.scan lstm.zeroState lstm.step sequence   // full sequence
let output = Tensor.stack (states |> List.map _.H, 1)    // [batch; seqLen; hidden]

Configure with LSTMConfig for custom weight initialization.

GRU

Gated Recurrent Unit. State type: GRUState = { H: Tensor }:

let gru = GRU.initDefault 128 256 F32 Cpu
let state0 = gru.zeroState batchSize
let states = RNN.scan gru.zeroState gru.step sequence

Attention and Transformer

MultiHeadAttention

Multi-head self-attention:

let attn = MultiHeadAttention.init 128 4 F32 Cpu  // dim=128, 4 heads
let out = attn.forward (input, ?mask = mask)

Supports an optional KvCache for autoregressive generation.

TransformerBlock

A pre-norm transformer block with multi-head attention and a feed-forward network:

let block = TransformerBlock.init 128 4 512 F32 Cpu  // dim=128, 4 heads, ffDim=512
let out = block.forward (input, ?mask = mask)

KvCache

Key-value cache for autoregressive inference. KvCache implements IDisposable and frees its cached tensors on disposal:

use cache = KvCache.create 2  // cache dim (typically the sequence dimension)
let k, v = cache.append (newK, newV)
cache.reset ()  // disposes cached tensors and resets the sequence length

Model Composition

Sequential

Use the sequential { } computation expression to compose IModule layers:

let l1 = Linear.init 784 256 F32 Cpu
let l2 = Linear.init 256 10 F32 Cpu

let model = sequential {
    l1
    Relu
    l2
}

let output = model.forward input

pipeline

Use pipeline { } when you need to compose layers and functions with heterogeneous signatures (e.g. train-aware layers, custom functions):

let l1 = Linear.init 784 256 F32 Cpu
let drop = Dropout.create 0.5
let l2 = Linear.init 256 10 F32 Cpu

let forward train = pipeline {
    l1
    _.relu ()
    drop.forwardT train
    l2
}

let trainOut = forward true input
let evalOut = forward false input

Weight Initialization

The Init discriminated union controls how layer weights are initialized. Most layers default to KaimingNormal. Other cases include Const, Randn, and Uniform. Use Init.toTensor to create a tensor, or Init.toParam to create a tensor with requiresGrad:

let w = Init.toParam [ 256; 128 ] F32 Cpu KaimingNormal

See the Init API reference for all cases.

Parameter Collection

Use Model.trainableVars to collect all trainable parameters from a model record:

type MyModel = {
    Layer1: Linear
    Layer2: Linear
}

let vars = Model.trainableVars myModel
let opt = AdamW.createWithLr 0.001 vars

Model.trainableVars uses reflection to find all Tensor values in the record. It traverses nested records and lists (including Sequential).

Model.namedParams returns all parameters with their names:

let namedVars = Model.namedParams myModel
// [("Layer1.Weight", tensor); ("Layer1.Bias", tensor); ...]

Model Persistence

Save and load model parameters as SafeTensors:

Model.save myModel "checkpoints/epoch10.safetensors"
let report = Model.loadInto myModel "checkpoints/epoch10.safetensors" Strict

Model.save writes all parameters to a single .safetensors file. Model.loadInto loads only the required parameters into memory, validates shape and dtype, and returns a LoadReport with Loaded, Missing, Unexpected, ShapeMismatches, and DTypeMismatches. Use Strict to fail on mismatches, or Lenient to allow partial loading.