Docs 2.0.0 · 11 min read

use layers — reference

use layers builds neural networks of any architecture out of small, composable layers: plain stacks, residual networks, any directed acyclic graph (several inputs and outputs, shared weights), convolutional and recurrent networks, transformers and graph neural networks. Models are ordinary Apsi values, so grad, value_and_grad, the optimizers of use nn, save_params / load_params and fit all work on them. Every model runs on NumPy (exact float64, always available) or on PyTorch (CPU, CUDA, ROCm, Apple MPS, Intel XPU) — see BACKENDS.md.

• 1. The model contract
• 2. Parameters
• 3. Layers
• 4. Composition: sequential, residual, parallel, dag
• 5. Graph neural networks
• 6. The model zoo
• 7. Training: fit, predict
• 8. Inspection, saving and loading
• 9. Writing your own layer
• 10. Recipes
• 11. Errors and how to fix them

1. The model contract

Every layer and every model is a value with two operations:

use layers
use nn

let model = sequential([linear(4, 32), activation("gelu"), linear(32, 3)])
var params = model.init(seed=0)
let logits = model(params, randn(8, 4))            # shape (8, 3)

Because the parameters are kept outside the model, a model value can be shared freely (between threads, between training runs, between a server and a trainer) and several parameter sets can be used with the same model.

2. Parameters

init returns a tree whose leaves are tensors:

• a layer with weights returns a record, e.g. linear(4, 3).init() → record(W=Tensor[4, 3], b=Tensor[3])
• sequential and parallel return a list with one entry per layer
• dag returns a record with one field per node name
• layers without weights (activations, pooling, dropout, …) return an empty record()

Read and change parameters like any other value: params[0].W, params.encoder[2].b. Every function that works on trees — tree_map, grad, optimizers, count_params, save_params — works on these.

Initialisation: linear uses He initialisation (init="glorot" for Glorot/Xavier), convolutions He, attention and GNN weights Glorot, embeddings N(0, 0.02²), LSTM/GRU uniform ±1/√hidden with the LSTM forget-gate bias set to 1.

3. Layers

Shapes: B batch, T sequence length, D features, C channels, H × W image size, N graph nodes.

Dense and embeddings

Convolution and pooling (images [B, C, H, W])

Sequences

Attention and transformers

Normalisation, regularisation, scaling

Activations

activation(name) with name one of relu, gelu (tanh approximation), silu / swish, mish, elu, leaky_relu (slope 0.01), tanh, sigmoid, softplus, softmax (last axis), linear / identity — or activation(any_function), e.g. activation(fn(x) => x * sigmoid(1.702 * x)). Inside sequential([...]) a name can be written directly: sequential([linear(2, 8), "relu", linear(8, 1)]).

Small helpers

4. Composition

sequential(layers)

Applies the layers in order. Parameters: a list.

let net = sequential([conv2d_layer(1, 16), "relu", max_pool(2), flatten_layer(), linear(16 * 14 * 14, 10)])

residual(inner, pre_norm=none)

x + inner(x); with pre_norm=layer_norm_layer(d) it is x + inner(norm(x)). Parameters record(inner, norm).

parallel(branches, combine="concat")

Runs every branch on the same input and combines the results: "concat" (last axis), "add", "mul", "mean" or "list" (keep them separate). Parameters: a list, one entry per branch.

dag(nodes, inputs=["x"], output=…) — any directed acyclic graph

Each node is [name, layer, input] or [name, layer, [input1, input2, …]]; a node with several inputs receives them as a list. Nodes may only use inputs or nodes defined before them, and every name is unique (both are checked when the model is created, with a clear message). output is a name or a list of names — the model then returns a list.

# two inputs, a shared trunk, two heads
let net = dag([
    ["img",   cnn(1, 64, image=28), "image"],
    ["tab",   mlp([12, 64, 64]), "table"],
    ["both",  concatenate(), ["img", "tab"]],
    ["trunk", mlp([128, 128, 64]), "both"],
    ["cls",   linear(64, 10), "trunk"],
    ["reg",   linear(64, 1), "trunk"],
], inputs=["image", "table"], output=["cls", "reg"])

let p = net.init(seed=0)
let outs = net(p, [images, tables])          # [class logits, regression output]

Weight sharing: siamese(encoder, combine, head) (section 6) applies one encoder — one parameter set — to two inputs [a, b]; use it as a node of a larger dag when you need shared weights inside your own graph.

5. Graph neural networks

The graph

let g = graph(edges=[[0, 1], [1, 2], [2, 0]], n=3, undirected=true)      # or graph(src=…, dst=…, n=…)

