Notebooks · lib/nn

The pieces a transformer is made of

A transformer block is four things: a matrix multiply, a normalization, a pointwise nonlinearity, and a softmax. This module has the forward pass of each, and of the pieces around the blocks: the embedding lookup and the cross-entropy. Training needs each of them backwards as well, and that comes from value_and_grad over the forward pass; the tests check the gradient it generates for every piece against central differences on the forward, which is the only way to know a gradient is right.

The choices here follow the nanogpt speedrun rather than the 2017 paper, because that is what a fast implementation looks like now:

RMSNorm with no learnable gain. Divide each row by the root mean square of its entries. There is nothing to train and nothing to store.
ReLU squared instead of GELU. max(x, 0)^2. Its derivative is 2 max(x, 0), which is continuous, and it costs one compare and one multiply.
Linear layers with no bias. The bias is redundant after a normalization and it is one more thing to move.

A row of X is a token and a column is a feature, so X is T by D for a sequence of T tokens in D dimensions, or B by T by D for a batch of B sequences: every def works along the last axis (the features) or the last two (the positions and the features), and reads T and D from its arguments.

The backward passes written by hand before value_and_grad existed are in lib/nn_adjoints, kept as references: its tests hold them to the generated ones, and the study/idioms ports that write a backward pass by hand, to compare it with the generated one, call them.

Used with import nn.

Look inside this program

Atlas
The program’s map: every function and which calls which.
There’s no recorded run of this program yet.

This program declares a record type (AttentionCase), and records run only in compiled C for now, so the Run buttons here are off. Open it in the workbench and use its C tab to compile it.

Shift+Enter in a cell runs it and moves to the next. Running a cell runs every cell above it too, as one program.

[1]lines 37–392

from linalg import test_matrix

NN_EPS = 1.0e-8

def row_rms(X: Arr) {
    return sqrt(mean(X * X, axis=-1, keepdims=true) + NN_EPS)
}

def rmsnorm(X: Arr) {
    return X / row_rms(X)
}

def relu2(X: Arr) {
    R = max(X, 0.0)
    return R * R
}

def softmax_rows(X: Arr) {
    e = exp(X - amax(X, axis=-1, keepdims=true))
    return e / sum(e, axis=-1, keepdims=true)
}

MASKED = -1.0e30

# The words attention is written in. The
# softmax of each row, the causal mask and the cross-entropy are phrases of
# Viv's table.
phrase "the attention scores of Q against K" -> (Q @ swapaxes(K, -1, -2)) * (1.0 / sqrt(1.0 * shape(Q)[-1]))
phrase "S masked where future" -> where(future, MASKED, S)

def attention(X: Arr, Wq: Arr, Wk: Arr, Wv: Arr) {
    T = shape(X)[-2]
    step project {
        Q = X @ Wq
        K = X @ Wk
        V = X @ Wv
    }
    step scores: S = take the attention scores of Q against K
    step mask {
        future = take the causal mask of size T
        S = take S masked where future
    }
    step weights: A = take the softmax of each row of S
    step mix: Y = A @ V
    return Y
}

# the weights alone, for a reader who wants to see where each position looked
# (the same operations as attention's first four steps, so the same numbers)
def attention_weights(X: Arr, Wq: Arr, Wk: Arr) {
    Q, K = X @ Wq, X @ Wk
    S = take the attention scores of Q against K
    future = take the causal mask of size shape(X)[-2]
    return take the softmax of each row of where(future, MASKED, S)
}

def cross_entropy(logits: Arr, target: Arr) {
    classes = numel(logits) / numel(target)
    top = amax(logits, axis=-1, keepdims=true)
    logp = logits - top - log(sum(exp(logits - top), axis=-1, keepdims=true))
    return -mean(take(logp, arange(numel(target)) * classes + ravel(target)))
}

def embedding(W: Arr, ids: Arr) {
    return take(W, ids, axis=0)
}

def argmax_rows(logits: Arr) {
    best = zeros(len(logits))
    for i to len(logits) {
        for j to shape(logits)[1] {
            if logits[i, j] > logits[i, round(best[i])] {
                best[i] = j
            }
        }
    }
    return best
}

