Automatic Differentiation¶

When a model doesn't fit the Sequential mold (weight sharing, custom losses, exotic architectures), build the computation directly and let reverse-mode autodiff derive the gradients. The engine is micrograd-style: a dynamically built graph of Nodes, define-by-run, single-use.

A Node holds an n-dimensional Tensor, so the same ops that run on a (rows, cols) matrix run on a (batch, sequence, model) activation: element-wise arithmetic broadcasts NumPy-style and MatMul multiplies whole stacks of matrices. A Matrix is taken as a zero-copy 2-D view wherever a leaf is built, so two-dimensional code reads the way it always did.

Params, Inputs, and the Trainer¶

w1 := autograd.Param(tensai.RandomMatrix(2, 8, rng))
b1 := autograd.Param(tensai.NewMatrix(1, 8))
w2 := autograd.Param(tensai.RandomMatrix(8, 1, rng))
trainer := autograd.NewTrainer(optim.NewAdam(0.05), w1, b1, w2)

for step := 0; step < 2000; step++ {
    loss := autograd.Input(x).MatMul(w1).AddRow(b1).Tanh().MatMul(w2).Sigmoid().MSELoss(y.Tensor())
    trainer.Step(loss) // backward + update + zero grads, returns the loss value
}

Param wraps a value whose gradient should be tracked and updated; Input wraps data. Both accept a *tensai.Matrix or a *tensai.Tensor — a matrix becomes a 2-D tensor view sharing the same backing array, so a parameter built from a matrix keeps updating that matrix.

For manual control, the pieces are public: loss.Backward(), p.Grad(), and autograd.ZeroGrads(params...). Backward starts from a single-element node and accumulates into the gradient, which is why a training step clears the gradients afterwards (Trainer.Step does it for you).

On a node, Value() and Grad() return *tensai.Tensor. They are methods rather than fields because a value does not have to live in host memory: on a device-resident graph they are where it comes back. Shape() returns the shape without fetching anything, Matrix() returns a matrix view of a 2-D node, Scalar() reads a single-element node, and Named("w1") labels a leaf for ToDot.

Ops¶

Op	Shapes
`MatMul(o)`	`(…, m, k) * (…, k, n)` → `(…, m, n)`; the leading axes broadcast, so one 2-D weight applies to a whole batch
`Add`, `Sub`, `Mul` (`MulElem`), `Div`	element-wise, NumPy broadcasting
`Scale(s)`, `Neg()`	element-wise by a scalar
`AddRow(row)`	`(m, n) + (1, n)`; `Add` already broadcasts, this only states the intent
`T()`, `Transpose(perm...)`	swap the last two axes / permute every axis
`Reshape(shape...)`	one dimension may be `-1`; shares the buffer
`Softmax()`	over the last axis
`LayerNorm(gain, bias, eps)`	over the last axis; `gain` and `bias` hold one element per feature and may be `nil` or trainable
`Embed(ids, shape...)`	`(vocab, d)` table plus `len(ids)` indices → `shape…, d`; repeated ids accumulate on backward
`Conv2D(w, bias, cfg)`	`(batch, channels, h, w)` convolved with `(channelskk, outChannels)` weights → `(batch, outChannels, outH, outW)`
`Im2Col(cfg)`	the patch expansion a convolution multiplies, on its own
`MaxPool2D(size)`, `AvgPool2D(size)`	square windows over the last two axes, stride equal to the window
`Sum()`, `Mean()`	reduce everything to one element
`SumAxis(axis, keepDims)`, `MeanAxis(axis, keepDims)`	reduce one axis; a negative axis counts from the end
`ReLU()`, `LeakyReLU(a)`, `Sigmoid()`, `Tanh()`, `GELU()`, `Exp()`, `Log()`, `Sqrt()`	element-wise
`Dropout(rate, rng)`	zeroes elements and scales the survivors; the mask is drawn once and the backward pass reuses it
`MSELoss(target)`, `SoftmaxCELoss(target)`, `CrossEntropy(labels []int)`	scalar losses; the cross-entropies read the last axis as classes

Graphs are built dynamically per step and are single-use. Shape mismatches panic during construction rather than returning an error: chaining would be unusable otherwise, and a wrong shape is a programming mistake. Every op's gradient is verified against finite differences in the test suite.

Convolutions¶

A convolution is an Im2Col and a product: the expansion turns each output pixel into a row of the patch that produced it, and the weights are one column per output channel.

// x is (batch, channels, height, width); w is (channels*k*k, outChannels).
h := x.Conv2D(w, bias, autograd.Conv{Kernel: 3, Pad: 1})
h = h.ReLU().MaxPool2D(2)

Conv2D composes the pieces -- x.Im2Col(cfg).MatMul(w).Transpose(0, 2, 1).Reshape(...) plus the bias -- so the product it spends its time in is the same one the GPU accelerates. Im2Col is exposed on its own for anything the composition does not cover, and MaxPool2D and AvgPool2D route their gradients to the position that won and over the window respectively. The forward pass matches layer.Conv2D, which the Sequential models use.

Broadcasting¶

An element-wise op aligns shapes at their trailing axes and stretches any axis of length 1, so a (1, 1, d) bias adds to a (batch, seq, d) activation. Gradients follow the same rule in reverse: whatever axes an operand was stretched along, its gradient is summed back over. That is why a bias, a per-feature LayerNorm gain, and a weight shared across a batch all collect the contributions of every position that used them without any special case in the op.

Building attention¶

Multi-head attention is a reshape, a transpose, and two batched products:

// x is (batch, seq, model); wq, wk, wv, wo are (model, model).
heads := func(t *autograd.Node) *autograd.Node {
    // (batch, seq, model) -> (batch, head, seq, headDim)
    return t.Reshape(batch, seq, nHeads, headDim).Transpose(0, 2, 1, 3)
}
q, k, v := heads(x.MatMul(wq)), heads(x.MatMul(wk)), heads(x.MatMul(wv))

// (batch, head, seq, seq) scores for every head of every sequence at once.
att := q.MatMul(k.T()).Scale(1 / float32(math.Sqrt(headDim))).Add(mask).Softmax()

y := att.MatMul(v).Transpose(0, 2, 1, 3).Reshape(batch, seq, model).MatMul(wo)

mask is a constant Input of shape (1, 1, seq, seq) holding 0 on and below the diagonal and math.Inf(-1) above it, broadcast over batch and head; softmax then gives future positions no weight. _example/tinygpt puts exactly this into a working character-level transformer — pre-norm blocks, a GELU feed-forward, and next-character cross-entropy — that memorizes a page of text in about a minute.

Reusing buffers with a tape¶

A graph is built and thrown away every step, so by default every step asks the allocator for every intermediate value and gradient it touches. A Tape recycles them. Bind the parameters once, then reset it at the end of each step:

tape := autograd.NewTape()
tape.Bind(w1, b1, w2, b2) // ops inherit the tape from their parents

for step := 0; step < steps; step++ {
    trainer.Step(forward(x).MSELoss(y))
    tape.Reset() // hands this step's buffers back to the pool
}

The rule is the one a training loop already follows: after Reset, nothing from the finished step may be read — no node's Value, no node's Grad. Parameter values are never recycled (the tape only owns what operations produce), so trained weights are always safe to keep; copy anything else with Clone before resetting. A Tape is not safe for concurrent use, so give each training goroutine its own.

On _example/charrnn, which unrolls 32 time steps per iteration, the tape takes a training step from 22 MB of allocation to 0.75 MB and cuts about a quarter off its wall time.

The same reuse is available one layer down: MatMulInto, MatMulTNInto, MatMulNTInto, AddInto, SubInto, MulInto and DivInto write into a tensor you already own, the way DotInto does for matrices.

Running the graph on a GPU¶

A tape can put the whole graph on a device. Values, gradients and the Adam update stay there; only the loss comes home each step:

dev, err := gpu.Open(gpu.HighPerformance)
if err == nil {
    defer dev.Close()
    tape.UseDevice(dev)
}
tape.Bind(params...) // parameters upload once and stay resident

Value() and Grad() download on demand, so reading a node still works — that is what makes them methods. Resident() reports where a node's value currently is.

Operations the device has kernels for run there; anything else falls back to the CPU, bringing home only what that operation needs. That covers the three products, element-wise arithmetic with a repeating operand, Scale, ReLU/tanh/sigmoid/GELU and their gradients, LayerNorm, Softmax, Transpose, Reshape, Embed (its scatter-add included), the sum a broadcast collects, and the Adam update — so a whole transformer block stays resident, and only the loss crosses the bus.

On an AMD 780M through -tags wgpu24, one step of x @ w1 -> GELU -> @ w2 -> MSE:

model width	CPU (AVX2)	accelerator hook	resident graph
512	31.1ms	30.8ms	13.1ms
1024	172.8ms	93.2ms	33.3ms
2048	1382ms	442ms	193ms

A transformer block -- pre-norm attention with a causal mask and a GELU feed-forward, batch 8, sequence 128, vocabulary 512 -- trains at:

model width	CPU (AVX2)	resident graph
256	83.5ms	47.0ms
512	209.3ms	54.3ms

The hook (tensai.UseAccelerator) sends each product to the device on its own, so it pays a round trip per product; residency pays one upload per step for the batch and one download for the loss. Both are opt-in, and they compose: with a tape on a device the hook is not consulted.

Visualizing the graph¶

loss.ToDot() returns Graphviz DOT (label leaves with .Named("w1")):

go run ./_example/dot | dot -Tsvg > graph.svg

Recurrent networks¶

rnn.Cell and rnn.LSTMCell are built on the autograd engine, so unrolling a sequence is a plain Go loop and backpropagation through time comes for free:

cell := rnn.NewLSTMCell(inSize, hidden, rng)
wOut := autograd.Param(tensai.RandomMatrix(hidden, numClasses, rng))
bOut := autograd.Param(tensai.NewMatrix(1, numClasses))
trainer := autograd.NewTrainer(optim.NewAdam(0.01), append(cell.Params(), wOut, bOut)...)

for step := 0; step < epochs; step++ {
    h, c := cell.InitState(batch)
    for _, x := range steps { // one (batch x inSize) matrix per time step
        h, c = cell.Step(autograd.Input(x), h, c)
    }
    logits := h.MatMul(wOut).AddRow(bOut)
    trainer.Step(logits.CrossEntropy(labels)) // labels is a []int of class indices
}

_example/charrnn trains a character-level LSTM on an embedded public-domain text and generates samples from the reloaded parameters.

rnn.SelfAttention is the single-head, one-sequence form: attn.Forward(x) computes softmax(Q*K^T/sqrt(d))*V on a (seqLen, inSize) node with learned projections, and the raw rnn.Attention(q, k, v) is also exposed. For batches and heads, write the block above.

Saving parameters¶

Autograd parameters are saved and restored positionally:

autograd.SaveParamsFile("cell.json", cell.Params()...)
// build the same cell, then
autograd.LoadParamsFile("cell.json", cell.Params()...)

Parameters of any rank round-trip; two-dimensional ones keep the encoding earlier checkpoints used, so files written before the engine went n-dimensional still load.