1. nanograd
    1. Description
    2. Requirements
    3. Modules
      1. nanograd-autograd
        1. Strict and Lazy Tensors
        2. Tensor Constructors
        3. Tensor Predicates
        4. Tensor Accessors
        5. Arithmetic Operations
        6. Linear Algebra Operations
        7. Reduction Operations
        8. Tensor Manipulation Operations
        9. Activation Functions
        10. Loss Functions
        11. Normalization Operations
        12. Convolution Operations
        13. Gradient Operations
        14. Utility Functions
      2. nanograd-layer
        1. Layer Predicates
        2. Dense Layer
        3. Convolutional Layer
        4. Batch Normalization Layer
        5. Layer Normalization
        6. Flatten Layer
        7. Pooling
        8. Sequential Container
        9. Layer Operations
        10. Activation Objects
        11. Checkpoints
        12. Utility Functions
      3. nanograd-optimizer
        1. Optimizer Predicates
        2. SGD Optimizer
        3. Adam Optimizer
        4. RMSprop Optimizer
        5. Optimizer Operations
        6. Optimizer Kernels
      4. nanograd-training
        1. Learning Rate Schedules
        2. Gradient Clipping
      5. nanograd-diagnostics
        1. Gradient Statistics
        2. Health Checks
        3. Gradient Monitor
        4. Remediation
      6. nanograd-array-morphisms
        1. Lazy Tensors
        2. Loss Functions
        3. AM-backed Layers
        4. Compiled Training and Inference
        5. Layer-by-layer Training with Contexts
        6. Kernel Backends
    4. Examples
      1. Batch Processing with Dense Layers
      2. Training an MLP
      3. Convolutional Network with Batch Normalization
      4. Saving and Restoring a Model
      5. Compiled Training of a CNN
    5. Performance Notes
    6. Limitations
    7. Troubleshooting
    8. Author
    9. Repository
    10. Version History
    11. See Also
    12. License
    13. References

nanograd

A lightweight automatic differentiation and neural network framework for CHICKEN Scheme, with BLAS-accelerated operations, batch processing, YASOS-based object abstractions, and an optional lazy-tensor backend built on array-morphisms that compiles a whole training step into a replayable program.

Description

nanograd provides a framework for building and training neural networks with automatic differentiation. It features:

Requirements

The unit tests also need test.

Modules
nanograd-autograd
Tensors, operations and the automatic differentiation engine.
nanograd-layer
Layers, containers, activation objects and checkpoints.
nanograd-optimizer
SGD, Adam and RMSprop.
nanograd-training
Learning rate schedules and gradient clipping.
nanograd-diagnostics
Gradient statistics, health checks and monitoring.
nanograd-array-morphisms
Lazy tensors, AM-backed layers and compiled training steps.

nanograd-autograd

Core automatic differentiation engine with tensor operations and batch support.

Strict and Lazy Tensors

nanograd has two kinds of tensor. Both answer the same tensor operations (tensor-data, tensor-shape, backward!, and so on), so layers and optimizers work with either.

Strict tensors
Created by make-tensor32 and make-tensor64. Their data is an SRFI-4 vector, and every operation computes its result at once.
Lazy tensors
Created by the nanograd-array-morphisms module. They wrap an array-morphisms variable, and their values are computed only when needed. See the nanograd-array-morphisms module below.
[procedure] (strict-tensor? obj) -> boolean
[procedure] (lazy-tensor? obj) -> boolean

Return #t if obj is a strict or a lazy tensor, respectively.

[procedure] (tensor-force tensor) -> tensor

Returns a strict tensor holding the values of tensor. A strict tensor is returned unchanged.

[procedure] (tensor-morphism-value tensor) -> morphism or #f

Returns the array morphism behind a lazy tensor, or #f for a strict tensor.

[procedure] (lazy-tensor-mv tensor) -> morph-variable

Returns the array-morphisms variable of a lazy tensor.

Tensor Constructors
[procedure] (make-tensor32 data shape #!key (requires-grad? #t)) -> tensor

Creates a 32-bit floating-point tensor.

data
f32vector holding the elements in row-major order
shape
list of dimensions, e.g. '(2 3) for a 2x3 matrix or '(10 2 3) for a batch of 10 matrices
requires-grad?
whether to track gradients (default #t)
; Single vector
(define x (make-tensor32 (f32vector 1.0 2.0 3.0) '(3) requires-grad?: #t))

; Batch of vectors
(define batch (make-tensor32 (make-f32vector 60) '(10 6) requires-grad?: #t))
[procedure] (make-tensor64 data shape #!key (requires-grad? #t)) -> tensor

Creates a 64-bit floating-point tensor from an f64vector.

[procedure] (make-base-tensor data shape dtype requires-grad?) -> tensor

Creates a strict tensor of dtype 'f32 or 'f64. This is the constructor behind make-tensor32 and make-tensor64.

Tensor Predicates
[procedure] (tensor? obj) -> boolean
[procedure] (tensor32? obj) -> boolean
[procedure] (tensor64? obj) -> boolean

Type predicates for tensors.

Tensor Accessors
[procedure] (tensor-data tensor) -> vector

Returns the f32vector or f64vector holding the tensor's data. For a lazy tensor the values are computed if needed.

[procedure] (tensor-grad tensor) -> vector or #f

Returns the gradient vector, or #f if there is no gradient.

[procedure] (tensor-shape tensor) -> list

Returns the shape as a list of dimensions.

[procedure] (tensor-dtype tensor) -> symbol

Returns the data type: 'f32 or 'f64.

[procedure] (tensor-requires-grad? tensor) -> boolean

Returns #t if the tensor tracks gradients.

[procedure] (tensor-children tensor) -> list

Returns the input tensors that tensor was computed from.

[procedure] (tensor-backward-fn tensor) -> procedure or #f

Returns the procedure that sends the gradient of tensor to its inputs, or #f for a leaf tensor.

Arithmetic Operations
[procedure] (add a b) -> tensor

Element-wise addition of tensors a and b. Both tensors must have the same shape and dtype.

(define z (add x y))  ; z = x + y

Gradient: dL/da = dL/dz, dL/db = dL/dz

[procedure] (sub a b) -> tensor

Element-wise subtraction: a - b.

Gradient: dL/da = dL/dz, dL/db = -dL/dz

[procedure] (mul a b) -> tensor

Element-wise multiplication (Hadamard product).

Gradient: dL/da = dL/dz * b, dL/db = dL/dz * a

[procedure] (div a b) -> tensor

Element-wise division: a / b.

Gradient: dL/da = dL/dz / b, dL/db = -dL/dz * a / b^2

[procedure] (safe-div a b #!key (epsilon 1e-8)) -> tensor

Element-wise division a / (b + epsilon), which avoids division by zero.

Linear Algebra Operations
[procedure] (matmul-op a b) -> tensor

Matrix multiplication using BLAS GEMM/GEMV operations with batch support. Supports:

; Matrix-vector multiplication
(define A (make-tensor32 (f32vector 1.0 2.0 3.0 4.0) '(2 2)))
(define b (make-tensor32 (f32vector 5.0 6.0) '(2)))
(define c (matmul-op A b))  ; 2x2 matrix times 2-vector = 2-vector

; Batch matrix multiplication
(define batch-A (make-tensor32 (make-f32vector 80) '(10 2 4)))  ; 10 samples
(define W (make-tensor32 (make-f32vector 12) '(4 3)))
(define batch-result (matmul-op batch-A W))  ; Shape: (10, 2, 3)

Gradient: dL/dA = dL/dC . B^T, dL/dB = A^T . dL/dC

[procedure] (dot-op a b) -> tensor

Dot product of two 1-D vectors using BLAS DOT.

(define result (dot-op x y))  ; scalar result

Gradient: dL/da = dL/dresult * b, dL/db = dL/dresult * a

[procedure] (scale-op tensor scalar) -> tensor

Scalar multiplication using BLAS SCAL.

Gradient: dL/dtensor = scalar * dL/dresult

Reduction Operations
[procedure] (reduce-tensor tensor reducer #!key (compute-gradient #f)) -> tensor

Generic reduction that keeps gradient flow. The reducer procedure is applied to each element in the forward pass. The optional compute-gradient procedure says how the gradient is distributed in the backward pass.

tensor
input tensor to reduce
reducer
procedure (element accumulator) -> new-accumulator
compute-gradient
optional procedure (grad-out index value all-values) -> grad-in. Without it, the gradient is distributed uniformly, as for a sum.

Returns a scalar tensor with the reduced value.

;; Sum all elements (uniform gradient distribution)
(define total (reduce-tensor x +))

;; Product of all elements (gradient uses the product rule)
(define prod (reduce-tensor x *
  compute-gradient: (lambda (grad-out idx val all-values)
                     ;; d(prod)/dx_i = prod / x_i
                     (let ((prod (fold * 1.0 all-values)))
                       (if (> val 0.0)
                           (* grad-out (/ prod val))
                           0.0)))))
[procedure] (sum-tensor tensor) -> tensor

Sums all elements. The gradient is distributed uniformly to all elements.

[procedure] (product-tensor tensor) -> tensor

Computes the product of all elements. The gradient uses the product rule: d(prod)/dx_i = prod / x_i.

[procedure] (mean-tensor tensor) -> tensor

Computes the mean of all elements.

[procedure] (compensated-sum dtype vec start end) -> number

Sums the elements start to end - 1 of the SRFI-4 vector vec with Kahan compensated summation, which reduces rounding error.

Tensor Manipulation Operations
[procedure] (slice-tensor tensor start length) -> tensor

Extracts a slice of a tensor along the first dimension. Gradients flow back to the matching positions of the original tensor.

tensor
input tensor with shape (n, ...)
start
starting index (0-based)
length
number of elements to extract

Returns a tensor with shape (length, ...).

;; Slice a batch of data
(define batch-data (make-tensor32 (make-f32vector 100) '(10 10)))
(define mini-batch (slice-tensor batch-data 2 5))  ; Shape: (5, 10)

;; Gradients flow back to the original positions
(backward! (sum-tensor mini-batch))
(tensor-grad batch-data)  ; Only rows 2-6 have non-zero gradients
[procedure] (reshape tensor new-shape) -> tensor

Reshapes the tensor. The total number of elements must not change.

[procedure] (squeeze tensor #!key (axis #f)) -> tensor

Removes dimensions of size 1. With axis, removes only that dimension, which must have size 1; negative values count from the end. If every dimension is removed, the result has shape (1).

[procedure] (transpose-tensor tensor axes) -> tensor

Permutes the dimensions of tensor. axes lists the source dimension of each result dimension; for example '(0 2 3 1) turns (N, C, H, W) into (N, H, W, C). The gradient is sent back through the inverse permutation.

[procedure] (flatten-tensor tensor) -> tensor

Flattens a multi-dimensional tensor to 1-D.

Activation Functions
[procedure] (relu tensor) -> tensor

Rectified Linear Unit: max(0, x).

Gradient: 1 if x > 0, else 0

[procedure] (tanh-op tensor) -> tensor

Hyperbolic tangent activation.

Gradient: 1 - tanh^2(x)

[procedure] (sigmoid tensor) -> tensor

Sigmoid (logistic) activation: s(x) = 1 / (1 + e^(-x)).

Gradient: s(x) * (1 - s(x))

[procedure] (sigmoid-stable tensor) -> tensor

Sigmoid computed so that it does not overflow for large negative inputs.

[procedure] (softmax x #!key (axis -1)) -> tensor

Numerically stable softmax with batch support.

Input shapes:

; Single sample
(define logits (make-tensor32 (f32vector 1.0 2.0 3.0) '(3)))
(define probs (softmax logits))  ; Sums to 1.0

; Batch of samples
(define batch-logits (make-tensor32 (make-f32vector 60) '(20 3)))
(define batch-probs (softmax batch-logits axis: -1))  ; Each row sums to 1.0

Gradient: dL/dx = softmax(x) * (dL/dy - sum(dL/dy * softmax(x)))

[procedure] (log-softmax x #!key (axis -1)) -> tensor

Log-softmax with batch support. It is more numerically stable than log(softmax(x)).

Input shapes:

Gradient: dL/dx = dL/dy - exp(log_softmax(x)) * sum(dL/dy)

[procedure] (leaky-relu tensor #!key (alpha 0.01)) -> tensor

Leaky ReLU: max(alpha * x, x).

[procedure] (softplus tensor #!key (beta 1.0)) -> tensor

Softplus activation: log(1 + e^(beta * x)) / beta.

[procedure] (gelu tensor) -> tensor

Gaussian Error Linear Unit, using the tanh approximation.

[procedure] (silu tensor) -> tensor

SiLU (Sigmoid Linear Unit), also known as Swish: x * s(x).

Loss Functions

The loss functions scale their gradient by the incoming gradient, so a loss can be scaled or added to other losses before backward!.

[procedure] (mse-loss pred target #!key (reduction 'mean)) -> tensor

Mean squared error loss with batch support.

pred
predictions tensor (any shape)
target
target tensor (same shape as pred)
reduction
'mean (average over all elements) or 'sum
; Single sample
(define loss (mse-loss predictions targets))

; Batch of samples
(define batch-pred (make-tensor32 pred-data '(32 10)))
(define batch-target (make-tensor32 target-data '(32 10)))
(define batch-loss (mse-loss batch-pred batch-target reduction: 'mean))
[procedure] (cross-entropy-loss pred target #!key (reduction 'mean) (from-logits #f)) -> tensor

Cross-entropy loss with batch support.

pred
predictions: probabilities (softmax already applied) when from-logits is #f, or raw scores when it is #t
target
one-hot targets of the same shape as pred, or class indices of shape (batch_size)
reduction
'mean (average over the batch) or 'sum
from-logits
if true, apply log-softmax to pred first

Input shapes:

; Single sample with one-hot target
(define loss (cross-entropy-loss probs target))

; Batch with one-hot targets
(define batch-probs (softmax logits axis: -1))
(define batch-loss (cross-entropy-loss batch-probs targets reduction: 'mean))

; Batch with class indices
(define class-indices (make-tensor32 (f32vector 0.0 2.0 1.0) '(3)))
(define batch-loss (cross-entropy-loss logits class-indices
                                       from-logits: #t
                                       reduction: 'mean))
Normalization Operations
[procedure] (rmsnorm x weight #!key (epsilon 1e-5)) -> tensor

Root Mean Square layer normalization with batch support.

Input shapes:

Formula: output[i] = (x[i] / RMS(x)) * weight[i], where RMS(x) = sqrt(mean(x^2) + epsilon)

; Single vector
(define x (make-tensor32 (make-f32vector 512) '(512)))
(define gamma (make-tensor32 (make-f32vector 512 1.0) '(512)))
(define normalized (rmsnorm x gamma))

; Batch of vectors
(define batch-x (make-tensor32 (make-f32vector (* 32 512)) '(32 512)))
(define batch-norm (rmsnorm batch-x gamma))  ; Normalized per row
[procedure] (l2-normalize tensor #!key (axis #f) (epsilon 1e-8)) -> tensor

L2 normalization.

axis
#f to normalize the entire tensor, or an axis to normalize along

For 2-D tensors:

; Normalize the entire tensor
(define normalized (l2-normalize x))

; Normalize each row of a batch
(define batch (make-tensor32 (make-f32vector 200) '(10 20)))
(define row-normalized (l2-normalize batch axis: 1))  ; Each row has L2 norm 1
[procedure] (cosine-similarity a b) -> tensor

Cosine similarity: (a . b) / (||a|| * ||b||).

Convolution Operations
[procedure] (conv2d input weight bias #!key (stride 1) (padding 0)) -> tensor

2-D convolution using im2col and GEMM, with batch support.

input
tensor of shape (C_in, H, W) or (N, C_in, H, W)
weight
tensor of shape (C_out, C_in, KH, KW)
bias
tensor of shape (C_out), or #f
stride
convolution stride (default 1)
padding
zero padding (default 0)

Output shapes:

; Single image
(define img (make-tensor32 (make-f32vector (* 3 32 32)) '(3 32 32)))
(define output (conv2d img weights bias stride: 2 padding: 1))

; Batch of images
(define batch-imgs (make-tensor32 (make-f32vector (* 16 3 32 32)) '(16 3 32 32)))
(define batch-output (conv2d batch-imgs weights bias))  ; Shape: (16, C_out, H_out, W_out)
Gradient Operations
[procedure] (zero-grad! tensor) -> void

Sets all gradient values to zero.

[procedure] (backward! tensor) -> void

Computes gradients by reverse-mode automatic differentiation. Sorts the computation graph topologically and runs each backward procedure in order. Signals an error if the graph has a cycle.

[procedure] (add-to-grad! tensor delta) -> void

Adds the vector delta to the tensor's gradient using BLAS AXPY.

[procedure] (set-grad! tensor grad) -> void

Replaces the tensor's gradient vector with grad.

[procedure] (set-backward-fn! tensor fn inputs) -> void

Makes tensor the result of a differentiable operation. fn is a procedure of no arguments that reads (tensor-grad tensor) and adds the gradient of each input with add-to-grad!. inputs is the list of input tensors. This is how new operations are defined:

;; y = 2x
(define (double x)
  (let ((y (make-base-tensor (f32vector-map (lambda (v) (* 2.0 v)) (tensor-data x))
                             (tensor-shape x) 'f32 #t)))
    (set-backward-fn! y
                      (lambda ()
                        (add-to-grad! x (f32vector-map (lambda (g) (* 2.0 g))
                                                       (tensor-grad y))))
                      (list x))
    y))

f32vector-map in this example stands for any element-wise map over an f32vector; it is not part of nanograd.

Utility Functions
[procedure] (tensor->list tensor) -> list

Converts tensor data to a list.

[procedure] (print-tensor tensor) -> void

Prints the tensor's shape, dtype, data and gradient.

[procedure] (vector-length-for-dtype vec dtype) -> integer

Returns the length of an f32vector ('f32) or f64vector ('f64).

The module also exports small helpers used by the other nanograd modules: f32vector-fold, f64vector-fold, and the BLAS-style vector routines ssub, dsub, sblit, dblit, scopy and dcopy.

nanograd-layer

Neural network layers and containers with batch processing support.

Layer Predicates
[procedure] (layer? obj) -> boolean
[procedure] (dense-layer? obj) -> boolean
[procedure] (conv2d-layer? obj) -> boolean
[procedure] (batch-norm-2d? obj) -> boolean
[procedure] (sequential? obj) -> boolean
[procedure] (flatten-layer? obj) -> boolean
Dense Layer
[procedure] (make-dense-layer input-size output-size #!key (activation (make-identity)) (use-bias #t) (dtype 'f32) (name "Dense")) -> layer

Creates a fully connected layer. Weights start as uniform random values scaled by sqrt(2 / input-size); biases start at zero.

input-size
number of input features
output-size
number of output features
activation
activation object (default identity)
use-bias
whether to include a bias term (default #t)
dtype
'f32 or 'f64 (default 'f32)
name
layer name

Input shapes:

For 2-D inputs, the layer uses one BLAS GEMM for the whole batch.

(define layer (make-dense-layer 784 128
                                activation: (make-relu)
                                name: "Hidden1"))

; Single input
(define x (make-tensor32 (make-f32vector 784) '(784)))
(define output (forward layer x))  ; Shape: (128)

; Batch input
(define batch-x (make-tensor32 (make-f32vector (* 32 784)) '(32 784)))
(define batch-output (forward layer batch-x))  ; Shape: (32, 128)
Convolutional Layer
[procedure] (make-conv2d-layer in-channels out-channels kernel-size #!key (stride 1) (padding 0) (activation (make-identity)) (dtype 'f32) (weight-values #f) (bias-values #f) (name "Conv2D")) -> layer

Creates a 2-D convolutional layer. Unless weight-values is given, weights start as uniform random values scaled by sqrt(2 / fan-in) (He scaling).

in-channels
number of input channels
out-channels
number of output channels
kernel-size
size of the square kernel
stride
convolution stride (default 1)
padding
zero padding (default 0)
activation
activation object
dtype
'f32 or 'f64
weight-values
optional SRFI-4 vector of initial weights, laid out as (out-channels, in-channels, kernel-size, kernel-size)
bias-values
optional SRFI-4 vector of initial biases
name
layer name

Input shapes:

Output shapes:

(define conv (make-conv2d-layer 3 32 3
                                stride: 1
                                padding: 1
                                activation: (make-relu)))

; Single image
(define img (make-tensor32 img-data '(3 32 32)))
(define features (forward conv img))  ; Shape: (32, 32, 32)

; Batch of images
(define batch (make-tensor32 batch-data '(16 3 32 32)))
(define batch-features (forward conv batch))  ; Shape: (16, 32, 32, 32)
Batch Normalization Layer
[procedure] (make-batch-norm-2d num-features #!key (epsilon 1e-5) (momentum 0.1) (dtype 'f32) (name "BatchNorm2d")) -> layer

Creates a 2-D batch normalization layer:

y = gamma * (x - mu) / sqrt(sigma^2 + epsilon) + beta

where mu and sigma^2 are the per-channel mean and variance of the batch in training mode, and running averages of them in evaluation mode.

num-features
number of channels (C)
epsilon
small constant for numerical stability (default 1e-5)
momentum
weight of the current batch in the running statistics (default 0.1)
dtype
'f32 or 'f64 (default 'f32)
name
layer name

Input shapes:

The output has the same shape as the input. gamma and beta are learned.

;; Batch norm for 64 channels
(define bn (make-batch-norm-2d 64 epsilon: 1e-5 momentum: 0.1))

;; Training mode: uses batch statistics
(set-training-mode! bn #t)
(define normalized (forward bn input))  ; Input shape: (N, 64, H, W)

;; Evaluation mode: uses running statistics
(set-eval-mode! bn)
(define test-normalized (forward bn test-input))
Layer Normalization
[procedure] (layer-norm x gamma beta #!key (epsilon 1e-5)) -> tensor

Normalizes each row of x over its last axis. Each row x_r of length D is mapped to

y_r = gamma * (x_r - mean(x_r)) / sqrt(var(x_r) + epsilon) + beta

where var is the biased variance. gamma and beta are tensors of shape (D). Gradients flow to x, gamma and beta.

Flatten Layer
[procedure] (make-flatten #!key (name "Flatten")) -> layer

Creates a layer that flattens each sample to a vector.

(define flatten (make-flatten name: "Flatten"))

; Flatten a batch of feature maps
(define features (make-tensor32 data '(32 64 8 8)))
(define flattened (forward flatten features))  ; Shape: (32, 4096)
Pooling
[procedure] (maxpool2d input kernel-size #!key (stride #f)) -> tensor

2-D max pooling with a square window of size kernel-size. The stride defaults to the kernel size, so windows do not overlap.

The gradient goes to the largest element of each window.

[procedure] (global-avg-pool2d input) -> tensor

Averages over the spatial dimensions.

The gradient is distributed uniformly over the spatial positions of each channel.

;; Batch of feature maps
(define batch-features (make-tensor32 (make-f32vector (* 32 128 8 8)) '(32 128 8 8)))
(define batch-pooled (global-avg-pool2d batch-features))  ; Shape: (32, 128)

;; Classification head
(define logits (forward fc-layer batch-pooled))  ; Shape: (32, num_classes)

Global average pooling can replace a large fully connected layer at the end of a CNN, with far fewer parameters.

Sequential Container
[procedure] (make-sequential layers #!key (name "Sequential")) -> layer

Creates a container that applies layers in order.

(define model
  (make-sequential
   (list
    (make-dense-layer 784 128 activation: (make-relu))
    (make-dense-layer 128 64 activation: (make-relu))
    (make-dense-layer 64 10 activation: (make-identity)))
   name: "MLP"))

; Works with both single and batch inputs
(define single-output (forward model single-input))
(define batch-output (forward model batch-input))
Layer Operations
[procedure] (forward layer input) -> tensor

Runs the layer on input, which may be a single sample or a batch.

[procedure] (parameters layer) -> list

Returns the list of trainable parameter tensors.

[procedure] (zero-grad-layer! layer) -> void

Zeros the gradients of all parameters of the layer.

[procedure] (set-training-mode! layer training?) -> void

Puts the layer, and every layer inside a container, in training mode (training? is #t) or evaluation mode (#f). In training mode batch normalization uses batch statistics and updates its running statistics; in evaluation mode it uses the running statistics and leaves them unchanged.

[procedure] (set-eval-mode! layer) -> void

Same as (set-training-mode! layer #f).

[procedure] (layer-input-size layer) -> integer or #f
[procedure] (layer-output-size layer) -> integer or #f
[procedure] (layer-activation layer) -> activation
[procedure] (layer-name layer) -> string

Layer properties. The sizes are #f for layers whose sizes depend on the input, such as flatten.

Activation Objects
[procedure] (make-relu) -> activation
[procedure] (make-tanh) -> activation
[procedure] (make-sigmoid) -> activation
[procedure] (make-gelu) -> activation
[procedure] (make-silu) -> activation
[procedure] (make-identity) -> activation

Create activation objects for use in layers.

[procedure] (activation? obj) -> boolean
[procedure] (activation-forward act x) -> tensor
[procedure] (activation-name act) -> string
Checkpoints

Layers and models are saved to and loaded from CBOR files.

[procedure] (save-model model filepath) -> void
[procedure] (save-layer layer filepath) -> void

Write model (a sequential container or a single layer) with all its weights to filepath. The two procedures are the same.

[procedure] (load-model filepath) -> layer
[procedure] (load-layer filepath) -> layer

Read a model written by save-model. They signal an error if the file is not a nanograd checkpoint or has an unsupported version.

[procedure] (layer->serializable layer) -> alist
[procedure] (serializable->layer alist) -> layer

Convert a layer to and from the association list that is stored in the file.

Dense, convolutional, flatten and sequential layers can be saved. The AM-backed layers of nanograd-array-morphisms cannot.

(save-model model "model.ngrd")
(define restored (load-model "model.ngrd"))

A checkpoint file holds one CBOR data item (RFC 8949), written with the cbor egg, after the CBOR self-describe tag. The item is an association list:

((format . nanograd-checkpoint)
 (version . 1)
 (model . <output of layer->serializable>))

Tensor data are stored as RFC 8746 typed arrays (tag 85 for little-endian f32, tag 86 for little-endian f64), so they are read and written as single blocks of bytes and round-trip bit for bit. Other CBOR decoders can read the files; for example, Python's cbor2 package decodes a checkpoint, and numpy.frombuffer turns each typed array into an array.

Checkpoints written by nanograd 3.0 and earlier (CHICKEN 5, s11n format) cannot be read.

Utility Functions
[procedure] (print-layer layer #!optional (indent 0)) -> void

Prints layer information with optional indentation.

[procedure] (summary model) -> void

Prints a model summary with its layers and parameters.

nanograd-optimizer

Optimization algorithms for neural network training.

Optimizer Predicates
[procedure] (optimizer? obj) -> boolean
[procedure] (sgd? obj) -> boolean
[procedure] (adam? obj) -> boolean
[procedure] (rmsprop? obj) -> boolean
SGD Optimizer
[procedure] (make-sgd parameters #!key (learning-rate 0.01) (momentum 0.0) (weight-decay 0.0) (nesterov #f)) -> optimizer

Stochastic gradient descent with optional momentum and Nesterov momentum.

parameters
list of parameter tensors to optimize
learning-rate
step size (default 0.01)
momentum
momentum factor (default 0.0, no momentum)
weight-decay
L2 regularization factor (default 0.0)
nesterov
use Nesterov momentum (default #f)
Adam Optimizer
[procedure] (make-adam parameters #!key (learning-rate 0.001) (beta1 0.9) (beta2 0.999) (epsilon 1e-8) (weight-decay 0.0)) -> optimizer

Adam (adaptive moment estimation) with bias correction.

beta1
decay rate of the first moment (default 0.9)
beta2
decay rate of the second moment (default 0.999)
epsilon
numerical stability constant (default 1e-8)
weight-decay
L2 regularization factor, added to the gradient (default 0.0)
RMSprop Optimizer
[procedure] (make-rmsprop parameters #!key (learning-rate 0.01) (alpha 0.99) (epsilon 1e-8) (weight-decay 0.0) (momentum 0.0)) -> optimizer

RMSprop with optional momentum.

alpha
smoothing constant (default 0.99)
Optimizer Operations
[procedure] (step! optimizer) -> void

Updates the parameters from their current gradients.

[procedure] (get-learning-rate optimizer) -> number

Returns the current learning rate.

[procedure] (set-learning-rate! optimizer lr) -> void

Sets the learning rate, for example from a learning rate schedule.

[procedure] (optimizer-state optimizer) -> alist

Returns the optimizer's settings as an association list (for Adam, also the step count).

Optimizer Kernels

An optimizer kernel performs an optimizer's whole update of one parameter vector in a single call, for example in compiled C code. When a kernel is registered, step! calls it instead of the Scheme update. Currently Adam uses registered kernels.

[procedure] (register-optimizer-kernel! name dtype kernel) -> void

Uses kernel for the updates of optimizer name (a symbol, such as 'adam) on parameters of dtype 'f32 or 'f64. A kernel of #f removes the registration.

An Adam kernel is a procedure (kernel n hyper param grad m v) that updates the first n elements of the SRFI-4 vectors param, m and v in place, and adds the weight-decay term to grad. hyper is an f64vector holding the learning rate, beta1, beta2, epsilon, the bias corrections 1 - beta1^t and 1 - beta2^t, and the weight decay.

[procedure] (optimizer-kernel name dtype) -> procedure or #f

Returns the kernel registered for optimizer name on dtype parameters, or #f.

nanograd-training

Learning rate schedules and gradient clipping.

Learning Rate Schedules

A schedule is a procedure that takes a step number and returns the learning rate for that step.

[procedure] (make-constant-schedule learning-rate) -> procedure

Always returns learning-rate.

[procedure] (make-linear-warmup-schedule init-lr final-lr warmup-steps) -> procedure

Rises linearly from init-lr to final-lr over warmup-steps steps, then stays at final-lr.

[procedure] (make-cosine-schedule max-lr min-lr total-steps #!key (warmup-steps 0)) -> procedure

Cosine annealing from max-lr down to min-lr over total-steps, after an optional linear warmup:

lr = min-lr + 0.5 * (max-lr - min-lr) * (1 + cos(pi * progress)), where progress = (step - warmup-steps) / (total-steps - warmup-steps)

[procedure] (make-cosine-warmup-schedule max-lr min-lr total-steps warmup-steps) -> procedure

Same as (make-cosine-schedule max-lr min-lr total-steps warmup-steps: warmup-steps).

[procedure] (make-step-decay-schedule init-lr decay-rate decay-steps #!key (warmup-steps 0)) -> procedure

Multiplies the learning rate by decay-rate every decay-steps steps, after an optional warmup:

lr = init-lr * decay-rate ^ floor(step / decay-steps)

[procedure] (learning-rate-at-step schedule step) -> number

Same as (schedule step).

(define schedule (make-cosine-schedule 1e-3 1e-5 10000 warmup-steps: 500))
(set-learning-rate! optimizer (schedule step))
Gradient Clipping
[procedure] (compute-gradient-norm parameters) -> number

Returns the global L2 norm of the gradients of parameters: the square root of the sum of all squared gradient elements.

[procedure] (clip-gradients-by-norm! parameters max-norm) -> number

If the global L2 norm is larger than max-norm, scales all gradients so that it equals max-norm, which keeps their direction. Returns the norm before clipping.

[procedure] (clip-gradients-by-value! parameters clip-value) -> void

Clamps each gradient element to the range [-clip-value, clip-value].

(backward! loss)
(clip-gradients-by-norm! (parameters model) 1.0)
(step! optimizer)

nanograd-diagnostics

Gradient statistics, health checks and monitoring over training.

Both this module and nanograd-training export compute-gradient-norm. To import both modules, rename or exclude one of them, for example:

(import nanograd-training
        (except nanograd-diagnostics compute-gradient-norm))
Gradient Statistics
[procedure] (gradient-stats parameters) -> alist

Returns an association list with the global L2 norm (global-l2-norm), the global infinity norm (global-inf-norm), and layer-stats: for each parameter with a gradient, its index (layer), l2-norm, inf-norm, mean absolute value (mean-abs) and size.

[procedure] (compute-gradient-norm parameters #!key (norm-type 'l2)) -> number

Returns the global gradient norm: L2 ('l2) or infinity ('inf).

[procedure] (gradient-l2-norm grad dtype) -> number
[procedure] (gradient-inf-norm grad dtype) -> number

Return the L2 norm or the largest absolute value of one gradient vector.

Health Checks
[procedure] (has-exploding-gradients? parameters #!key (threshold 10.0) (norm-type 'l2)) -> boolean

Returns #t if the global gradient norm is above threshold, or is NaN or infinite.

[procedure] (has-vanishing-gradients? parameters #!key (threshold 1e-7) (norm-type 'l2)) -> boolean

Returns #t if the global gradient norm is below threshold.

[procedure] (check-gradient-health parameters #!key (exploding-threshold 10.0) (vanishing-threshold 1e-7) (imbalance-threshold 1000.0) (warn? #t)) -> alist

Checks for exploding, vanishing and unbalanced gradients. Gradients are unbalanced when the largest per-parameter L2 norm is more than imbalance-threshold times the smallest one; this check is skipped when some gradient is zero. With warn?, each problem found is printed.

Returns an association list with healthy? (#t if no problem was found), issues (a list of symbols, such as gradient-imbalance) and stats (the result of gradient-stats).

[procedure] (print-gradient-report parameters #!optional (label "")) -> void

Prints the health status and gradient statistics.

Gradient Monitor
[procedure] (make-gradient-monitor #!key (exploding-threshold 10.0) (vanishing-threshold 1e-7) (imbalance-threshold 1000.0) (history-size 100)) -> monitor

Creates an object that records gradient statistics over training. It keeps the last history-size steps.

[procedure] (gradient-monitor? obj) -> boolean
[procedure] (record-step! monitor step parameters) -> alist

Records the gradient norms and health of parameters at training step step. Returns the health check result, as from check-gradient-health.

[procedure] (get-gradient-history monitor) -> list

Returns the recorded steps, oldest first. Each entry is an association list with step, global-l2-norm, global-inf-norm, healthy? and issues.

[procedure] (diagnose-training monitor) -> alist

Summarizes the last 10 recorded steps: mean, largest and smallest gradient norm, the number of unhealthy steps and warnings, a stability ratio (smallest norm divided by largest), and a list of recommendations as strings.

Remediation
[procedure] (clip-gradients! parameters #!key (max-norm 1.0) (norm-type 'l2) (verbose? #t)) -> number or #f

Scales the gradients down if their global norm is above max-norm. Returns the norm before clipping if clipping was applied, and #f otherwise.

[procedure] (suggest-learning-rate parameters current-lr gradient-norm) -> number

Suggests a new learning rate from the gradient norm: lower for large norms, higher for very small ones. Prints the suggestion.

(import (except nanograd-diagnostics compute-gradient-norm))

(define monitor (make-gradient-monitor))

;; In the training loop, after backward!
(record-step! monitor step (parameters model))

;; After training
(diagnose-training monitor)

nanograd-array-morphisms

Lazy tensors backed by array-morphisms, layers built on them, and training and inference steps that are compiled once and replayed.

A lazy tensor wraps an array-morphisms variable. Operations on lazy tensors build an expression instead of computing values; the values are computed when they are read with tensor-data, or when a step procedure below runs the computation. Because array-morphisms sees the whole computation, it can fuse operations, reuse buffers, and run matrix products and element-wise operations on registered C backends.

The module provides three ways to train:

am-training-step/ssa
Compiles the forward pass, the loss and the backward pass of the model into one SSA program on the first call, and replays it on later calls. This is the fastest way.
am-training-step
Runs the forward and backward passes layer by layer through two morphism contexts, which reuse buffers after the first step.
backward! on a lazy loss
Builds the gradients as lazy expressions, without contexts.

For inference, am-forward/ssa compiles and replays the forward pass alone.

Morphism contexts are created with make-morphism-context from the array-morphisms-context module. Each context holds one compiled program, and a compiled program is fixed to the shapes of the inputs it was compiled for. Use one context per model and input shape, for example one per batch size.

Lazy Tensors
[procedure] (make-lazy-tensor mv) -> tensor

Creates a lazy tensor wrapping the array-morphisms variable mv.

[procedure] (get-or-make-lazy mv) -> tensor

Returns the lazy tensor for mv, creating it if needed. Use this procedure to wrap input data:

(import array-morphisms-core (prefix array-morphisms-grad am:))

;; A batch of 32 grayscale 28x28 images, without copying the data
(define x-lt
  (get-or-make-lazy
   (am:make-var (make-morphism image-data '(32 1 28 28) 'f32)
                #f)))  ; no gradient for the input
[procedure] (lazy-tensor-morph-variable tensor) -> morph-variable or #f

Returns the array-morphisms variable of a lazy tensor, or #f if tensor is not lazy.

[procedure] (activation-forward-am act tensor) -> tensor

Applies the activation object act (from make-relu, make-gelu, and so on) to a lazy tensor.

Loss Functions
[procedure] (am-mse-loss pred target) -> tensor

Mean squared error of two lazy tensors. Returns a scalar lazy tensor.

[procedure] (am-cross-entropy-loss logits target) -> tensor

Numerically stable softmax cross-entropy, averaged over the batch. logits is a (batch, C) lazy tensor of raw scores and target a (batch, C) lazy tensor of one-hot labels. Returns a scalar lazy tensor.

AM-backed Layers

These layers take and return lazy tensors, and answer the same layer operations as the layers of nanograd-layer (forward, parameters, set-training-mode!, and so on). Their parameters are lazy tensors, so the optimizers of nanograd-optimizer work with them. Their weights start from He initialization. They cannot be saved with save-model.

[procedure] (make-am-dense-layer in-size out-size #!key (activation (make-identity)) (dtype 'f64)) -> layer

Fully connected layer for inputs of shape (batch, in-size).

[procedure] (make-am-conv2d-layer in-channels out-channels kernel-size #!key (stride 1) (padding 0) (activation (make-identity)) (dtype 'f64)) -> layer

2-D convolution, run as a single operation that the SSA compiler differentiates. kernel-size is a number, or a pair (KH . KW).

The input has shape (N, C, H, W), or (N, H, W, C) when its last dimension equals in-channels. The output has shape (N, OH, OW, C_out), with channels last, so that the next convolution can read it without a copy. A flatten layer after it therefore orders the features as (H, W, C), not (C, H, W) as for make-conv2d-layer.

[procedure] (make-am-batch-norm num-features #!key (epsilon 1e-5) (momentum 0.1) (dtype 'f64)) -> layer

Batch normalization over the batch axis (axis 0), with learned gamma and beta of length num-features.

[procedure] (make-am-flatten #!key (name "Flatten")) -> layer

Flattens each sample: (N, ...) -> (N, product of the other dimensions). 1-D and 2-D inputs are passed through unchanged.

[procedure] (make-am-sequential layers) -> layer

Applies layers in order.

[procedure] (am-parameters layer) -> list

Returns the parameters of layer as lazy tensors.

[procedure] (am-zero-grad! layer) -> void

Zeros the gradients of all parameters of layer.

Compiled Training and Inference
[procedure] (am-training-step/ssa ctx optimizer model loss-fn x . extra-inputs) -> tensor

Performs one training step: runs model on the lazy tensor x, computes the loss with (loss-fn output), computes the gradients, calls step! on optimizer and zeros the gradients. Returns the loss as a lazy tensor.

The first call with a context ctx compiles the forward pass, the loss and the backward pass into one SSA program and plans its buffers; later calls replay that program with the new input. Parameters updated by the optimizer are seen by the next replay automatically.

extra-inputs are further lazy tensors that change every step, such as the targets. Pass each one here and also use it inside loss-fn, so that the compiled program reads its new value on every step. The procedure signals an error if some parameter of the model gets no gradient.

[procedure] (am-forward/ssa ctx model x) -> tensor

Runs model on the lazy tensor x for inference, through a compiled forward-only program. The first call with ctx compiles the program; later calls replay it with the new input. The program reads the model's current weights on every call.

Use one context per input shape, and do not use a context for both am-forward/ssa and am-training-step/ssa. The result's data lives in the context's buffers and is overwritten by the next call with the same context, so read it with tensor-data before calling again.

[procedure] (am-replay-plan-stats ctx) -> alist

Returns an association list of (instruction-tag . count) for the replay plan compiled for ctx.

Layer-by-layer Training with Contexts
[procedure] (make-am-training-context) -> (values ctx-fwd ctx-bwd)

Returns two new morphism contexts, for the forward and the backward pass.

[procedure] (am-training-step ctx-fwd ctx-bwd optimizer model loss-fn x) -> tensor

Performs one training step through the two contexts: the forward pass layer by layer, the loss, the backward pass, the optimizer step, and zeroing of the gradients. After the first step both contexts reuse the buffers they planned. Returns the loss as a lazy tensor.

[procedure] (am-forward/ctx ctx layer x) -> tensor

Runs layer on x through ctx, computing each layer's output before the next layer runs.

[procedure] (am-backward!/ctx ctx loss) -> void

Runs the backward pass from the lazy scalar loss through ctx, computing each gradient before it is passed on.

[procedure] (am-realize-grads/ctx ctx parameters) -> void

Computes the gradients of all parameters through ctx, so that buffers can be shared among them.

[procedure] (am-finalize-context! ctx) -> void

Plans the buffers of ctx from the allocations recorded so far.

[procedure] (am-reset-context! ctx) -> void

Prepares ctx for the next step, which reuses the planned buffers.

[procedure] (fresh-copy-morphism m) -> morphism

Copies the concrete array m into a new row-major buffer that does not belong to any context.

Kernel Backends

array-morphisms can run matrix products, convolutions and element-wise operations on registered C backends; see register-blas-backend! and register-activation-backend! in array-morphisms. nanograd uses whatever backends are registered. Together with an Adam kernel registered with register-optimizer-kernel!, this moves nearly all of the work of a compiled training step out of Scheme.

Examples

Batch Processing with Dense Layers

 
(import nanograd-autograd nanograd-layer)

(define batch-size 32)
(define input-dim 784)
(define batch-data (make-f32vector (* batch-size input-dim)))

;; Fill with data...

(define batch-input (make-tensor32 batch-data (list batch-size input-dim)))

(define layer (make-dense-layer input-dim 128 activation: (make-relu)))
(define output (forward layer batch-input))  ; Shape: (32, 128)

Training an MLP

 
(import nanograd-autograd nanograd-layer nanograd-optimizer)

(define model
  (make-sequential
   (list
    (make-dense-layer 784 256 activation: (make-relu))
    (make-dense-layer 256 128 activation: (make-relu))
    (make-dense-layer 128 10 activation: (make-identity)))
   name: "MLP"))

(define optimizer (make-adam (parameters model) learning-rate: 0.001))

;; train-batches is a list of (x . y) pairs:
;; x of shape (batch_size, 784), y one-hot of shape (batch_size, 10)
(define (train-epoch train-batches)
  (set-training-mode! model #t)
  (for-each
   (lambda (batch)
     (let* ((logits (forward model (car batch)))
            (loss (cross-entropy-loss logits (cdr batch)
                                      from-logits: #t
                                      reduction: 'mean)))
       (backward! loss)
       (step! optimizer)
       (zero-grad-layer! model)))
   train-batches))

Convolutional Network with Batch Normalization

 
(import nanograd-autograd nanograd-layer nanograd-optimizer)

(define cnn
  (make-sequential
   (list
    (make-conv2d-layer 3 32 3 stride: 1 padding: 1)
    (make-batch-norm-2d 32)
    (make-conv2d-layer 32 64 3 stride: 1 padding: 1)
    (make-batch-norm-2d 64)
    (make-flatten)
    (make-dense-layer (* 64 32 32) 128 activation: (make-relu))
    (make-dense-layer 128 10))
   name: "CNN"))

(define batch-images (make-tensor32 image-data '(16 3 32 32)))  ; 16 RGB images
(set-training-mode! cnn #t)
(define predictions (forward cnn batch-images))  ; Shape: (16, 10)

Saving and Restoring a Model

 
(import nanograd-layer)

(save-model cnn "cnn.ngrd")

(define restored (load-model "cnn.ngrd"))
(set-eval-mode! restored)

Compiled Training of a CNN

 
(import srfi-4
        array-morphisms-core array-morphisms-context
        (prefix array-morphisms-grad am:)
        nanograd-autograd nanograd-layer nanograd-optimizer
        nanograd-array-morphisms)

(define num-classes 4)

(define model
  (make-am-sequential
   (list
    (make-am-conv2d-layer 1 16 3 padding: 1 activation: (make-relu) dtype: 'f32)
    (make-am-conv2d-layer 16 32 3 stride: 2 padding: 1 activation: (make-relu) dtype: 'f32)
    (make-am-flatten)
    (make-am-dense-layer (* 14 14 32) 64 activation: (make-relu) dtype: 'f32)
    (make-am-dense-layer 64 num-classes dtype: 'f32))))

(define optimizer (make-adam (am-parameters model) learning-rate: 0.001))

;; Wraps an f32vector as a lazy input tensor, without copying
(define (lazy-input data shape)
  (get-or-make-lazy (am:make-var (make-morphism data shape 'f32) #f)))

;; One context for training at batch size 32
(define train-ctx (make-morphism-context))

;; images: f32vector of 32*1*28*28 values; labels: f32vector of 32*4 one-hot values
(define (train-batch images labels)
  (let* ((x (lazy-input images '(32 1 28 28)))
         (y (lazy-input labels (list 32 num-classes)))
         (loss (am-training-step/ssa train-ctx optimizer model
                                     (lambda (logits) (am-cross-entropy-loss logits y))
                                     x y)))
    (f32vector-ref (tensor-data loss) 0)))

;; A separate context for inference at batch size 32
(define eval-ctx (make-morphism-context))

(define (predict images)
  (tensor-data (am-forward/ssa eval-ctx model (lazy-input images '(32 1 28 28)))))

Performance Notes

Limitations

Troubleshooting

Shape mismatch errors

Make sure tensor shapes match. For batched operations, the batch dimensions must agree.

(define x (make-tensor32 (make-f32vector 200) '(10 20)))
(define y (make-tensor32 (make-f32vector 300) '(15 20)))
(add x y)  ; Error: shape mismatch

A compiled step rejects a batch

am-forward/ssa signals an error when the input shape differs from the shape the context was compiled for, as happens with a smaller last batch. Use a separate context for each batch size.

Batch normalization mode not set

Set the mode before training and before evaluation:

(set-training-mode! model #t)
(train-epoch model)

(set-eval-mode! model)
(evaluate model)

Author

Ivan Raikov

Repository

https://github.com/iraikov/nanograd

Version History
3.1
Support for CHICKEN 6
3.0
array-morphisms integration
2.3
Iris dataset example
2.2
Module nanograd-training: learning rate schedules and gradient clipping
2.1
Unit tests use the test egg
2.0
Batch processing support
1.2
Additional operations
1.1
Correction in the mul operation
1.0
Initial release

See Also

License

LGPL-3

References