# ---- checks ----

def worst_rel(A: Arr, B: Arr) {
    return max(abs(A - B) / max(max(abs(A), abs(B)), 1.0e-8))
}

def central(f, X: Arr, C: Arr) {
    h = 1.0e-6
    num = X * 0.0
    for i to len(X) {
        for j to shape(X)[1] {
            was = X[i, j]
            X[i, j] = was + h
            up = f(X, C)
            X[i, j] = was - h
            dn = f(X, C)
            X[i, j] = was
            num[i, j] = (up - dn) / (2.0 * h)
        }
    }
    return num
}

# the central difference of batch_attention_probe in every entry of sequence b
def central_slab(X: Arr, C: Arr, b) {
    h = 1.0e-6
    rows, cols = shape(X[b, :, :])
    num = zeros(rows, cols)
    for i to rows {
        for j to cols {
            was = X[b, i, j]
            X[b, i, j] = was + h
            up = batch_attention_probe(X, C)
            X[b, i, j] = was - h
            dn = batch_attention_probe(X, C)
            X[b, i, j] = was
            num[i, j] = (up - dn) / (2.0 * h)
        }
    }
    return num
}

def nn_test_matrix(rows, cols, seed) {
    M = test_matrix(max(rows, cols), seed)
    return M[0:rows, 0:cols]
}

def rmsnorm_probe(X: Arr, C: Arr) {
    return sum(C * rmsnorm(X))
}

def relu2_probe(X: Arr, C: Arr) {
    return sum(C * relu2(X))
}

def softmax_probe(X: Arr, C: Arr) {
    return sum(C * softmax_rows(X))
}

def embedding_probe(W: Arr, C: Arr) {
    return sum(C * embedding(W, [2.0, 0.0, 2.0]))
}

record AttentionCase { X: Arr, Wq: Arr, Wk: Arr, Wv: Arr }

def attention_case() {
    return AttentionCase {
        X: nn_test_matrix(5, 5, 7.0),
        Wq: nn_test_matrix(5, 3, 19.0),
        Wk: nn_test_matrix(5, 3, 37.0),
        Wv: nn_test_matrix(5, 3, 59.0),
    }
}

def attention_probe(a: AttentionCase, C: Arr) {
    return sum(C * attention(a.X, a.Wq, a.Wk, a.Wv))
}

# the central difference of attention_probe in every entry of w, a matrix of
# the case `a`, into the same entries of nw
def central_into(w: Arr, nw: Arr, C: Arr, a: AttentionCase) {
    h = 1.0e-6
    for i to len(w) {
        for j to shape(w)[1] {
            was = w[i, j]
            w[i, j] = was + h
            up = attention_probe(a, C)
            w[i, j] = was - h
            dn = attention_probe(a, C)
            w[i, j] = was
            nw[i, j] = (up - dn) / (2.0 * h)
        }
    }
}

def central_case(a: AttentionCase, C: Arr) {
    num = zeros_like(a)
    for_leaves(central_into, a, num, C, a)
    return num
}

test "RMSNorm gives each row a root mean square of one, and its gradient matches central differences" {
    X = nn_test_matrix(4, 5, 67.0)
    C = nn_test_matrix(4, 5, 89.0)
    Y = rmsnorm(X)
    # every row came out with unit root mean square, up to the epsilon in the
    # divisor -- which is there so that a row of zeros does not divide by zero,
    # and which costs about 1e-7 here
    for i to len(Y) {
        assert_near(sum(Y[i, :] * Y[i, :]) / shape(Y)[1], 1.0, 1.0e-6)
    }
    v, g = value_and_grad(rmsnorm_probe, X, C)
    assert_lt(worst_rel(g, central(rmsnorm_probe, X, C)), 1.0e-6)
}

test "ReLU squared's gradient matches central differences, away from the kink" {
    X = nn_test_matrix(4, 5, 13.0)
    C = nn_test_matrix(4, 5, 29.0)
    v, g = value_and_grad(relu2_probe, X, C)
    assert_lt(worst_rel(g, central(relu2_probe, X, C)), 1.0e-6)
    # the derivative is continuous at zero, which is the reason for the square
    assert_near(2.0 * max(0.0, 0.0), 0.0, 1.0e-15)
}

