29. Softmax and Attention

What you will understand: how a Mountain Goat program computes softmax (turning scores into probabilities) and scaled dot-product attention, the operation at the heart of a transformer. The one thing the language lacked was an exponential, so this chapter adds a single new operation, exp, and then writes everything else with what already existed: the matrix product, reductions, broadcasting, transpose. Along the way it shows a real numerical trap (the textbook softmax produces nan on large scores), the one-line fix, causal masking, and a complete self-attention block, each checked against an independent Python implementation.

What you need to know first: Chapter 22 (reductions such as row_max and row_sum, and broadcasting a size-1 axis), Chapter 21 (@, transpose) and the language tour. Chapter 28 is not needed (everything here is rank 2).

Compile and run

cd docs/part29/code && ./build.sh                 # once: the newest compiler (adds mg.exp)
./mgc run examples/06_attention.mg                # compile and run one example
./mgc mlir examples/04_softmax_stable.mg          # show the MLIR the front end produces
./run_examples.sh > examples_out.txt              # every example and every error example on this page
./check_attention.py                              # softmax and attention against an independent Python version
./mutation.sh > mutation_out.txt                  # breaks the compiler on purpose, eight ways
cd ../../part15/code && ./run_lit.sh              # the whole test suite

Every listing and output on this page comes from these commands.

Why this page shows nan, FAIL lines and error messages, and why that is good

Example 3 is supposed to print inf and nan: it demonstrates why the textbook softmax formula is wrong on a computer. The four error examples are programs the compiler must reject (each ends exit status 1). The mutation output is a negative control: the compiler deliberately broken and run against the tests, where "caught" is the wanted result and "NOT CAUGHT" would be a test gap.

What is new: exp

exp(m) applies e^x to every element of a matrix, and nothing else is new in the language. Adding it touched the same three layers as every earlier operation (Chapter 22):

• the dialect: mg.exp in MgOps.td (one input, one result of the same shape) with the same verifier as mg.neg and mg.relu;
• the lowering: a loop nest that loads each element, calls MLIR's own math.exp, and stores the result. The math dialect is then converted to LLVM (--convert-math-to-llvm, a new step in mgc), which turns it into a call to the C library's exp, so mgc also links the math library (-lm);
• the front end: exp joins relu, row_sum and the others as a built-in function that needs a matrix.

Because it is an ordinary elementwise operation, exp works on the dynamic shapes of Chapter 14 too, though softmax itself does not (see the limits).

################ mgc run 01_exp.mg
  | # exp(x) applies e^x to every element. Its results are not whole numbers, so the output shows six significant digits.
  | let a = [[0, 1], [2, -1]]
  | print exp(a)                              # e^0 = 1, e^1 = 2.71828, e^2 = 7.38906, e^-1 = 0.367879
  | # exp turns sums into products: e^(x+y) = e^x * e^y. Both lines print the same numbers (to the printed digits).
  | let x = [[0.5, 1.5]]
  | let y = [[1, -2]]
  | print exp(x + y)
  | print exp(x) * exp(y)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data = 
[[1,   2.71828], 
 [7.38906,   0.367879]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[4.48169,   0.606531]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data = 
[[4.48169,   0.606531]]
  ----- exit status 0

The first matrix is e^0 = 1, e^1 = 2.71828, e^2 = 7.38906, e^-1 = 0.367879. The last two lines show the rule that makes softmax work: exp(x + y) and exp(x) * exp(y) are the same numbers (here [4.48169, 0.606531] for both), because e^x+y = e^x·e^y. A sum of scores becomes a product of positive numbers.

Softmax

Given a row of scores, softmax produces a row of positive numbers that add up to 1, bigger scores getting bigger shares:

softmax(row)_j = e^row_j / (e^row₀ + e^row₁ + …)

In Mountain Goat that is one line. exp(x) has shape 2x3 and row_sum(exp(x)) has shape 2x1 (one total per row); dividing the first by the second uses Chapter 22's broadcasting to stretch the total across the three columns:

################ mgc run 02_softmax_naive.mg
  | # softmax turns each row of scores into probabilities: positive numbers that add up to 1, bigger score = bigger share.
  | #   softmax(row)_j = e^(row_j) / (e^(row_0) + e^(row_1) + ...)
  | # exp(x) is 2x3, row_sum(...) is 2x1: the size-1 column is stretched across the three columns (Chapter 22's broadcasting).
  | def softmax_naive(x: tensor[2x3]) = exp(x) / row_sum(exp(x))
  | let scores = [[1, 2, 3], [0, 0, 0]]
  | print softmax_naive(scores)                        # row 0: [0.0900306, 0.244728, 0.665241]; row 1: three equal scores -> a third each
  | print row_sum(softmax_naive(scores))               # every row of probabilities sums to 1
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0.0900306,   0.244728,   0.665241], 
 [0.333333,   0.333333,   0.333333]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data = 
[[1], 
 [1]]
  ----- exit status 0

Row 0, scores 1 2 3, becomes 0.0900306, 0.244728, 0.665241: the score 3 takes about two thirds. Row 1, three equal scores, gives a third each. The second matrix is row_sum of the first and is exactly 1 in every row, as it must be.

The textbook formula breaks: overflow

The function above is called softmax_naive for a reason. A double-precision number cannot be larger than about 1.8 × 10³⁰⁸, and e⁷¹⁰ already exceeds it. Scores in the hundreds or thousands are normal in real models.

################ mgc run 03_softmax_overflow.mg
  | # THE GOTCHA the "naive" name was warning about. e^1000 is larger than any double can hold, so exp gives infinity,
  | # and infinity / infinity is "not a number" (nan). The formula is right on paper and wrong on a computer.
  | def softmax_naive(x: tensor[1x3]) = exp(x) / row_sum(exp(x))
  | print exp([[1000, 1001, 1002]])
  | print softmax_naive([[1000, 1001, 1002]])
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[inf,   inf,   inf]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[-nan,   -nan,   -nan]]
  ----- exit status 0

Read this output as the expected one. exp of 1000 is inf (infinity), so the formula computes inf / inf, which is "not a number": nan (the minus sign printed in front of nan is an accident of the hardware; a nan has no meaningful sign). The formula is right on paper and wrong in a computer. Notice that the compiler did nothing wrong, nothing aborted, and the program printed a result. Silent nan is the failure mode.

The fix: subtract the row maximum

Divide the top and bottom of the fraction by e^max (which changes nothing, because it cancels), and the formula becomes

softmax(row)_j = e^{row_j − max} / Σ e^{row_k − max}

Now every exponent is at most 0, so every exp is at most 1: nothing can overflow. The largest score contributes exactly e^0 = 1, so the denominator is at least 1: nothing divides by zero. A function is a single expression (no let inside), so the subexpression is written twice:

################ mgc run 04_softmax_stable.mg
  | # The fix: subtract each row's largest score first. That changes nothing mathematically (the factor e^(-max) cancels between the top
  | # and the bottom of the fraction) but it makes every exponent <= 0, so every exp is <= 1 and nothing can overflow.
  | def softmax(x: tensor[2x3]) = exp(x - row_max(x)) / row_sum(exp(x - row_max(x)))
  | print softmax([[1000, 1001, 1002], [1, 2, 3]])     # both rows have the same gaps between scores, so the same probabilities
  | print softmax([[1, 2, 3], [1, 2, 3]] + 100)        # adding 100 to every score changes nothing (shift invariance)
  | print softmax([[-1000, -1001, -1002], [0, 0, 0]])  # very negative scores: exp underflows to 0 harmlessly, the largest one still wins
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0.0900306,   0.244728,   0.665241], 
 [0.0900306,   0.244728,   0.665241]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0.0900306,   0.244728,   0.665241], 
 [0.0900306,   0.244728,   0.665241]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data = 
[[0.665241,   0.244728,   0.0900306], 
 [0.333333,   0.333333,   0.333333]]
  ----- exit status 0

Three observations, each visible in the output:

1. Scores 1000 1001 1002 and scores 1 2 3 give the same probabilities, 0.0900306, 0.244728, 0.665241, because only the gaps between scores matter.
2. Adding 100 to every score changes nothing (the second matrix): softmax is shift invariant, which is exactly why subtracting the maximum is allowed.
3. Very negative scores (the third matrix) are fine: the exponentials underflow to 0 harmlessly, and the largest score still gets the largest share. A row of zeros gives a third each.

What the compiler built for the stable version is exactly the recipe, and a test checks it (this is mgc mlir, abridged):

mg.reduce ... kind = "max" : tensor<2x3xf64> -> tensor<2x1xf64>      # the row maximum
mg.broadcast ...           : tensor<2x1xf64> -> tensor<2x3xf64>      # stretched across the columns
mg.sub ...                 : tensor<2x3xf64>, tensor<2x3xf64> -> tensor<2x3xf64>
mg.exp ...                 : tensor<2x3xf64> -> tensor<2x3xf64>
mg.reduce ... kind = "sum" : tensor<2x3xf64> -> tensor<2x1xf64>      # the row total
mg.div ...                 : tensor<2x3xf64>, tensor<2x3xf64> -> tensor<2x3xf64>

Temperature

Dividing the scores by a number T (a compile-time scalar) before the softmax sharpens it when T < 1 and flattens it when T > 1:

################ mgc run 05_temperature.mg
  | # Dividing the scores by a "temperature" T before the softmax sharpens it (T < 1) or flattens it (T > 1). T is a compile-time number.
  | def softmax(x: tensor[1x3]) = exp(x - row_max(x)) / row_sum(exp(x - row_max(x)))
  | let scores = [[1, 2, 3]]
  | print softmax(scores)               # T = 1
  | print softmax(scores / 0.5)         # T = 0.5: sharper, the top score takes about 87% (exp(4) against exp(2) and 1)
  | print softmax(scores / 4)           # T = 4: flatter, close to a third each
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[0.0900306,   0.244728,   0.665241]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[0.0158762,   0.11731,   0.866813]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data = 
[[0.254275,   0.326496,   0.419229]]
  ----- exit status 0

With T = 0.5 the top score takes 86.7% (the gaps double, from 1 and 2 to 2 and 4). With T = 4 the three shares are 0.254, 0.326, 0.419: nearly equal.

Attention

Attention lets each position in a sequence take a weighted average of every position's content. There are three matrices with one row per position: queries q (what each position is looking for), keys k (what each position offers) and values v (what each position hands over). The score of position i looking at position j is the dot product of q[i] and k[j], so all the scores at once are q @ transpose(k), an n x n matrix. Softmax turns each row of scores into weights, and the output is the weights times v. Dividing the scores by √d (where d is the width of a query) keeps their size from growing with d:

attention(q, k, v) = softmax(q · kᵀ / √d) · v

################ mgc run 06_attention.mg
  | # Scaled dot-product attention for one head. Row i of q asks a question, row j of k advertises what position j holds, row j of v is what it
  | # hands over. scores[i][j] = how well question i matches key j; softmax turns each row of scores into weights; the output row i is the
  | # weighted average of the rows of v.   attention(q, k, v) = softmax(q @ transpose(k) / sqrt(d)) @ v      (here d = 2)
  | def softmax(s: tensor[3x3]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
  | def weights(q: tensor[3x2], k: tensor[3x2]) = softmax(q @ transpose(k) * 0.7071067811865476)
  | def attention(q: tensor[3x2], k: tensor[3x2], v: tensor[3x2]) = weights(q, k) @ v
  | let q = [[1, 0], [0, 1], [1, 1]]
  | let k = [[1, 0], [0, 1], [1, 1]]
  | let v = [[10, 0], [0, 10], [5, 5]]
  | print weights(q, k)                      # 3x3, each row sums to 1
  | print attention(q, k, v)                 # 3x2: row i is the weights of row i times the three rows of v
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1] data = 
[[0.401112,   0.197776,   0.401112], 
 [0.197776,   0.401112,   0.401112], 
 [0.248255,   0.248255,   0.50349]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 
[[6.01668,   3.98332], 
 [3.98332,   6.01668], 
 [5,   5]]
  ----- exit status 0

Here q, k and v have three positions of width d = 2, and the scale is 1/√2 = 0.7071067811865476, a compile-time number. The first matrix is the attention weights: row 0 (0.401, 0.198, 0.401) says position 0 listens mostly to positions 0 and 2 (their keys, [1,0] and [1,1], point the same way as its query [1,0]) and least to position 1. The second matrix is the output: row 0 is 0.401·[10,0] + 0.198·[0,10] + 0.401·[5,5] = [6.017, 3.983], which you can check by hand. The functions call each other (attention calls weights, which calls softmax), and every call is shape-checked: the compiler knows softmax's argument must be 3x3.

Causal attention: not looking ahead

A model that generates text must not let position i see positions after it. The standard trick is to add a mask to the scores before the softmax: 0 where looking is allowed (j ≤ i) and a huge negative number where it is not. e^(-1000000000) is exactly 0, so masked positions get weight exactly 0 and the row still sums to 1:

################ mgc run 07_causal_attention.mg
  | # Causal (masked) attention: position i may only look at positions 0..i. Add a huge negative number to the scores above the diagonal;
  | # e^(-1000000000) is exactly 0, so those positions get weight 0 and the rows still sum to 1.
  | def softmax(s: tensor[3x3]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
  | def causal_weights(q: tensor[3x2], k: tensor[3x2], mask: tensor[3x3]) = softmax(q @ transpose(k) * 0.7071067811865476 + mask)
  | let mask = [[0, -1000000000, -1000000000], [0, 0, -1000000000], [0, 0, 0]]
  | let q = [[1, 0], [0, 1], [1, 1]]
  | let k = [[1, 0], [0, 1], [1, 1]]
  | let v = [[10, 0], [0, 10], [5, 5]]
  | print causal_weights(q, k, mask)                   # upper triangle exactly 0; row 0 is [1, 0, 0]: position 0 can only see itself
  | print causal_weights(q, k, mask) @ v               # row 0 of the output is row 0 of v unchanged
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1] data = 
[[1,   0,   0], 
 [0.330238,   0.669762,   0], 
 [0.248255,   0.248255,   0.50349]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 
[[10,   0], 
 [3.30238,   6.69762], 
 [5,   5]]
  ----- exit status 0

The weights are lower triangular. Row 0 is [1, 0, 0]: position 0 can only see itself, so the first output row is exactly row 0 of v ([10, 0]). Row 1 mixes only positions 0 and 1. Row 2 is the same as the unmasked one, because the last position was allowed to see everything.

A self-attention block

In a transformer the queries, keys and values are not given; they are made from the same input by three learned matrices. One more matrix maps the attention output back to the input's width, and the block's result is added to its input (a residual connection), so each token keeps what it had and gains what it attended to:

################ mgc run 08_self_attention_block.mg
  | # Self-attention from token vectors: one matrix of 3 tokens (rows) of 4 numbers each is projected into queries, keys and values by three
  | # learned 4x2 matrices, attended over, and projected back to 4 numbers per token by a 2x4 matrix. The block's output is added to its input
  | # (a "residual connection"), so every token keeps what it had and gains what it attended to. (A def is one line, so it is long.)
  | def softmax(s: tensor[3x3]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
  | def self_attention(x: tensor[3x4], wq: tensor[4x2], wk: tensor[4x2], wv: tensor[4x2], wo: tensor[2x4]) = x + softmax((x @ wq) @ transpose(x @ wk) * 0.7071067811865476) @ (x @ wv) @ wo
  | let x  = [[1, 0, 1, 0], [0, 1, 0, 0], [1, 1, 0, 1]]
  | let wq = [[1, 0], [0, 1], [1, 0], [0, 1]]
  | let wk = [[1, 0], [0, 1], [0, 1], [1, 0]]
  | let wv = [[1, 0], [0, 1], [1, 1], [0, 0]]
  | let wo = [[1, 0, 0, 1], [0, 1, 1, 0]]
  | print self_attention(x, wq, wk, wv, wo)
  ----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 4] strides = [4, 1] data = 
[[2.14131,   1,   2,   1.14131], 
 [1,   2,   1,   1], 
 [2.14397,   2,   1,   2.14397]]
  ----- exit status 0

Three tokens of four numbers become queries, keys and values of width two (x @ wq, x @ wk, x @ wv), are attended over, and projected back to width four by wo, then added to x. The definition is one long line because a def is one line. The result was checked against an independent Python calculation (the attention weights are different for each token here: 0.187, 0.045, 0.768 for token 0, a third each for token 1 whose query matches every key equally, and 0.284, 0.140, 0.576 for token 2).

Mistakes the compiler reports

################ mgc run errors/e1_wrong_sequence_length.mg
  | # attention was defined for 3 tokens of width 2. Passing 4 tokens does not fit the parameter shapes, and the compiler says so.
  | def softmax(s: tensor[3x3]) = exp(s - row_max(s)) / row_sum(exp(s - row_max(s)))
  | def weights(q: tensor[3x2], k: tensor[3x2]) = softmax(q @ transpose(k))
  | let four = [[1, 0], [0, 1], [1, 1], [2, 2]]
  | print weights(four, four)
  ----- output
errors/e1_wrong_sequence_length.mg:5: error: argument shape 4x2 does not fit parameter 3x2
  ----- exit status 1

################ mgc run errors/e2_exp_of_a_scalar.mg
  | # exp is applied to a matrix. A bare number in the source is a compile-time scalar with no shape to print.
  | print exp(2)
  ----- output
errors/e2_exp_of_a_scalar.mg:2: error: exp needs a matrix, not a scalar
  ----- exit status 1

################ mgc run errors/e3_q_and_k_widths_differ.mg
  | # q @ transpose(k) needs q and k to have the same width (the inner sizes of the product). Here q is 3x2 and k is 3x3.
  | let q = [[1, 0], [0, 1], [1, 1]]
  | let k = [[1, 0, 0], [0, 1, 0], [1, 1, 1]]
  | print q @ transpose(k)
  ----- output
errors/e3_q_and_k_widths_differ.mg:4: error: cannot multiply (matmul) shapes 3x2 and 3x3: inner dimensions differ
  ----- exit status 1

################ mgc run errors/e4_softmax_over_a_variable_length.mg
  | # A softmax over a sequence of any length would need x - row_max(x) with a '?' row count, and a '?' size is never stretched (Chapter 22):
  | # the compiler refuses it. Attention in this chapter is for a fixed sequence length.
  | def softmax(x: tensor[?x3]) = exp(x - row_max(x)) / row_sum(exp(x - row_max(x)))
  | print softmax([[1, 2, 3]])
  ----- output
errors/e4_softmax_over_a_variable_length.mg:3: error: broadcasting needs static shapes; this operand has a '?' dimension
  ----- exit status 1

In order: passing four tokens to a function defined for three; exp of a bare number (a compile-time scalar has no shape); a query width that does not match the key width in q @ transpose(k); and softmax over a sequence of any length. The last one is a real limit, not a bug: x - row_max(x) with a ? row count needs a ? size to be stretched, and the language never stretches a ? size (Chapter 22). Attention in this chapter is for a fixed sequence length.

Tests

Five new lit files in test/attention29/ (the suite is now 117 tests, counting the tour's new test):

Test	What it checks
examples	the eight examples print the values on this page (including inf and nan for example 3)
errors	the four mistakes give the exact messages, with file and line
verifiers	mg.exp given a result of a different shape is rejected by the dialect, whatever front end built it
ir	the stable softmax is a max-reduce, broadcast, subtract, exp, sum-reduce, divide, in that order
against-reference	check_attention.py: softmax of eight score matrices, including scores of ±500 and of −1500 to −1000 (where the naive formula gives nan), attention and causal attention for five sizes, against an independent Python implementation, plus the properties below

The properties that must hold whatever the numbers are, checked on every case: each softmax row is positive and sums to 1; adding a constant to every score changes nothing; causal weights above the diagonal are exactly 0. Values are compared with a relative tolerance of 10^-5, because mgc prints six significant digits (Chapter 28's contractions could be compared exactly because integers and quarters are exact in binary; an exponential is not).

-- softmax against the reference, rows sum to 1, shift invariance
  ok   softmax 1x1, scores in [-3, 3]: equals the reference
  ok   softmax 1x1, scores in [-3, 3]: rows are positive and sum to 1
  ok   softmax 1x1, scores in [-3, 3]: adding 37 to every score changes nothing
  ok   softmax 1x5, scores in [-3, 3]: equals the reference
  ok   softmax 1x5, scores in [-3, 3]: rows are positive and sum to 1
  ok   softmax 1x5, scores in [-3, 3]: adding 37 to every score changes nothing
  ok   softmax 4x4, scores in [-10, 10]: equals the reference
  ok   softmax 4x4, scores in [-10, 10]: rows are positive and sum to 1
  ok   softmax 4x4, scores in [-10, 10]: adding 37 to every score changes nothing
  ok   softmax 3x7, scores in [-20, 20]: equals the reference
  ok   softmax 3x7, scores in [-20, 20]: rows are positive and sum to 1
  ok   softmax 3x7, scores in [-20, 20]: adding 37 to every score changes nothing
  ok   softmax 2x6, scores in [480, 500]: equals the reference
  ok   softmax 2x6, scores in [480, 500]: rows are positive and sum to 1
  ok   softmax 2x6, scores in [480, 500]: adding 37 to every score changes nothing
  ok   softmax 2x6, scores in [-500, -480]: equals the reference
  ok   softmax 2x6, scores in [-500, -480]: rows are positive and sum to 1
  ok   softmax 2x6, scores in [-500, -480]: adding 37 to every score changes nothing
  ok   softmax 2x6, scores in [-1500, -1000]: equals the reference
  ok   softmax 2x6, scores in [-1500, -1000]: rows are positive and sum to 1
  ok   softmax 2x6, scores in [-1500, -1000]: adding 37 to every score changes nothing
  ok   softmax 5x3, scores in [-500, 500]: equals the reference
  ok   softmax 5x3, scores in [-500, 500]: rows are positive and sum to 1
  ok   softmax 5x3, scores in [-500, 500]: adding 37 to every score changes nothing
-- attention (n tokens of width d) against the reference, with and without a causal mask
  ok   attention n=1 d=1: output equals the reference
  ok   causal attention n=1 d=1: weights and output equal the reference
  ok   causal attention n=1 d=1: weights above the diagonal are exactly 0, rows sum to 1
  ok   attention n=2 d=3: output equals the reference
  ok   causal attention n=2 d=3: weights and output equal the reference
  ok   causal attention n=2 d=3: weights above the diagonal are exactly 0, rows sum to 1
  ok   attention n=3 d=2: output equals the reference
  ok   causal attention n=3 d=2: weights and output equal the reference
  ok   causal attention n=3 d=2: weights above the diagonal are exactly 0, rows sum to 1
  ok   attention n=4 d=4: output equals the reference
  ok   causal attention n=4 d=4: weights and output equal the reference
  ok   causal attention n=4 d=4: weights above the diagonal are exactly 0, rows sum to 1
  ok   attention n=5 d=3: output equals the reference
  ok   causal attention n=5 d=3: weights and output equal the reference
  ok   causal attention n=5 d=3: weights above the diagonal are exactly 0, rows sum to 1
all checks pass

Are the tests good enough? Break the compiler eight ways.

#!/bin/sh
# READ THIS FIRST: this script breaks the Chapter 29 code on purpose, one way at a time, and re-runs the Chapter 29 tests (test/attention29).
#   "caught by ..." (a test failed) is the EXPECTED, wanted result: it shows the tests can notice that mistake.
#   "NOT CAUGHT" would be a test gap. Part 1 edits the front end (mgfront.py) and the driver (mgc) in place; part 2 rebuilds mg-opt with one bug
#   injected into the dialect or the lowering (a full rebuild each, in work/mut/N). Every edited file is restored afterwards. Output: mutation_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code/run_lit.sh; W=$HERE/work/mut; rm -rf $W; mkdir -p $W
FILTER=':: attention29/'
echo "NOTE: this script deliberately breaks the compiler code. 'caught' lines are EXPECTED: they show the tests can detect the mistake."
F=$HERE/mgfront.py; G=$HERE/mgc; cp $F $W/mgfront.orig; cp $G $W/mgc.orig
restore() { cp $W/mgfront.orig $F; cp $W/mgc.orig $G; }
trap restore EXIT
fails_of() { grep '^FAIL' | sed 's/.*:: //; s/ (.*//' | tr '\n' ' '; }
out=$(MG_TEST_TMP=$W/out sh $LIT --filter "$FILTER" 2>&1); echo "baseline (nothing broken): $(echo "$out" | grep -q '^FAIL' && echo 'UNEXPECTED FAILURE' || echo 'tests pass (expected)')"
echo; echo "== Part 1: the front end and the driver"
fe() {  # label file sed-expr
  restore; sed -i "$3" $2
  if cmp -s $F $W/mgfront.orig && cmp -s $G $W/mgc.orig; then echo "$1: MUTATION DID NOT APPLY"; return; fi
  f=$(MG_TEST_TMP=$W/out sh $LIT --filter "$FILTER" 2>&1 | fails_of)
  if [ -n "$f" ]; then echo "$1: caught by $f"; else echo "$1: NOT CAUGHT (a test gap)"; fi
  restore
}
fe "exp is compiled as relu by the front end"                            $F 's/{r} = mg.exp {v}/{r} = mg.relu {v}/'
fe "the driver forgets to lower the math dialect"                        $G 's/--convert-math-to-llvm //'
fe "the driver does not link the math library (-lm)"                     $G 's/ -lm / /'
restore

echo; echo "== Part 2: the dialect verifier and the lowering (a full rebuild per mutation)"
n=0
try() {  # label file sed-expr
  n=$((n+1)); d=$W/$n; mkdir -p $d/build; cp -r $HERE/tree $d/tree; rm -rf $d/tree/build
  tgt=$d/tree/lib/$2; cp $tgt $d/before; sed -i "$3" $tgt
  if cmp -s $d/before $tgt; then echo "$1: MUTATION DID NOT APPLY"; return; fi
  ( cd $d/build && cmake $d/tree -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF >/dev/null 2>&1 && mkdir -p include/mg && make -j8 >/dev/null 2>&1 )
  [ -x $d/build/mg-opt ] || { echo "$1: BUILD FAILED"; return; }
  f=$(MG_OPT=$d/build/mg-opt MG_TEST_TMP=$d/out sh $LIT --filter "$FILTER" 2>&1 | fails_of)
  if [ -n "$f" ]; then echo "$1: caught by $f"; else echo "$1: NOT CAUGHT (a test gap)"; fi
}
try "exp computes e^(-x)"                                               LowerToAffine.cpp 's/b.create<math::ExpOp>(l, x)/b.create<math::ExpOp>(l, b.create<arith::NegFOp>(l, x))/'
try "exp copies its input (computes nothing)"                           LowerToAffine.cpp 's/b.create<math::ExpOp>(l, x)/x/'
try "mg.exp's verifier accepts anything"                                MgDialect.cpp 's/return verifyElementwise(\*this, {getInput().getType(), getResult().getType()}, "mg.exp"); }/return mlir::success(); }/'
try "a row maximum is computed as a minimum (Chapter 22's reduce)"       LowerToAffine.cpp 's/Value(b.create<arith::MaximumFOp>(l, acc, x))/Value(b.create<arith::MinimumFOp>(l, acc, x))/'
try "a maximum starts from 0, not -infinity (wrong for all-negative rows)" LowerToAffine.cpp 's/double initVal = isMax ? -std::numeric_limits<double>::infinity() : 0.0;/double initVal = 0.0;/'
restore
echo; echo "restored: $(cmp -s $F $W/mgfront.orig && cmp -s $G $W/mgc.orig && echo 'mgfront.py and mgc are back to their originals' || echo 'RESTORE FAILED')"

NOTE: this script deliberately breaks the compiler code. 'caught' lines are EXPECTED: they show the tests can detect the mistake.
baseline (nothing broken): tests pass (expected)

== Part 1: the front end and the driver
exp is compiled as relu by the front end: caught by attention29/ir.mlir attention29/examples.mlir attention29/against-reference.mlir 
the driver forgets to lower the math dialect: caught by attention29/examples.mlir attention29/against-reference.mlir 
the driver does not link the math library (-lm): caught by attention29/examples.mlir attention29/against-reference.mlir 

== Part 2: the dialect verifier and the lowering (a full rebuild per mutation)
exp computes e^(-x): caught by attention29/examples.mlir attention29/against-reference.mlir 
exp copies its input (computes nothing): caught by attention29/examples.mlir attention29/against-reference.mlir 
mg.exp's verifier accepts anything: caught by attention29/verifiers.mlir 
a row maximum is computed as a minimum (Chapter 22's reduce): caught by attention29/examples.mlir attention29/against-reference.mlir 
a maximum starts from 0, not -infinity (wrong for all-negative rows): caught by attention29/examples.mlir attention29/against-reference.mlir 

restored: mgfront.py and mgc are back to their originals

All eight are caught. Two details are worth noticing.

• "A maximum starts from 0, not -infinity" is caught only because a test uses scores below −745. Scores such as −490 survive the mistake: e^−490 is tiny but not zero, so the answer stays right. Only when the scores are so negative that the exponentials underflow (the [-1500, -1000] case here, and example 4's -1000 row) does the wrong starting value produce 0/0. An earlier version of check_attention.py used only [-500, -480] and this mutation was caught by the examples alone; the extra case was added after the first run of this script showed that, and the script was re-run in full.
• Dropping -lm is caught. Without the math library the link step fails, so mgc produces nothing and every example fails. It would have been an equivalent mutant had the C library's exp lived in the main C library on this platform; it does not.

The full suite:

-- Testing: 117 tests, 4 workers --
PASS: IR Forge / Mountain Goat :: tour/gallery.mlir (1 of 117)
PASS: IR Forge / Mountain Goat :: attention29/against-reference.mlir (2 of 117)
PASS: IR Forge / Mountain Goat :: contraction/tensor-contraction.mlir (3 of 117)
PASS: IR Forge / Mountain Goat :: attention29/examples.mlir (4 of 117)
PASS: IR Forge / Mountain Goat :: learning/gradient-descent.mlir (5 of 117)
PASS: IR Forge / Mountain Goat :: contraction28/examples.mlir (6 of 117)
PASS: IR Forge / Mountain Goat :: contraction28/against-definition.mlir (7 of 117)
PASS: IR Forge / Mountain Goat :: tour/more-examples.mlir (8 of 117)
PASS: IR Forge / Mountain Goat :: perf/options-same-results.mlir (9 of 117)
PASS: IR Forge / Mountain Goat :: tour/tour-examples.mlir (10 of 117)
PASS: IR Forge / Mountain Goat :: loop-order/results-identical.mlir (11 of 117)
PASS: IR Forge / Mountain Goat :: contraction/mismatch-trap.mlir (12 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/values.mlir (13 of 117)
PASS: IR Forge / Mountain Goat :: loop-order/bitwise-identical.mlir (14 of 117)
PASS: IR Forge / Mountain Goat :: ops/values.mlir (15 of 117)
PASS: IR Forge / Mountain Goat :: ops/elementwise-mismatch-aborts.mlir (16 of 117)
PASS: IR Forge / Mountain Goat :: ops/cpp-interop-ops.mlir (17 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/cpp-interop-reductions.mlir (18 of 117)
PASS: IR Forge / Mountain Goat :: frontend/more-examples.mlir (19 of 117)
PASS: IR Forge / Mountain Goat :: tour/scientific-literals.mlir (20 of 117)
PASS: IR Forge / Mountain Goat :: loop-order/cpp-ikj.mlir (21 of 117)
PASS: IR Forge / Mountain Goat :: frontend/cpp-interop.mlir (22 of 117)
PASS: IR Forge / Mountain Goat :: execution/gpu-host-with-stub-runtime.mlir (23 of 117)
PASS: IR Forge / Mountain Goat :: frontend/runtime-mismatch-aborts.mlir (24 of 117)
PASS: IR Forge / Mountain Goat :: perf/flags-reach-clang.mlir (25 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/dynamic-reductions.mlir (26 of 117)
PASS: IR Forge / Mountain Goat :: perf/tiling-structure.mlir (27 of 117)
PASS: IR Forge / Mountain Goat :: frontend/dynamic-runs.mlir (28 of 117)
PASS: IR Forge / Mountain Goat :: ops/dynamic-ops.mlir (29 of 117)
PASS: IR Forge / Mountain Goat :: loop-order/flag-reaches-pass.mlir (30 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/identities.mlir (31 of 117)
PASS: IR Forge / Mountain Goat :: ops/composed-programs.mlir (32 of 117)
PASS: IR Forge / Mountain Goat :: frontend/hello-and-static.mlir (33 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/networks.mlir (34 of 117)
PASS: IR Forge / Mountain Goat :: ops/negate-and-dynamic-divide.mlir (35 of 117)
PASS: IR Forge / Mountain Goat :: execution/mismatch-columns-aborts.mlir (36 of 117)
PASS: IR Forge / Mountain Goat :: execution/dynamic-shapes-run.mlir (37 of 117)
PASS: IR Forge / Mountain Goat :: contraction28/errors.mlir (38 of 117)
PASS: IR Forge / Mountain Goat :: ops/precedence-and-identities.mlir (39 of 117)
PASS: IR Forge / Mountain Goat :: ops/ptx-new-ops.mlir (40 of 117)
PASS: IR Forge / Mountain Goat :: execution/static-add-runs.mlir (41 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-on-stderr.mlir (42 of 117)
PASS: IR Forge / Mountain Goat :: execution/mismatch-aborts.mlir (43 of 117)
PASS: IR Forge / Mountain Goat :: tour/tour-errors.mlir (44 of 117)
PASS: IR Forge / Mountain Goat :: ops/matmul-mismatch-aborts.mlir (45 of 117)
PASS: IR Forge / Mountain Goat :: chapters/ch09-execution-engine.mlir (46 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/runtime-check-survives-transforms.mlir (47 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-runs.mlir (48 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-runs.mlir (49 of 117)
PASS: IR Forge / Mountain Goat :: attention29/errors.mlir (50 of 117)
PASS: IR Forge / Mountain Goat :: ops/division-ieee.mlir (51 of 117)
PASS: IR Forge / Mountain Goat :: execution/mixed-static-dynamic.mlir (52 of 117)
PASS: IR Forge / Mountain Goat :: execution/strided-view.mlir (53 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-after-tiling.mlir (54 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/nested-assert-runs.mlir (55 of 117)
PASS: IR Forge / Mountain Goat :: chapters/ch01-pipeline.mlir (56 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/valid-shapes-still-pass.mlir (57 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/message-bytes-exact.mlir (58 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/ptx-layer.mlir (59 of 117)
PASS: IR Forge / Mountain Goat :: frontend/ptx-from-source.mlir (60 of 117)
PASS: IR Forge / Mountain Goat :: contraction28/ir.mlir (61 of 117)
PASS: IR Forge / Mountain Goat :: gpu/pass-order-crash.mlir (62 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-needs-constant-trip-counts.mlir (63 of 117)
PASS: IR Forge / Mountain Goat :: gpu/ptx-bare-pointers.mlir (64 of 117)
PASS: IR Forge / Mountain Goat :: gpu/ptx-index32.mlir (65 of 117)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls.mlir (66 of 117)
PASS: IR Forge / Mountain Goat :: gpu/host-runtime-calls-bare.mlir (67 of 117)
PASS: IR Forge / Mountain Goat :: gpu/ptx-default.mlir (68 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-full-needs-constant-trip-count.mlir (69 of 117)
PASS: IR Forge / Mountain Goat :: loop-order/structure.mlir (70 of 117)
PASS: IR Forge / Mountain Goat :: frontend/front-end-errors.mlir (71 of 117)
PASS: IR Forge / Mountain Goat :: ops/matmul-lowering-structure.mlir (72 of 117)
PASS: IR Forge / Mountain Goat :: ops/front-end-errors.mlir (73 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/front-end-errors.mlir (74 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/roundtrip-and-lowering.mlir (75 of 117)
PASS: IR Forge / Mountain Goat :: chapters/ch02-example-programs.mlir (76 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/guard-nested-assert.mlir (77 of 117)
PASS: IR Forge / Mountain Goat :: gpu/cuda-comparison.mlir (78 of 117)
PASS: IR Forge / Mountain Goat :: loops/fusion.mlir (79 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/static-unchanged.mlir (80 of 117)
PASS: IR Forge / Mountain Goat :: gpu/nvvm-lowering.mlir (81 of 117)
PASS: IR Forge / Mountain Goat :: ops/elementwise-static-mismatch.mlir (82 of 117)
PASS: IR Forge / Mountain Goat :: gpu/kernel-outlining.mlir (83 of 117)
PASS: IR Forge / Mountain Goat :: attention29/ir.mlir (84 of 117)
PASS: IR Forge / Mountain Goat :: frontend/emit-mlir.mlir (85 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/tiling-structure.mlir (86 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-declines.mlir (87 of 117)
PASS: IR Forge / Mountain Goat :: ops/scalar-bad-op.mlir (88 of 117)
PASS: IR Forge / Mountain Goat :: loops/tiling.mlir (89 of 117)
PASS: IR Forge / Mountain Goat :: ops/dynamic-accepted.mlir (90 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/rewrite-structure.mlir (91 of 117)
PASS: IR Forge / Mountain Goat :: verifier/add-static-mismatch.mlir (92 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/unroll-epilogue.mlir (93 of 117)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-add-affine.mlir (94 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-broadcast.mlir (95 of 117)
PASS: IR Forge / Mountain Goat :: attention29/verifiers.mlir (96 of 117)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice.mlir (97 of 117)
PASS: IR Forge / Mountain Goat :: dynamic-loops/fusion-fires-when-static.mlir (98 of 117)
PASS: IR Forge / Mountain Goat :: canonicalize/fold-constant-add.mlir (99 of 117)
PASS: IR Forge / Mountain Goat :: contraction28/verifiers.mlir (100 of 117)
PASS: IR Forge / Mountain Goat :: verifier/transpose-rank3.mlir (101 of 117)
PASS: IR Forge / Mountain Goat :: verifier/transpose-bad-shape.mlir (102 of 117)
PASS: IR Forge / Mountain Goat :: verifier/unranked-rejected.mlir (103 of 117)
PASS: IR Forge / Mountain Goat :: assert-stderr/run-twice-no-collision.mlir (104 of 117)
PASS: IR Forge / Mountain Goat :: loop-order/bad-option.mlir (105 of 117)
PASS: IR Forge / Mountain Goat :: loops/unroll-full.mlir (106 of 117)
PASS: IR Forge / Mountain Goat :: broadcast/verifier-reduce.mlir (107 of 117)
PASS: IR Forge / Mountain Goat :: ops/matmul-inner-mismatch.mlir (108 of 117)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-dynamic-add.mlir (109 of 117)
PASS: IR Forge / Mountain Goat :: lowering/dynamic-transpose-no-check.mlir (110 of 117)
PASS: IR Forge / Mountain Goat :: lowering/mixed-add-one-check.mlir (111 of 117)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-exact.mlir (112 of 117)
PASS: IR Forge / Mountain Goat :: lowering/static-add-affine.mlir (113 of 117)
PASS: IR Forge / Mountain Goat :: verifier/dynamic-accepted.mlir (114 of 117)
PASS: IR Forge / Mountain Goat :: canonicalize/transpose-twice-dynamic-mixed.mlir (115 of 117)
PASS: IR Forge / Mountain Goat :: lowering/bufferize-static-add.mlir (116 of 117)
PASS: IR Forge / Mountain Goat :: ops/matmul-bad-result-shape.mlir (117 of 117)

Testing Time: 10.97s

Total Discovered Tests: 117
  Passed: 117 (100.00%)

Limits and what is not established

• Fixed sequence length. Softmax and attention are written for static shapes. A function over tensor[?x3] is rejected (example error 4) because a ? size is never stretched. Attention over a sequence of any length is not available.
• One head, no batch. Multi-head attention and a batch of sequences need a batched matrix product (an axis shared by both inputs and the output). Chapter 28's contract does not support batch axes, and nothing here adds them. A multi-head block can be written as separate heads added together (each head's output through its own slice of the output projection), which is mathematically the same; that is not demonstrated here.
• No normalization layers, no feed-forward network, no training. A transformer has more than attention. This chapter stops at the attention block with a residual connection.
• Numbers printed to six significant digits. All comparisons with Python are to a relative 10^-5. The compiled exp calls the platform C library's exp; its accuracy (typically within one unit in the last place) was not measured here, and a different C library could print different last digits.
• A large negative number is not minus infinity. The mask uses -1000000000. A score of that size in real data would be masked incorrectly; exp of it is exactly 0 only because it underflows. (The language has no literal for infinity.)
• The stable softmax computes the exponentials twice (a def cannot name a subexpression), so it does about twice the exp work it needs to. Speed was not measured, and the compiler does not eliminate the duplicate.
• exp on the GPU path is not covered. mgc ptx was not run with exp; the math dialect's route to PTX was not tried.
• Overflow is demonstrated, not defended. Nothing in the language prevents a user writing the naive softmax; example 3 shows what happens. A mistake that produces nan produces no error.
• Platform: x86-64 Linux, LLVM 18.1.3.

Reproducing

cd docs/part29/code && ./build.sh
./run_examples.sh > examples_out.txt
./check_attention.py > check_attention_out.txt
./mutation.sh > mutation_out.txt
cd ../../part15/code && ./run_lit.sh          # 117 tests

Chapter summary

• One new operation, exp, was enough to write softmax and scaled dot-product attention in Mountain Goat: everything else was already there.
• The textbook softmax overflows to nan for scores in the hundreds; subtracting the row maximum fixes it without changing the answer (shift invariance), and the examples show both.
• Attention is softmax(q @ transpose(k) * scale) @ v; a causal mask is an added matrix of 0 and a huge negative number; a self-attention block adds three projections and a residual connection.
• Everything agrees with an independent Python version to the six printed digits, including scores of plus and minus 500, and the structural properties (rows sum to 1, masked weights exactly 0) hold on every case.
• Static shapes only, one head, no batching: the transformer chapter that follows has to deal with the limits listed above.

Self-check questions

Each answer is collapsed; try the question first.

1. Why is softmax_naive([[1000, 1001, 1002]]) all nan rather than a plausible wrong answer?

   Answer

   exp(1000) exceeds the largest double and becomes infinity, and so do exp(1001) and exp(1002). The row total is infinity, and each element is infinity divided by infinity, which is undefined, so the result is not-a-number. The stable version never makes a number bigger than 1, so it cannot overflow.

2. Why does subtracting the row maximum not change the result?

   Answer

   e^{x − m} = e^x·e^−m, so every term of the numerator and of the denominator is multiplied by the same factor e^−m, which cancels in the fraction. This is the "shift invariance" that example 4's second matrix shows directly.

3. In exp(x) / row_sum(exp(x)) the shapes are 2x3 and 2x1. What makes that legal, and what would stop it if the row count were ??

   Answer

   Chapter 22's broadcasting: a static dimension of size 1 is stretched to match the other operand (a 2x1 becomes 2x3). A ? size is not known when compiling, so it cannot be stretched, and the compiler rejects the program (example error 4). The same rule is why attention here needs a fixed length.

4. What do the three matrices q, k, v each contribute, and what shape is q @ transpose(k) for n positions of width d?

   Answer

   Queries say what each position is looking for, keys say what each position offers, values say what each position hands over. For q and k of shape n x d, transpose(k) is d x n, so q @ transpose(k) is n x n: entry [i][j] is the dot product of query i and key j, the score of position i looking at position j.

5. Why divide the scores by the square root of d?

   Answer

   A dot product of two vectors of width d with entries of ordinary size grows in size with d (the typical size of a sum of d products grows like √d). Large scores push the softmax toward giving almost all the weight to one position. Dividing by √d keeps the scores about the same size whatever d is. In the examples d = 2 and the scale is 1/√2 = 0.7071067811865476, written as a compile-time number.

6. The causal mask uses -1000000000. Why does a masked position get weight exactly 0, not a tiny positive number, and when would this go wrong?

   Answer

   After the maximum is subtracted, a masked score is about −10⁹ below the largest, and exp of that is far smaller than the smallest positive double, so it underflows to exactly 0. It would go wrong if real scores were as large as the mask (or if a whole row were masked, which makes a row of equal very negative numbers and a meaningless result); the language has no literal for −infinity, which is the usual choice.

7. Why can row 0 of the causal output in example 7 be read off without computing anything?

   Answer

   The only position row 0 may attend to is position 0, so its weights are [1, 0, 0] and its output is 1·v[0] + 0·v[1] + 0·v[2] = v[0]: [10, 0].

8. Why does check_attention.py compare with a tolerance when Chapter 28's check compared exactly?

   Answer

   Chapter 28 used whole numbers and multiples of a quarter, whose sums and products are exact in binary, so any difference at all meant a bug. An exponential's result is almost never exactly representable, and mgc prints only six significant digits, so exact comparison is impossible; a relative tolerance of 10^-5 is a little looser than the printing precision.

9. Which mutation in mutation.sh is caught only because a test uses very negative scores?

   Answer

   "A maximum starts from 0, not −infinity" gives the right answer for any row with a non-negative score, so only an all-negative row can expose it, and even scores around −490 survive it (e^−490 is tiny but not zero). It takes scores below about −745, where exp underflows to exactly 0, to turn the wrong maximum into 0/0: the [−1500, −1000] case and example 4's third matrix. The lesson is the same as Chapter 22's 13_negative_max example: a test with only friendly numbers cannot see a wrong starting value, and a somewhat unfriendly one may not either.

10. What would multi-head attention need that this chapter's language does not have?

    Answer

    A batched matrix product: the head index appears in both inputs and in the output, which is a batch axis, and Chapter 28's contract rejects that (an axis must be paired and summed away). Without it, heads can be written as separate functions and their outputs combined by adding each head's product with its own slice of the output matrix, which is the same arithmetic but not a single batched operation.