Standard library · Machine learning

autograd

Reverse mode automatic differentiation for training models.

import autograd

On this page

Types
Functions

Training a model means lowering the loss with gradients. This package records every computation on a Tape, and Backward works out the gradient of the loss for every input. The functions return a node index rather than a tensor; read the value with Value and the gradient with Grad.

Call reset on the tape at every training step. The ready made layers are in nn and the optimisers in optim.

Types

`Node`

struct Node

One node on the tape.

Name	Type	Description
`op`	`int = 0`	The operation.
`a`	`int = -1`	The first input.
`b`	`int = -1`	The second input.
`out`	`int = -1`	The output.
`aux`	`int = -1`	The auxiliary data.
`s0`	`int = 0`	The first integer parameter.
`s1`	`int = 0`	The second integer parameter.
`f0`	`f64 = 0`	The float parameter.

autograd-Basics.salam

import autograd
import tensor

func rounded(t: tensor.Tensor): str:
    mut text := ""
    repeat t.numel() in i:
        if i > 0:
            text += " "
        end
        text += "" + ((t.at(i) * 100.0 + 10000.5) as int - 10000)
    end
    ret text
end

func main:
    mut tape := autograd.NewTape(1)
    mut x_value := tensor.Arange(3)
    mut three_value := tensor.Full(tensor.Shape1(3), 3.0)
    x := autograd.Watch(tape, x_value)
    three := autograd.Constant(tape, three_value)
    squared := autograd.Square(tape, x)
    linear := autograd.Mul(tape, three, x)
    loss := autograd.Sum(tape, autograd.Add(tape, squared, linear))
    autograd.Backward(tape, loss)
    println rounded(autograd.Value(tape, loss)), rounded(autograd.Grad(tape, x)), tape.len()

    tape.reset()
    a := autograd.Input(tape, x_value)
    one := autograd.AddScalar(tape, a, 1.0)
    every := [
        autograd.Sub(tape, a, one), autograd.Div(tape, a, one), autograd.Neg(tape, a), autograd.Scale(tape, a, 2.0),
        autograd.Exp(tape, a), autograd.Log(tape, one), autograd.Tanh(tape, a), autograd.Sigmoid(tape, a),
        autograd.Relu(tape, a), autograd.Gelu(tape, a), autograd.Silu(tape, a), autograd.Sqrt(tape, one),
        autograd.PowScalar(tape, a, 3.0), autograd.Softmax(tape, a), autograd.LogSoftmax(tape, a), autograd.Mean(tape, a)
    ]
    each node in every:
        println rounded(autograd.Value(tape, node))
    end

    tape.reset()
    mut m_value := tensor.Reshape(tensor.Arange(6), tensor.Shape2(2, 3))
    mut m_grad := tensor.Zeros2(2, 3)
    m := autograd.Param(tape, m_value, m_grad)
    t := autograd.Transpose2D(tape, m)
    product := autograd.MatMul(tape, m, t)
    println rounded(autograd.Value(tape, product))
    println rounded(autograd.Value(tape, autograd.SumAxis(tape, m, 0, false)))
    println rounded(autograd.Value(tape, autograd.MeanAxis(tape, m, 1, false)))
    mut shape := tensor.Shape1(6)
    println tensor.ShapeString(autograd.Value(tape, autograd.Reshape(tape, m, shape)))
    autograd.Backward(tape, autograd.Sum(tape, product))
    println rounded(m_grad)
    tape.free()
end

Output1400 300 500 700 6 -100 -100 -100 0 50 67 0 -100 -200 0 200 400 100 272 739 0 69 110 0 76 96 50 73 88 0 100 200 0 84 195 0 73 176 100 141 173 0 100 800 9 24 67 -241 -141 -41 100 500 1400 1400 5000 300 500 700 100 400 6 600 1000 1400 600 1000 1400

`Tape`

struct Tape

The tape: it records every computation so backpropagation can work out the gradients. Each node is known by an integer index.

Name	Type	Description
`vals`	`Vector`	The values of the nodes.
`grads`	`Vector`	The gradients of the nodes.
`owns_val`	`Vector`	The tape owns the value.
`owns_grad`	`Vector`	The tape owns the gradient.
`needs`	`Vector`	The node needs a gradient.
`nodes`	`Vector`	The nodes.
`training`	`bool = true`	Training mode; dropout only acts in training.
`rng`	`tensor.Rng`	The random generator of the tape.

The example for this item is shown with the ones above: autograd-Basics.salam

`len`

func len(): int

The number of nodes.

The example for this item is shown with the ones above: autograd-Basics.salam