test "softmax rows sum to one and its gradient matches central differences" {
    X = nn_test_matrix(4, 5, 53.0)
    C = nn_test_matrix(4, 5, 71.0)
    Y = softmax_rows(X)
    for i to len(Y) {
        assert_near(sum(Y[i, :]), 1.0, 1.0e-12)
        assert_lt(0.0, min(Y[i, :]))
    }
    v, g = value_and_grad(softmax_probe, X, C)
    assert_lt(worst_rel(g, central(softmax_probe, X, C)), 1.0e-6)
}

test "softmax does not overflow when the scores are large" {
    rows = 2
    cols = 4
    X = zeros(rows, cols)
    for i to rows {
        X[i, :] = 800.0 + 3.0 * linspace(0.0, cols - 1.0, cols)
    }
    Y = softmax_rows(X)
    for i to len(Y) {
        assert_near(sum(Y[i, :]), 1.0, 1.0e-12)
    }
    # exp(800) is infinite in double precision; shifting by the row max is what
    # keeps this finite
    assert_lt(0.0, Y[0, 3])
}

test "causal attention looks only backwards, and its weights are a distribution" {
    a = attention_case()
    A = attention_weights(a.X, a.Wq, a.Wk)
    for i to len(A) {
        assert_near(sum(A[i, :]), 1.0, 1.0e-12)
        # nothing after position i may be attended to
        for j to len(A) {
            if j > i {
                assert_near(A[i, j], 0.0, 1.0e-300)
            }
        }
    }
    # the first token can only attend to itself
    assert_near(A[0, 0], 1.0, 1.0e-12)
}

test "attention's four gradients all match central differences" {
    # one gradient of the whole case: an AttentionCase holding the gradient
    # with respect to each of its four matrices
    a = attention_case()
    C = nn_test_matrix(5, 3, 83.0)
    va, ga = value_and_grad(attention_probe, a, C)
    num = central_case(a, C)
    assert_lt(worst_rel(ga.X, num.X), 1.0e-6)
    assert_lt(worst_rel(ga.Wq, num.Wq), 1.0e-6)
    assert_lt(worst_rel(ga.Wk, num.Wk), 1.0e-6)
    assert_lt(worst_rel(ga.Wv, num.Wv), 1.0e-6)
}

test "cross-entropy is the loss it claims to be, and its gradient checks out" {
    V = 4
    logits = 3.0 * nn_test_matrix(5, V, 97.0)
    target = [mod(1.0 * i, 1.0 * V) for i to len(logits)]
    # a row that is certain and right costs nothing; one that is uniform costs log V
    sure = zeros(1, V)
    sure[0, 2] = 60.0
    assert_near(cross_entropy(sure, [2.0]), 0.0, 1.0e-12)
    assert_near(cross_entropy(zeros(1, V), [2.0]), log(1.0 * V), 1.0e-12)
    loss, d = value_and_grad(cross_entropy, logits, target)
    assert_lt(worst_rel(d, central(cross_entropy, logits, target)), 1.0e-6)
    # the gradient of a whole row sums to zero: probability has to go somewhere
    for i to len(d) {
        assert_near(sum(d[i, :]), 0.0, 1.0e-14)
    }
}

test "an embedding looks rows up, and its gradient adds into the rows it looked up" {
    W = nn_test_matrix(3, 4, 31.0)
    C = nn_test_matrix(3, 4, 43.0)
    X = embedding(W, [2.0, 0.0, 2.0])
    assert_near(max(abs(X[1, :] - W[0, :])), 0.0, 0.0)
    v, g = value_and_grad(embedding_probe, W, C)
    assert_lt(worst_rel(g, central(embedding_probe, W, C)), 1.0e-6)
    # row 2 was looked up twice, row 1 never
    assert_near(max(abs(g[2, :] - C[0, :] - C[2, :])), 0.0, 1.0e-15)
    assert_near(max(abs(g[1, :])), 0.0, 0.0)
}

# the sequences in the batch tests
NN_TEST_BATCH = 2