A graph has .n, .edges, .src, .dst, .in_degree, .out_degree and .adjacency().

Graph layers (node features [N, D]; apply with model(p, x, graph=g))

Differentiable graph operations

segment_sum(data, ids, n), segment_mean, segment_max, segment_softmax(scores, ids, n) (softmax within each segment, e.g. over the incoming edges of every node), segment_count(ids, n), gather(x, ids). With these you can write any message-passing layer yourself (see section 9).

Node classification with a training mask

let r = fit(gat_net(8, 16, classes, heads=4), …, features, labels, loss="cross_entropy",
            graph=g, mask=train_mask, val=…, epochs=200)

The whole graph is used in every step; only nodes where mask is true count in the loss.

6. The model zoo

Every zoo model is an ordinary composition of the layers above, so summary(model) shows exactly what is inside and you can copy and change it.

7. Training

fit

let r = fit(model, params, X, y, loss="cross_entropy", opt=adam(lr=1e-3), epochs=20, batch_size=64)
print(r.history[-1])            # loss of the last epoch
let trained = r.params

fit returns record(params, history, val_history, device) — device is "numpy" or the torch device used.

Your own training loop

fit is a convenience. Any loop works, because the model is a function of its parameters:

let opt = adamw(lr=3e-4, weight_decay=0.01)
var p = model.init(seed=0)
for step in 0..1000:
    let xb, yb = next_batch()
    let l, g = value_and_grad(fn(q) => cross_entropy(model(q, xb, train=true), yb))(p)
    p = opt.step(p, clip_grad_norm(g, 1.0))

predict

predict(model, params, X, graph=…, backend="numpy") applies the model with train=false.

8. Inspection, saving and loading

• summary(model, params) prints a table of every layer and its parameter count, and returns it as a string
• count_params(params) — the number of trainable numbers
• save_params("model.npz", params) / load_params("model.npz", model.init()) — portable NumPy archives
• quantize_model(params, int8) and qlinear shrink a trained model (see AI.md)
• on the torch backend, torch_module(model, params) gives a real torch.nn.Module (state_dict, ONNX export, …)

9. Writing your own layer

Any function fn(params, x) already is a model you can differentiate. To make it a reusable layer with its own parameters, put it in a dag/sequential through lambda_layer (no parameters), or keep the parameters next to the model in a record:

use layers
use nn

# a gated message-passing layer written with the segment operations
fn gated_mp(p, x, g):
    let msg = gather(x, g.src) @ p.W
    let gate = sigmoid(gather(x, g.dst) @ p.G)
    return x + segment_mean(msg * gate, g.dst, g.n)

let p = record(W=randn(8, 8, seed=0) * 0.1, G=randn(8, 8, seed=1) * 0.1)
let l, grads = value_and_grad(fn(q) => mean(gated_mp(q, features, g) ** 2))(p)

Mix both styles freely: record(backbone=net.init(seed=0), head=dense(64, 1)) is a valid parameter tree, and fn model(p, x) = net(p.backbone, x) @ p.head.W + p.head.b is a valid model.

10. Recipes

Image classifier

let net = cnn(1, 10, image=28, channels=[32, 64], dense_width=128)
let r = fit(net, net.init(seed=0), images, labels, loss="cross_entropy", opt=adam(lr=1e-3), epochs=5, batch_size=128)

Character-level GPT

let lm = transformer_lm(vocab_size, d=128, heads=4, layers=4, context=128, dropout=0.1)
let r = fit(lm, lm.init(seed=0), tokens[:, 0:-1], tokens[:, 1:], loss="cross_entropy", epochs=10, batch_size=32,
            schedule=fn(s) => cosine_lr(s, 5000, base=3e-4, warmup=200), clip=1.0)

Sequence regression with an LSTM

let net = sequential([lstm_layer(3, 64), lstm_layer(64, 64, return_sequences=false), linear(64, 1)])

Wide & deep

let net = dag([["wide", linear(20, 1), "x"], ["deep", mlp([20, 128, 64, 1]), "x"], ["out", add(), ["wide", "deep"]]])

Variational autoencoder

let v = vae(784, 16, hidden=[256])
let r = fit(v, v.init(seed=0), images_flat, loss="vae", epochs=20, batch_size=128)
let recon = predict(v, r.params, images_flat)[0]

Graph classification — many small graphs in one batch:

let g = graph(src=all_src, dst=all_dst, n=total_nodes, batch=graph_id_of_node)
let net = gin_net(n_features, 64, classes, layers=3)
let r = fit(net, net.init(seed=0), node_features, graph_labels, loss="cross_entropy", graph=g, epochs=100)

11. Errors and how to fix them

Found a mistake? Report it · Markdown source