`reset`

func reset()

Empties the tape for the next step.

The example for this item is shown with the ones above: autograd-Basics.salam

`free`

func free()

Frees the tape.

The example for this item is shown with the ones above: autograd-Basics.salam

Functions

`MSELoss`

func MSELoss(tp &: Tape, pred: int, target: int): int

The mean squared error loss; for regression.

autograd-Layers.salam

import autograd
import tensor

func rounded(t: tensor.Tensor): str:
    mut text := ""
    repeat t.numel() in i:
        if i > 0:
            text += " "
        end
        text += "" + ((t.at(i) * 100.0 + 10000.5) as int - 10000)
    end
    ret text
end

func cube(t: tensor.Tensor): f64:
    mut total := 0.0
    repeat t.numel() in i:
        total += t.at(i) * t.at(i) * t.at(i)
    end
    ret total
end

func main:
    mut tape := autograd.NewTape(7)
    mut image := tensor.Reshape(tensor.Arange(16), tensor.Shape4(1, 1, 4, 4))
    mut kernel := tensor.Ones(tensor.Shape4(2, 1, 2, 2))
    mut bias := tensor.Zeros1(2)
    img := autograd.Input(tape, image)
    weight := autograd.Constant(tape, kernel)
    b := autograd.Constant(tape, bias)
    conv := autograd.Conv2D(tape, img, weight, b, 2, 2, 0, 0)
    println tensor.ShapeString(autograd.Value(tape, conv))
    println rounded(autograd.Value(tape, autograd.MaxPool2D(tape, img, 2, 2, 2, 2)))
    println rounded(autograd.Value(tape, autograd.AvgPool2D(tape, img, 2, 2, 2, 2)))
    println tensor.ShapeString(autograd.Value(tape, autograd.GroupNorm(tape, conv, autograd.Constant(tape, tensor.Ones(tensor.Shape1(2))), autograd.Constant(tape, tensor.Zeros1(2)), 2, 0.00001)))
    mut running_mean := tensor.Zeros1(2)
    mut running_var := tensor.Ones(tensor.Shape1(2))
    norm := autograd.BatchNorm(tape, conv, autograd.Constant(tape, tensor.Ones(tensor.Shape1(2))), autograd.Constant(tape, tensor.Zeros1(2)), running_mean, running_var, 0.1, 0.00001)
    println tensor.ShapeString(autograd.Value(tape, norm)), rounded(running_mean)

    mut row := tensor.Reshape(tensor.Arange(8), tensor.Shape2(2, 4))
    r := autograd.Input(tape, row)
    gamma := autograd.Constant(tape, tensor.Ones(tensor.Shape1(4)))
    beta := autograd.Constant(tape, tensor.Zeros1(4))
    println rounded(autograd.Value(tape, autograd.LayerNorm(tape, r, gamma, beta, 0.00001)))
    println tensor.ShapeString(autograd.Value(tape, autograd.RMSNorm(tape, r, gamma, 0.00001)))
    println rounded(autograd.Value(tape, autograd.Narrow(tape, r, 1, 1, 2)))
    println rounded(autograd.Value(tape, autograd.SliceAxis0(tape, r, 1, 1)))
    println tensor.ShapeString(autograd.Value(tape, autograd.Concat2(tape, r, r, 0)))
    mut order := Vector {} as Vector<int>
    order.push(1)
    order.push(0)
    println tensor.ShapeString(autograd.Value(tape, autograd.Permute(tape, r, order)))
    mut three_dim := tensor.Ones(tensor.Shape3(2, 2, 3))
    s := autograd.Input(tape, three_dim)
    println tensor.ShapeString(autograd.Value(tape, autograd.BatchMatMul(tape, s, autograd.Permute(tape, s, tensor.Shape3(0, 2, 1)))))
    println tensor.ShapeString(autograd.Value(tape, autograd.RoPE(tape, r, 2, 1, 0, 10000.0)))
    println rounded(autograd.Value(tape, autograd.Dropout(tape, r, 0.0)))

    mut vocab := tensor.Reshape(tensor.Arange(6), tensor.Shape2(3, 2))
    mut indices := Vector {} as Vector<int>
    indices.push(2)
    indices.push(0)
    println rounded(autograd.Value(tape, autograd.Embedding(tape, autograd.Constant(tape, vocab), indices)))

    mut prediction := tensor.Full(tensor.Shape1(2), 1.0)
    mut target := tensor.Zeros1(2)
    println rounded(autograd.Value(tape, autograd.MSELoss(tape, autograd.Input(tape, prediction), autograd.Constant(tape, target))))
    println rounded(autograd.Value(tape, autograd.BCEWithLogits(tape, autograd.Input(tape, target), autograd.Constant(tape, target))))
    mut logits := tensor.Zeros2(2, 3)
    println rounded(autograd.Value(tape, autograd.CrossEntropyLogits(tape, autograd.Input(tape, logits), indices)))

    mut point := tensor.Arange(3)
    analytic := tensor.Mul(tensor.Square(point), tensor.Full(tensor.Shape1(3), 3.0))
    println autograd.GradCheck((t: tensor.Tensor): ret cube(t) end, point, analytic, 0.001, 0.01)
    tape.free()
end

Output1x2x2x2 500 700 1300 1500 250 450 1050 1250 1x2x2x2 1x2x2x2 300 300 -134 -45 45 134 -134 -45 45 134 2x4 100 200 500 600 400 500 600 700 4x4 4x2 2x2x2 2x4 0 100 200 300 400 500 600 700 400 500 0 100 100 69 110 true

`CrossEntropyLogits`

func CrossEntropyLogits(tp &: Tape, logits: int, targets &: Vector<int>): int

The cross entropy loss from logits and the right class; for classification.

The example for this item is shown with the ones above: autograd-Layers.salam

`BCEWithLogits`

func BCEWithLogits(tp &: Tape, logits: int, target: int): int

Binary cross entropy from a logit; for binary classification.

The example for this item is shown with the ones above: autograd-Layers.salam

`Conv2D`

func Conv2D(tp &: Tape, x: int, w: int, b: int, sh: int, sw: int, ph: int, pw: int): int

A two dimensional convolution.

The example for this item is shown with the ones above: autograd-Layers.salam

`MaxPool2D`

func MaxPool2D(tp &: Tape, x: int, kh: int, kw: int, sh: int, sw: int): int

Two dimensional max pooling.

The example for this item is shown with the ones above: autograd-Layers.salam

`AvgPool2D`

func AvgPool2D(tp &: Tape, x: int, kh: int, kw: int, sh: int, sw: int): int

Two dimensional average pooling.

The example for this item is shown with the ones above: autograd-Layers.salam

`Add`

func Add(tp &: Tape, a: int, b: int): int

Addition with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Sub`

func Sub(tp &: Tape, a: int, b: int): int

Subtraction with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Mul`

func Mul(tp &: Tape, a: int, b: int): int

Elementwise multiplication with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Div`

func Div(tp &: Tape, a: int, b: int): int

Elementwise division with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Neg`

func Neg(tp &: Tape, a: int): int

Negation with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Scale`

func Scale(tp &: Tape, a: int, s: f64): int

Multiplication by a number.

The example for this item is shown with the ones above: autograd-Basics.salam

`AddScalar`

func AddScalar(tp &: Tape, a: int, s: f64): int

Addition of a number.

The example for this item is shown with the ones above: autograd-Basics.salam

`Exp`

func Exp(tp &: Tape, a: int): int

The exponential with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Log`

func Log(tp &: Tape, a: int): int

The logarithm with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Tanh`

func Tanh(tp &: Tape, a: int): int

The hyperbolic tangent with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Sigmoid`

func Sigmoid(tp &: Tape, a: int): int

The sigmoid with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Relu`

func Relu(tp &: Tape, a: int): int

The rectifier with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Gelu`

func Gelu(tp &: Tape, a: int): int

GELU with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Sqrt`

func Sqrt(tp &: Tape, a: int): int

The square root with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Square`

func Square(tp &: Tape, a: int): int

The square with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`PowScalar`

func PowScalar(tp &: Tape, a: int, p: f64): int

A power.

The example for this item is shown with the ones above: autograd-Basics.salam

`Silu`

func Silu(tp &: Tape, a: int): int

SiLU with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`MatMul`

func MatMul(tp &: Tape, a: int, b: int): int

Matrix multiplication with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Sum`

func Sum(tp &: Tape, a: int): int

The sum of everything with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Mean`

func Mean(tp &: Tape, a: int): int

The mean of everything with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`SumAxis`

func SumAxis(tp &: Tape, a: int, axis: int, keepdim: bool): int

The sum along one dimension.

The example for this item is shown with the ones above: autograd-Basics.salam

`MeanAxis`

func MeanAxis(tp &: Tape, a: int, axis: int, keepdim: bool): int

The mean along one dimension.

The example for this item is shown with the ones above: autograd-Basics.salam

`Softmax`

func Softmax(tp &: Tape, a: int): int

The softmax with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`LogSoftmax`

func LogSoftmax(tp &: Tape, a: int): int

The log softmax with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Reshape`

func Reshape(tp &: Tape, a: int, shape &: Vector<int>): int

A reshape.

The example for this item is shown with the ones above: autograd-Basics.salam

`Transpose2D`

func Transpose2D(tp &: Tape, a: int): int

The transpose with a gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`LayerNorm`

func LayerNorm(tp &: Tape, x: int, gamma: int, beta: int, eps: f64): int

Layer normalisation with gamma and beta.

The example for this item is shown with the ones above: autograd-Layers.salam

`RMSNorm`

func RMSNorm(tp &: Tape, x: int, gamma: int, eps: f64): int

Root mean square normalisation, as Llama uses.

The example for this item is shown with the ones above: autograd-Layers.salam

`BatchNorm`

func BatchNorm(tp &: Tape, x: int, gamma: int, beta: int, running_mean &: tensor.Tensor, running_var &: tensor.Tensor, momentum: f64, eps: f64): int

Batch normalisation; it updates the running mean and variance during training.

The example for this item is shown with the ones above: autograd-Layers.salam

`GroupNorm`

func GroupNorm(tp &: Tape, x: int, gamma: int, beta: int, groups: int, eps: f64): int

Group normalisation over the channels.

The example for this item is shown with the ones above: autograd-Layers.salam

`BatchMatMul`

func BatchMatMul(tp &: Tape, a: int, b: int): int

Batched matrix multiplication with a gradient.

The example for this item is shown with the ones above: autograd-Layers.salam

`Permute`

func Permute(tp &: Tape, a: int, perm &: Vector<int>): int

A permutation of the dimensions.

The example for this item is shown with the ones above: autograd-Layers.salam

`Embedding`

func Embedding(tp &: Tape, weight: int, indices &: Vector<int>): int

The rows of an embedding table, by index.

The example for this item is shown with the ones above: autograd-Layers.salam

`Dropout`

func Dropout(tp &: Tape, a: int, p: f64): int

In training, zeroes the elements with probability p.

The example for this item is shown with the ones above: autograd-Layers.salam

`Narrow`

func Narrow(tp &: Tape, a: int, axis: int, start: int, len: int): int

A range of one dimension.

The example for this item is shown with the ones above: autograd-Layers.salam

`SliceAxis0`

func SliceAxis0(tp &: Tape, a: int, start: int, len: int): int

A range of the first dimension.

The example for this item is shown with the ones above: autograd-Layers.salam

`Concat2`

func Concat2(tp &: Tape, a: int, b: int, axis: int): int

Joins two nodes.

The example for this item is shown with the ones above: autograd-Layers.salam

`RoPE`

func RoPE(tp &: Tape, a: int, seq: int, heads: int, offset: int, theta: f64): int

Rotary position embedding, for attention.

The example for this item is shown with the ones above: autograd-Layers.salam

`NewTape`

func NewTape(seed: i64): Tape

A tape with a random generator seed.

The example for this item is shown with the ones above: autograd-Basics.salam

`Input`

func Input(tp &: Tape, t &: tensor.Tensor): int

Puts a tensor on the tape with no gradient.

The example for this item is shown with the ones above: autograd-Basics.salam

`Constant`

func Constant(tp &: Tape, t &: tensor.Tensor): int

A constant; as an input.

The example for this item is shown with the ones above: autograd-Basics.salam

`Watch`

func Watch(tp &: Tape, t &: tensor.Tensor): int

Records a tensor with a gradient, so it can be read with Grad after backpropagation.

The example for this item is shown with the ones above: autograd-Basics.salam

`Param`

func Param(tp &: Tape, val &: tensor.Tensor, grad &: tensor.Tensor): int

A parameter with a gradient tensor that backpropagation accumulates into.

The example for this item is shown with the ones above: autograd-Basics.salam

`Value`

func Value(tp &: Tape, v: int): tensor.Tensor

The value of a node.

The example for this item is shown with the ones above: autograd-Basics.salam

`Grad`

func Grad(tp &: Tape, v: int): tensor.Tensor

The gradient of a node after backpropagation.

The example for this item is shown with the ones above: autograd-Basics.salam

`Backward`

func Backward(tp &: Tape, loss: int)

Backpropagates from the loss node.

The example for this item is shown with the ones above: autograd-Basics.salam

`GradCheck`

func GradCheck(f: func(tensor.Tensor): f64, x &: tensor.Tensor, analytic &: tensor.Tensor, h: f64, tol: f64): bool

Measures the analytic gradient against a numeric difference; for testing a new layer.

The example for this item is shown with the ones above: autograd-Layers.salam