# a batch of test matrices, one per sequence, each from its own seed
def test_batch(rows, cols, seed) {
    S = zeros(NN_TEST_BATCH, rows, cols)
    for b to NN_TEST_BATCH {
        S[b, :, :] = nn_test_matrix(rows, cols, seed + b)
    }
    return S
}

def batch_attention_probe(X: Arr, C: Arr) {
    a = attention_case()
    return sum(C * attention(X, a.Wq, a.Wk, a.Wv))
}

test "on a batch of sequences, each sequence gets what it gets alone" {
    X = test_batch(5, 5, 7.0)
    a = attention_case()
    Y, A = attention(X, a.Wq, a.Wk, a.Wv), attention_weights(X, a.Wq, a.Wk)
    R = rmsnorm(X)
    P = softmax_rows(X)
    for b to len(X) {
        Xb = X[b, :, :]
        assert_near(max(abs(Y[b, :, :] - attention(Xb, a.Wq, a.Wk, a.Wv))), 0.0, 0.0)
        assert_near(max(abs(A[b, :, :] - attention_weights(Xb, a.Wq, a.Wk))), 0.0, 0.0)
        assert_near(max(abs(R[b, :, :] - rmsnorm(X[b, :, :]))), 0.0, 0.0)
        assert_near(max(abs(P[b, :, :] - softmax_rows(X[b, :, :]))), 0.0, 0.0)
    }
    # the batch's cross-entropy is the mean over every row of every sequence
    target = [[0.0, 1.0, 2.0, 3.0, 4.0], [4.0, 4.0, 0.0, 1.0, 2.0]]
    each = 0.0
    for b to len(X) {
        each = each + cross_entropy(X[b, :, :], target[b, :])
    }
    assert_near(cross_entropy(X, target), each / len(X), 1.0e-15)
    # an embedding of a matrix of ids is one row per id
    E = embedding(nn_test_matrix(3, 4, 31.0), [[2.0, 0.0], [1.0, 2.0]])
    assert_near(max(abs(E[1, 0, :] - nn_test_matrix(3, 4, 31.0)[1, :])), 0.0, 0.0)
}

test "attention's gradient on a batch matches central differences" {
    X = test_batch(5, 5, 7.0)
    C = test_batch(5, 3, 83.0)
    v, g = value_and_grad(batch_attention_probe, X, C)
    for b to len(X) {
        assert_lt(worst_rel(g[b, :, :], central_slab(X, C, b)), 1.0e-6)
    }
}

test "argmax picks the largest entry of each row, and the first on a tie" {
    rows = 3
    cols = 4
    logits = zeros(rows, cols)
    logits[0, 2] = 5.0
    logits[1, 0] = -1.0
    logits[1, 3] = 0.5
    # row 2 is all zeros: a four-way tie, won by column 0
    guess = argmax_rows(logits)
    assert_near(guess[0], 2.0, 0.0)
    assert_near(guess[1], 3.0, 0.0)
    assert_near(guess[2], 0.0, 0.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors

from linalg import test_matrix

NN_EPS = 1.0e-8

def row_rms(X: Arr) {
    return sqrt(mean(X * X, axis=-1, keepdims=true) + NN_EPS)
}

def rmsnorm(X: Arr) {
    return X / row_rms(X)
}

def relu2(X: Arr) {
    R = max(X, 0.0)
    return R * R
}

def softmax_rows(X: Arr) {
    e = exp(X - amax(X, axis=-1, keepdims=true))
    return e / sum(e, axis=-1, keepdims=true)
}

MASKED = -1.0e30

# The words attention is written in. The
# softmax of each row, the causal mask and the cross-entropy are phrases of
# Viv's table.
phrase "the attention scores of Q against K" -> (Q @ swapaxes(K, -1, -2)) * (1.0 / sqrt(1.0 * shape(Q)[-1]))
phrase "S masked where future" -> where(future, MASKED, S)

def attention(X: Arr, Wq: Arr, Wk: Arr, Wv: Arr) {
    T = shape(X)[-2]
    step project {
        Q = X @ Wq
        K = X @ Wk
        V = X @ Wv
    }
    step scores: S = take the attention scores of Q against K
    step mask {
        future = take the causal mask of size T
        S = take S masked where future
    }
    step weights: A = take the softmax of each row of S
    step mix: Y = A @ V
    return Y
}

# the weights alone, for a reader who wants to see where each position looked
# (the same operations as attention's first four steps, so the same numbers)
def attention_weights(X: Arr, Wq: Arr, Wk: Arr) {
    Q, K = X @ Wq, X @ Wk
    S = take the attention scores of Q against K
    future = take the causal mask of size shape(X)[-2]
    return take the softmax of each row of where(future, MASKED, S)
}

def cross_entropy(logits: Arr, target: Arr) {
    classes = numel(logits) / numel(target)
    top = amax(logits, axis=-1, keepdims=true)
    logp = logits - top - log(sum(exp(logits - top), axis=-1, keepdims=true))
    return -mean(take(logp, arange(numel(target)) * classes + ravel(target)))
}

def embedding(W: Arr, ids: Arr) {
    return take(W, ids, axis=0)
}

def argmax_rows(logits: Arr) {
    best = zeros(len(logits))
    for i to len(logits) {
        for j to shape(logits)[1] {
            if logits[i, j] > logits[i, round(best[i])] {
                best[i] = j
            }
        }
    }
    return best
}

# ---- checks ----

def worst_rel(A: Arr, B: Arr) {
    return max(abs(A - B) / max(max(abs(A), abs(B)), 1.0e-8))
}

def central(f, X: Arr, C: Arr) {
    h = 1.0e-6
    num = X * 0.0
    for i to len(X) {
        for j to shape(X)[1] {
            was = X[i, j]
            X[i, j] = was + h
            up = f(X, C)
            X[i, j] = was - h
            dn = f(X, C)
            X[i, j] = was
            num[i, j] = (up - dn) / (2.0 * h)
        }
    }
    return num
}

# the central difference of batch_attention_probe in every entry of sequence b
def central_slab(X: Arr, C: Arr, b) {
    h = 1.0e-6
    rows, cols = shape(X[b, :, :])
    num = zeros(rows, cols)
    for i to rows {
        for j to cols {
            was = X[b, i, j]
            X[b, i, j] = was + h
            up = batch_attention_probe(X, C)
            X[b, i, j] = was - h
            dn = batch_attention_probe(X, C)
            X[b, i, j] = was
            num[i, j] = (up - dn) / (2.0 * h)
        }
    }
    return num
}

def nn_test_matrix(rows, cols, seed) {
    M = test_matrix(max(rows, cols), seed)
    return M[0:rows, 0:cols]
}

def rmsnorm_probe(X: Arr, C: Arr) {
    return sum(C * rmsnorm(X))
}

def relu2_probe(X: Arr, C: Arr) {
    return sum(C * relu2(X))
}

def softmax_probe(X: Arr, C: Arr) {
    return sum(C * softmax_rows(X))
}

def embedding_probe(W: Arr, C: Arr) {
    return sum(C * embedding(W, [2.0, 0.0, 2.0]))
}

record AttentionCase { X: Arr, Wq: Arr, Wk: Arr, Wv: Arr }

def attention_case() {
    return AttentionCase {
        X: nn_test_matrix(5, 5, 7.0),
        Wq: nn_test_matrix(5, 3, 19.0),
        Wk: nn_test_matrix(5, 3, 37.0),
        Wv: nn_test_matrix(5, 3, 59.0),
    }
}

def attention_probe(a: AttentionCase, C: Arr) {
    return sum(C * attention(a.X, a.Wq, a.Wk, a.Wv))
}

# the central difference of attention_probe in every entry of w, a matrix of
# the case `a`, into the same entries of nw
def central_into(w: Arr, nw: Arr, C: Arr, a: AttentionCase) {
    h = 1.0e-6
    for i to len(w) {
        for j to shape(w)[1] {
            was = w[i, j]
            w[i, j] = was + h
            up = attention_probe(a, C)
            w[i, j] = was - h
            dn = attention_probe(a, C)
            w[i, j] = was
            nw[i, j] = (up - dn) / (2.0 * h)
        }
    }
}

def central_case(a: AttentionCase, C: Arr) {
    num = zeros_like(a)
    for_leaves(central_into, a, num, C, a)
    return num
}

test "RMSNorm gives each row a root mean square of one, and its gradient matches central differences" {
    X = nn_test_matrix(4, 5, 67.0)
    C = nn_test_matrix(4, 5, 89.0)
    Y = rmsnorm(X)
    # every row came out with unit root mean square, up to the epsilon in the
    # divisor -- which is there so that a row of zeros does not divide by zero,
    # and which costs about 1e-7 here
    for i to len(Y) {
        assert_near(sum(Y[i, :] * Y[i, :]) / shape(Y)[1], 1.0, 1.0e-6)
    }
    v, g = value_and_grad(rmsnorm_probe, X, C)
    assert_lt(worst_rel(g, central(rmsnorm_probe, X, C)), 1.0e-6)
}

test "ReLU squared's gradient matches central differences, away from the kink" {
    X = nn_test_matrix(4, 5, 13.0)
    C = nn_test_matrix(4, 5, 29.0)
    v, g = value_and_grad(relu2_probe, X, C)
    assert_lt(worst_rel(g, central(relu2_probe, X, C)), 1.0e-6)
    # the derivative is continuous at zero, which is the reason for the square
    assert_near(2.0 * max(0.0, 0.0), 0.0, 1.0e-15)
}

test "softmax rows sum to one and its gradient matches central differences" {
    X = nn_test_matrix(4, 5, 53.0)
    C = nn_test_matrix(4, 5, 71.0)
    Y = softmax_rows(X)
    for i to len(Y) {
        assert_near(sum(Y[i, :]), 1.0, 1.0e-12)
        assert_lt(0.0, min(Y[i, :]))
    }
    v, g = value_and_grad(softmax_probe, X, C)
    assert_lt(worst_rel(g, central(softmax_probe, X, C)), 1.0e-6)
}

test "softmax does not overflow when the scores are large" {
    rows = 2
    cols = 4
    X = zeros(rows, cols)
    for i to rows {
        X[i, :] = 800.0 + 3.0 * linspace(0.0, cols - 1.0, cols)
    }
    Y = softmax_rows(X)
    for i to len(Y) {
        assert_near(sum(Y[i, :]), 1.0, 1.0e-12)
    }
    # exp(800) is infinite in double precision; shifting by the row max is what
    # keeps this finite
    assert_lt(0.0, Y[0, 3])
}

test "causal attention looks only backwards, and its weights are a distribution" {
    a = attention_case()
    A = attention_weights(a.X, a.Wq, a.Wk)
    for i to len(A) {
        assert_near(sum(A[i, :]), 1.0, 1.0e-12)
        # nothing after position i may be attended to
        for j to len(A) {
            if j > i {
                assert_near(A[i, j], 0.0, 1.0e-300)
            }
        }
    }
    # the first token can only attend to itself
    assert_near(A[0, 0], 1.0, 1.0e-12)
}

test "attention's four gradients all match central differences" {
    # one gradient of the whole case: an AttentionCase holding the gradient
    # with respect to each of its four matrices
    a = attention_case()
    C = nn_test_matrix(5, 3, 83.0)
    va, ga = value_and_grad(attention_probe, a, C)
    num = central_case(a, C)
    assert_lt(worst_rel(ga.X, num.X), 1.0e-6)
    assert_lt(worst_rel(ga.Wq, num.Wq), 1.0e-6)
    assert_lt(worst_rel(ga.Wk, num.Wk), 1.0e-6)
    assert_lt(worst_rel(ga.Wv, num.Wv), 1.0e-6)
}

test "cross-entropy is the loss it claims to be, and its gradient checks out" {
    V = 4
    logits = 3.0 * nn_test_matrix(5, V, 97.0)
    target = [mod(1.0 * i, 1.0 * V) for i to len(logits)]
    # a row that is certain and right costs nothing; one that is uniform costs log V
    sure = zeros(1, V)
    sure[0, 2] = 60.0
    assert_near(cross_entropy(sure, [2.0]), 0.0, 1.0e-12)
    assert_near(cross_entropy(zeros(1, V), [2.0]), log(1.0 * V), 1.0e-12)
    loss, d = value_and_grad(cross_entropy, logits, target)
    assert_lt(worst_rel(d, central(cross_entropy, logits, target)), 1.0e-6)
    # the gradient of a whole row sums to zero: probability has to go somewhere
    for i to len(d) {
        assert_near(sum(d[i, :]), 0.0, 1.0e-14)
    }
}

test "an embedding looks rows up, and its gradient adds into the rows it looked up" {
    W = nn_test_matrix(3, 4, 31.0)
    C = nn_test_matrix(3, 4, 43.0)
    X = embedding(W, [2.0, 0.0, 2.0])
    assert_near(max(abs(X[1, :] - W[0, :])), 0.0, 0.0)
    v, g = value_and_grad(embedding_probe, W, C)
    assert_lt(worst_rel(g, central(embedding_probe, W, C)), 1.0e-6)
    # row 2 was looked up twice, row 1 never
    assert_near(max(abs(g[2, :] - C[0, :] - C[2, :])), 0.0, 1.0e-15)
    assert_near(max(abs(g[1, :])), 0.0, 0.0)
}

# the sequences in the batch tests
NN_TEST_BATCH = 2

# a batch of test matrices, one per sequence, each from its own seed
def test_batch(rows, cols, seed) {
    S = zeros(NN_TEST_BATCH, rows, cols)
    for b to NN_TEST_BATCH {
        S[b, :, :] = nn_test_matrix(rows, cols, seed + b)
    }
    return S
}

def batch_attention_probe(X: Arr, C: Arr) {
    a = attention_case()
    return sum(C * attention(X, a.Wq, a.Wk, a.Wv))
}

test "on a batch of sequences, each sequence gets what it gets alone" {
    X = test_batch(5, 5, 7.0)
    a = attention_case()
    Y, A = attention(X, a.Wq, a.Wk, a.Wv), attention_weights(X, a.Wq, a.Wk)
    R = rmsnorm(X)
    P = softmax_rows(X)
    for b to len(X) {
        Xb = X[b, :, :]
        assert_near(max(abs(Y[b, :, :] - attention(Xb, a.Wq, a.Wk, a.Wv))), 0.0, 0.0)
        assert_near(max(abs(A[b, :, :] - attention_weights(Xb, a.Wq, a.Wk))), 0.0, 0.0)
        assert_near(max(abs(R[b, :, :] - rmsnorm(X[b, :, :]))), 0.0, 0.0)
        assert_near(max(abs(P[b, :, :] - softmax_rows(X[b, :, :]))), 0.0, 0.0)
    }
    # the batch's cross-entropy is the mean over every row of every sequence
    target = [[0.0, 1.0, 2.0, 3.0, 4.0], [4.0, 4.0, 0.0, 1.0, 2.0]]
    each = 0.0
    for b to len(X) {
        each = each + cross_entropy(X[b, :, :], target[b, :])
    }
    assert_near(cross_entropy(X, target), each / len(X), 1.0e-15)
    # an embedding of a matrix of ids is one row per id
    E = embedding(nn_test_matrix(3, 4, 31.0), [[2.0, 0.0], [1.0, 2.0]])
    assert_near(max(abs(E[1, 0, :] - nn_test_matrix(3, 4, 31.0)[1, :])), 0.0, 0.0)
}

test "attention's gradient on a batch matches central differences" {
    X = test_batch(5, 5, 7.0)
    C = test_batch(5, 3, 83.0)
    v, g = value_and_grad(batch_attention_probe, X, C)
    for b to len(X) {
        assert_lt(worst_rel(g[b, :, :], central_slab(X, C, b)), 1.0e-6)
    }
}

test "argmax picks the largest entry of each row, and the first on a tie" {
    rows = 3
    cols = 4
    logits = zeros(rows, cols)
    logits[0, 2] = 5.0
    logits[1, 0] = -1.0
    logits[1, 3] = 0.5
    # row 2 is all zeros: a four-way tie, won by column 0
    guess = argmax_rows(logits)
    assert_near(guess[0], 2.0, 0.0)
    assert_near(guess[1], 3.0, 0.0)
    assert_near(guess[2], 0.0, 0.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors