$cat ~/posts/bend-ml

bilingual · en / pt

← /blog

a pytorch where a shape error doesn't compile

2026-10-04·en·47 min read·bendlanguagesmlproofsperformance

On September 29, Taelin, the creator of Bend, published a long info dump about the language's future. His company was raising a funding round, was going to set up two new teams, and in the middle of the text there was a wishlist for the ecosystem. One item on the list read: "AI framework / PyTorch / llama.bend / etc.".

Someone asked, in the replies, what a person should do to join one of those teams. The answer was short: build something in Bend, something that shows investors traction and impresses the team.

I had already written about the language's idea and built the editor extensions. What was missing was using it for real, on something big. So I took that wishlist item and spent October 3 and 4 on it. The result is called bend-ml: a machine learning library in Bend, five published packages, a handwritten-digit classifier and GPT-2 running, all with results checked number by number against PyTorch.

This post is the whole story. The first half assumes you know nothing about machine learning, types or proofs: I explain each piece before using it. The second half goes down to the code, the experiments and the numbers, including the numbers where the project loses badly.

the project in 60 seconds. the video is generated, not recorded: every terminal output and every number comes from the real runs.
  1. 2026-09-19the manifestoI write about bend2's bet without having run a single line.
  2. 2026-09-29the wishlisttaelin publishes the info dump: "AI framework / PyTorch / llama.bend".
  3. 2026-10-03v0.1 to v1.0lemmas, tokenizer, tensors, autograd, MNIST and GPT-2. all correct, all slow.
  4. 2026-10-04v2 and v2.1the performance experiments, the GPU, CI and the missing proofs.
each bar is the gap until the next milestone. the whole project fit in the last two rows.

I have no background in type theory or formal proofs. My day job is AI engineering, with agents and NLP, but on top of the models: their insides, the training, the gradients and the tensors, are something I have been studying on the side. The project was also my crash course in all of it. When I say "proved" in this post, I am not the one vouching for it: the language's checker is, and anyone can run it again.

what training a model means

Before any code, it helps to understand what a framework like PyTorch does, because that is what I wanted to rebuild.

A machine learning model is a function with many knobs. You give it an input (a photo of a digit, a piece of text) and it returns an output (which digit it is, which word comes next). The knobs are numbers, called weights, and every combination of weights gives a different function. Training is turning those knobs, little by little, until the function gets things right.

The way you turn them is always the same. You show an example, measure how wrong the function was, figure out which way each knob should turn to be less wrong, and turn it a tiny bit. Then repeat with the next example, millions of times. That "which way to turn" is the gradient, and I will come back to it shortly.

The weights are not loose: they are organized in tables. A table of numbers with rows and columns is a matrix; when it has more dimensions, the generic name is tensor. And the operation a model does all the time, billions of times, is multiplying matrices.

multiplying matrices, and the shape error

Multiplying a matrix A by a matrix B works like this: for each row of A and each column of B, you multiply the numbers one by one and add everything up. The result is one number of matrix C.

each number of C comes from one row of A and one column of B. that is why A's rows and B's columns must have the same length.

Notice the hidden rule: to combine a row of A with a column of B number by number, both must have the same length. In other words, the number of columns of A must equal the number of rows of B. A 2×3 matrix multiplies a 3×4 one, but not a 4×5 one. That pair of numbers, rows by columns, is the matrix's shape.

In a real model, shapes flow through dozens of layers, each transforming the previous one's shape. And getting a shape wrong is the most common bug when writing a model. In PyTorch, the error shows up like this, at run time, after the program is already running:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x3 and 4x5)

Sometimes it shows up in seconds. Sometimes it shows up after hours of training, in a layer that only runs at the end of an epoch. The question that drove the project was: what if that error could never get to run?

what a type is

Almost every language has the idea of a type: the label that says what kind of thing a value is. 3 is an integer, "hello" is text. In Python types exist, but nobody checks them before running: if you add a number to a string, the error only shows up when that line executes. In statically checked languages, a program called the type checker reads the code before it runs and rejects what does not make sense. Adding a number to a string does not even compile.

But a regular type checker knows nothing about shapes. To it, a 2×3 matrix and a 4×5 one have the same type: "matrix". The shape error gets through.

Bend2 has dependent types. The name is scary, but the idea is simple: a type can depend on a value. The label can carry numbers. So instead of "matrix", the type can be "matrix with 2 rows and 3 columns", written Mat<2, 3>. And then the multiplication rule becomes a type rule: Mat<n, k> times Mat<k, m> gives Mat<n, m>, and the k must be the same on both sides.

Try it yourself. Pick the shapes of A and B:

matrix A
matrix B
Mat<2, 3>
·
Mat<4, 5>
=
?—
SOME PROOFS FAIL
Error:
- expected : Mat<3n, 5n>
- observed : Mat<4n, 5n>
the rule is the same as in bend-ml: A's columns must equal B's rows. when they don't, the message mimics the real compiler's.

A framework in Bend can offer this and PyTorch cannot, so it became the project's pitch: a PyTorch where a shape error doesn't compile.

what a proof is

Bend2 goes beyond dependent types. Besides functions (def) and data (type), it has a third thing: the law (law). A law is a statement about the program, like "for every number a and every number b, a + b equals b + a". And the language demands a proof of every law.

If you have never seen a formal proof, the best picture is a row of dominoes. If you show that the first domino falls, and show that any domino that falls knocks over the next one, then all of them fall, no matter how many there are. That is called induction, and it is how you prove something about infinitely many numbers with a finite amount of work: you prove it for zero, and prove that if it holds for one number, it holds for the next.

In Bend, a proof is an ordinary function. The function's type is the statement, and its body is the argument. The zero case is the first domino; the function's recursive call is the "if it holds for the previous one". If the function passes the type checker, the statement is proved.

What checks this is the kernel, the part of the checker that decides whether a proof is right. Bend has two: the regular one, which is fast, and a second one, run with --verdict, which was written and proved correct in Lean, another proof language, one that professional mathematicians already use. Bend's compiler was written almost entirely by AI and has not been fully audited; the --verdict kernel is the part humans audited. That is why, in this project, every law had to pass both.

And this is where Bend's thesis comes in: "a language evolved by AI, secured by math". Whoever writes the code can make mistakes; the law, checked by the kernel, guarantees that certain properties always hold. ML is exactly the kind of code where nobody fully trusts what is written, and nobody proves anything.

reconnaissance

With the idea in hand, the first task was not code. Bend2 is new and changes fast, and its syntax has nothing to do with Bend1's, which was Python-style. So I read the documentation, the examples and the standard library, and wrote down everything I would need. A few findings changed the plan:

  • pin the version: everything was done on 2.0.35, and the project refuses to run on anything else. In a language that changes every week, that is survival.
  • numbers: there are only Nat (naturals: 0, 1, 2...), U32 (32-bit integers) and F32 (32-bit decimal numbers). No F64. For ML, F32 is fine, but it comes back when it is time to prove things.
  • the standard library has almost no lemmas: Base has numbers, lists and strings, but no proof that addition is commutative, for example. A lemma is a small law that serves as a step for proving others.
  • BendHub, the package registry: GitHub-only login, four-number versions (0.1.2.0), and publishing is public and permanent. A published version cannot be deleted.
  • what Bend compiles to: C, JavaScript and CUDA, which is what runs on NVIDIA graphics cards.

The first thing I wrote was a small proof of concept: a matrix with the shape in its type, a matmul, and two programs that should not compile. When they failed, for the right reason, the rest of the plan made sense.

the shape lives in the type

From here on the post gets more technical. The central type of the first package, bend-ml-tensor, is this:

type Mat<-r: Nat, -c: Nat> is Data:
  Mat{rows: List<&2, List<&2, F32>>}

A matrix with r rows and c columns stores a list of rows, each row a list of F32. The - before r and c marks the dimensions as erased: they exist only for the checker and vanish from the compiled program. So the guarantee costs nothing at run time. With that, multiplication gets the signature the math asks for:

def Mat.matmul(-n: Nat, +k: Nat, +m: Nat, a: Mat<n, k>, b: Mat<k, m>) -> Mat<n, m>:

The k shows up in both arguments. If one matrix's columns don't match the other's rows, no k satisfies both, and the program doesn't compile:

bend rejecting the product of a 2x3 and a 4x5 matrix: expected Mat<3n, 5n>, observed Mat<4n, 5n>
the compiler's real output. the program does not run: it never gets to exist.

The message says exactly what happened: the second argument should have been Mat<3, 5>, to match the first one's 3 columns, and a Mat<4, 5> arrived.

reshape is more interesting. It rearranges the same numbers into a different shape, say 12 numbers from 2×6 to 3×4. That only makes sense if the total count does not change. In bend-ml, reshape demands a proof of that as an argument:

def Mat.reshape(-r1: Nat, -c1: Nat, r2: Nat, +c2: Nat,
                -p: {Nat.mul(r1, c1) == Nat.mul(r2, c2) : Nat},
                a: Mat<r1, c1>) -> Mat<r2, c2>:

The parameter p has a strange type: {r1*c1 == r2*c2}. A value of that type is a proof that the two products are equal. To call the function, you have to hand over that proof. From 2×6 to 3×4, the proof is just {==}, which means "compute both sides and see that they match": 12 and 12. From 2×6 to 5×3, there is no proof that 12 equals 15, and the compiler refuses. And since the parameter is also erased (-p), the proof costs nothing at run time.

When the dimensions are not fixed numbers, the proof becomes a real law. Transposing r×c into c×r requires knowing that r*c == c*r for any r and c, and that is commutativity of multiplication, which had to be proved first.

The same holds inside training. In backprop, which I explain further down, the gradient of a 784×128 weight matrix W is Xᵀ·dY, and its type is Mat<784, 128>. Asking for that gradient as 128×784, the classic mistake of transposing on the wrong side, does not compile either. The whole MNIST training step is checked this way, layer by layer.

the missing proofs

Since Base has no lemmas, the first published package was not even ML: it was bend-ml-nat-lemmas, with 15 laws about Nat and lists. Commutative and associative addition, commutative, associative and distributive multiplication, identity elements, list concatenation and its length.

Commutativity of addition looks like this:

law add_comm:
  for  a: Nat
  for +b: Nat
  {Nat.add(a, b) == Nat.add(b, a) : Nat}

def add_comm(a, b):
  match a:
    case 0n:
      %add_zero(b) : {_ == Nat.add(b, 0n) : Nat}
      {==}
    case 1n+p:
      %add_succ(b, p) : {1n+Nat.add(p, b) == _ : Nat}
      %add_comm(p, b) : {1n+Nat.add(p, b) == 1n+_ : Nat}
      {==}

It is worth reading line by line, because it is the same structure as every proof in the project.

The law is the statement: for every a and every b, a + b == b + a. I explain the + before b later; for now, it says b may be used more than once.

The def with the same name is the proof. match a splits the two cases of a natural number: either it is zero (0n), or it is "one plus something" (1n+p), where p is the previous number. Those are the two dominoes.

In the zero case, the goal is 0 + b == b + 0. Bend already knows how to simplify the left side to b, but not the right one. %add_zero(b) uses another lemma, b + 0 == b, to rewrite the goal, and then both sides become b and {==} closes it.

In the 1+p case, the goal is (1+p) + b == b + (1+p). %add_succ rewrites the right side to 1 + (b + p). %add_comm(p, b) is the recursive call: the law itself, for the smaller number p. It is the "if it holds for the previous one". It swaps p + b for b + p, both sides become equal, and {==} closes it.

There is no tactic language, as in Lean or Coq: you prove by writing code. Coming from Python like me, reading this for the first time looks like magic; after the tenth proof, it looks like recursion with a very demanding type. The addition and multiplication proofs follow Bend's official numeric-proofs demo, adapted to Base's Nat.add and Nat.mul.

The lemma that mattered most was product_append, which says the product of a concatenated list is the product of each part. A list of dimensions is a tensor's shape, and its product is the number of elements: this lemma is what holds up reshape for arbitrary shapes.

This was also the first package on BendHub. Publishing is one command, but since there is no undo, it became a rule: I only publish after ALL PROOFS CHECK, after --verdict passes, with no @unsafe (the language's "trust me", which turns checking off) and no ?TODO (a hole left for later) in the file, and after importing the package from BendHub itself in a clean folder and using one of its laws.

text becomes numbers: the tokenizer

A language model does not read letters: it reads numbers. So all text going into GPT-2 first goes through a tokenizer, which cuts the text into pieces and swaps each piece for a number. Those pieces are the tokens.

The simplest way would be one token per letter. But then a sentence becomes a huge sequence. The other extreme, one token per word, would need a dictionary with every word of every language. GPT-2 uses a middle ground called BPE (byte pair encoding). It starts with the 256 possible bytes, which can represent any text, and learns merges: the pair of tokens that appears together most often becomes a new token. Then the next most frequent pair, and so on. Common words end up as a single token; rare words stay in pieces.

Type anything and move the slider to watch the training happen:

tokens (5)
aaabdaaabac
rules (oldest first)
  1. a + a → aa
  2. aa + a → aaa
  3. aaa + b → aaab
decoded
aaabdaaabac

✓ decode(encode(text)) == text

the same algorithm as the package: the most frequent pair becomes a new token, ties go to the pair seen first, and decoding expands each token back into bytes. the last line checks that the text came back whole.

Look at the last line: decoding what was encoded gives the original text back. It seems obvious, but it is exactly the kind of thing that breaks in production, on a strange character, a different rule order, an emoji. And it is the main law of the second package, bend-ml-bpe-tokenizer.

In the package, a token is a raw byte, B{n}, or a token created by the rule with id k, M{k}. The table is a list of rules from oldest to newest, just like GPT-2's merges.txt. And the law is this:

the roundtrip law as bend states it, followed by ALL PROOFS CHECK
the statement the kernel checked, exactly as it is in the published package.

decode(encode(s)) == s, for any byte sequence, as long as the table is well formed, meaning no rule id repeats. The idea of the proof is that when rule k merges the pair (a, b) into M{k}, M{k} expands back to the expansion of a followed by that of b, so decoding changes nothing. That only holds if k did not already exist in the table, which is exactly what "well formed" guarantees. The helper lemmas prove that stacking the new rule on top of the table does not change the expansion of the tokens that already existed.

Alongside it came vocab_bound (every token encode emits is a byte or was created by a rule of the table, never an unknown id) and dec_append (decoding two lists together is decoding each and concatenating the bytes).

Then came the part I like most: proving that train always produces a well-formed table (train_wf). The proof is by induction over the training loop, with the invariant "every id already used is smaller than the next one". An invariant is something that holds before and after every step of a loop; finding the right invariant is half of any proof about loops. Putting the two laws together gives roundtrip_trained: train a table on any corpus, with as many rules as you want, and the roundtrip holds, with no hypothesis at all.

What is not proved is which pair train chooses to merge. And it does not need to be: the roundtrip's correctness does not depend on it. That part is tested against a reference implementation in Python, with English text, accented text (coração, ação), emoji and Chinese characters. train, encode and decode give exactly the same results on both sides.

how a network learns: the gradient

Back to the knobs. To know which way to turn each weight, training computes the gradient: for each weight, how much the error changes if that weight goes up a tiny bit. If raising the weight raises the error, you turn it down; if it lowers the error, you turn it up. It is the slope of the terrain, and training is walking down the hill of error.

Computing the gradient of a function with millions of weights sounds impossible, but there is a trick. Every function in a model is made of small operations (sums, products), chained together. And there is a rule from calculus, the chain rule, that says how to combine the slopes of the parts into the slope of the whole. A framework does this on its own, and that is called automatic differentiation, or autograd.

There are two ways to apply the chain rule. Forward mode carries the slope along with the value, from input to output. Reverse mode computes all the values first and then walks back from the output to the input, handing the gradient backwards. Reverse mode is what every ML framework uses, because in a single pass back it computes the gradient of every weight at once. That is what is called backpropagation.

on the way back, each product hands on the gradient times the value of the other factor, and each sum hands it on unchanged. whatever reaches x along different paths adds up.

The third package, bend-ml-autograd, implements both modes and proves they agree. The expression is a tree of constants, the variable, sums and products:

type NE is Data:
  NCst{n: Nat}
  NX{}
  NAdd{a: NE, b: NE}
  NMul{a: NE, b: NE}

Reverse mode receives the gradient g coming from above and hands it to the children. In a product, each side receives g times the value of the other side, exactly as in the diagram:

def nbwd(e: NE, +x: Nat, +g: Nat) -> Nat:
  match e:
    case NCst{n}:
      0n
    case NX{}:
      g
    case NAdd{a, b}:
      Nat.add(nbwd(a, x, g), nbwd(b, x, g))
    case NMul{+a, +b}:
      Nat.add(nbwd(a, x, Nat.mul(g, nval(b, x))), nbwd(b, x, Nat.mul(g, nval(a, x))))

And the law:

law reverse_eq_forward:
  for +e: NE
  for +x: Nat
  {nbwd(e, x, 1n) == nfwd(e, x) : Nat}

The proof goes through a stronger lemma, nbwd(e, x, g) == g * nfwd(e, x) for any g, because the induction only closes if the gradient coming from above is generic. With g = 1, the law follows. It is a common pattern: sometimes the statement you want is not provable directly, but a more general one is.

This is where the rule I wrote at the start of the project, and kept to the end, shows up: floats are not reals. F32 rounds: 0.1 + 0.2 is not exactly 0.3, and addition is not always associative. Proving numeric properties about F32 means trying to prove false things. So the law is proved over Nat, where arithmetic is exact, and the real numbers are tested against PyTorch with gradient checking: you compute the gradient with autograd and compare it with a numeric approximation, nudging the input slightly and measuring the change in the output. Proofs for structure, tests for numbers. That split runs through the whole project.

MNIST: teaching it to read digits

With the packages in place, the two demos came next. The first is machine learning's "hello world": MNIST, a set of 70 thousand photos of handwritten digits, 0 to 9, each 28×28 grayscale pixels. The task is to look at the photo and say which digit it is.

The model is an MLP (multilayer perceptron), the simplest neural network there is. The 784 pixels (28×28) go in, pass through a layer of 128 neurons, and 10 numbers come out, one per digit; the largest is the guess. Each layer is a matrix product plus an adjustment vector (the bias), followed by a function that bends the space, ReLU (which zeroes the negatives).

Training looks at 100 photos at a time, a batch, and the shapes of one step look like this:

X      : Mat<100, 784>     100 photos, 784 pixels each
W1     : Mat<784, 128>     first layer's weights
H = X·W1 + b1 : Mat<100, 128>
W2     : Mat<128, 10>      second layer's weights
Z = H·W2 + b2 : Mat<100, 10>   10 scores per photo

dW2 = Hᵀ·dZ : Mat<128, 10>     the gradient has the weight's shape
dW1 = Xᵀ·dH : Mat<784, 128>

At the end of each step, every weight moves a tiny bit against its gradient (learning rate 0.1). An epoch is going through the 60 thousand training photos once: 600 batches.

To check the result, I made the initial weights in the PyTorch reference script and exported them to Bend, and both sides read the same batches in the same order. That way you can compare number to number, not just "similar accuracy". And it matched: the first epoch's loss is 0.5204771 on both, with 9129 hits out of 10000 test photos. Then 0.27043572 and 9298, then 0.2156194 and 9418. Identical to the seventh digit.

GPT-2: predicting the next word

The second demo is GPT-2 small, the language model OpenAI published in 2019, with 124 million weights. It does one thing: given a text, it predicts the next token. To generate a sentence, you ask for the next token, append it to the text and ask again.

Inside, it is a stack of 12 identical layers, each with two parts. Attention lets each position in the text look at the earlier positions and decide which ones matter (in "the capital of France is", the word "France" matters a lot for what comes next). The MLP transforms the result, as in MNIST. Between the parts come normalizations (LayerNorm) and a smooth function similar to ReLU (GELU).

Getting this to run in Bend took more pieces than I expected:

  • the weights come from the official file, model.safetensors, which a Python script splits into one binary file per tensor. Bend reads each file in blocks with File.read_at and turns every 4 bytes into an F32.
  • the tokenizer uses my BPE package with GPT-2's 50 thousand rules. But GPT-2 does not apply the rules in order: it always looks for the pair with the lowest position in the table. I applied them in order and checked against tiktoken, the official tokenizer, that for trained tables both give the same result.
  • the pre-tokenizer: before BPE, GPT-2 splits the text with a regular expression that separates contractions, letters, numbers, punctuation and whitespace. Bend has no regular expressions, so I wrote the pre-tokenizer by hand, byte by byte, imitating each alternative of the regex.
  • the byte permutation: GPT-2 renumbers the 256 bytes with its own table, which becomes one more file.
  • the cache: to avoid recomputing the whole text for each new token, every layer keeps the attention keys and values of the earlier positions, the KV cache.
  • generation: greedy, always the highest-scoring token, which is the easiest to compare.

The result: the same tokens as PyTorch, with each token's scores (the logits) equal to the fourth decimal.

$ ./gpt2_fast "The capital of France is" 8
text: The capital of France is the capital of the French Republic, and

(GPT-2 small is not great at geography. But it is wrong in exactly the same way on both sides, which is what matters here.)

v1 worked. and it was 1800 times slower

At the end of the first day, everything was correct and everything was slow.

One MNIST epoch took 544 seconds. PyTorch takes 0.3. About 1800 times slower. GPT-2 took 3 seconds per token, and around 10 seconds just to load the weights.

The temptation was to conclude that Bend is slow and stop there. But I had not measured anything, only timed the total. So v2 started with a different plan: optimize nothing without first having a number that said where the problem was. The second day became a series of experiments, each with a hypothesis and a measurement before any change to the code. All of them are in the project notes, with the command to reproduce them.

why lists are slow

The first suspect was the data structure. v1 stored a matrix as a list of rows, and each row was a linked list: each number sits anywhere in memory, together with the address of the next one. To reach the tenth number, you go through the nine before it.

The alternative is an array: the numbers side by side, and you reach any of them directly by position, the index.

in a list, reading a number means following arrows. in an array, it means doing arithmetic on the index.

For the processor, the difference is huge. It fetches memory in blocks and tries to guess what you will read next. With an array, the next read sits right next to the previous one and already came in the same block. With a list, each number can be anywhere, and each read waits for the previous one to finish to learn the next address.

Experiment 1. I rewrote the same matrix product (100 million multiply-adds, one thread) with a flat Array<F32>, accessed by index:

timemultiply-adds/s
lists2.274 s44 M
Array<F32> by index0.046 s2,200 M

Forty-nine times, just by changing the data structure. v1's conclusion ("Bend is 1800 times slower") was an effect of the data structure, not of the language.

Bend's Array has a catch: underneath, it is not a contiguous block of memory, it is a tree, and the index picks the path through the branches. That showed up in two more numbers. Walking the whole tree instead of indexing was 70 times worse. And an array with 2^20 slots instead of 2^17 made the same product 60% slower, because the tree's depth enters every read. The rule that came out of it: allocate the smallest possible array.

affine types, in practice

The array brought in a Bend concept I had only read about: affine values. An affine value can be used at most once. That is what lets Bend manage memory without a garbage collector: if the compiler knows nobody else points to a value, it can free it right away or reuse the space.

For an array, that has a curious consequence: reading an element "uses up" the array. So Array.get returns the element and the array back, so you can keep using it. That leaked into the package's API: every operation returns what it read, together with the result, in a record with its own type. The matrix product returns this:

type MMul<-n: Nat, -k: Nat, -m: Nat> is Type:
  MMul{a: Mat<n, k>, b: Mat<k, m>, c: Mat<n, m>}

Both factors back, and the product. It felt strange for the first hour; after that it felt natural, because it makes explicit who owns each matrix at each point. And the parameter markers complete the story: + allows reusing a Data value (numbers, plain lists), - erases the parameter from the compiled program, and a parameter with no marker can be used only once.

parallelism: when splitting helps and when it hurts

My laptop's processor has 22 threads, meaning it can do 22 things at once. Bend is built for that: it automatically runs in parallel the parts of a program that do not depend on each other. A matrix product is perfect for it, because each row of the result is independent of the others.

Experiment 2. I split the product into row blocks and sent each block to a task. Since values are affine, each task needs its own copy of the matrices (Array.clone). My hypothesis was that the copies would share memory underneath and the tasks would fight over access. I measured it: 16 tasks reading copies scale exactly like 16 tasks with their own arrays. Hypothesis refuted, and I was glad I measured instead of optimizing for a problem that did not exist. The gain on MNIST was about 2 times.

Experiment 3. In GPT-2, generating one token at a time, the math is matrix times vector: a single input row. I parallelized it the same way and it got 10 times slower, with 1, 8 or 16 threads.

what decides whether the copy pays off is how many times each weight is used after being copied.

The cause is simple arithmetic. Cloning costs about 0.1 nanoseconds per number, per task. In matrix times vector, each weight is read only once, and the math with it costs about 0.4 nanoseconds. Copying the matrix to several tasks costs more than doing the whole computation alone. In MNIST, matrix times matrix with 100 photos, each weight is read 100 times, and the copy disappears in the noise. I only found that out by measuring.

what moved the needle: lists to array 49x faster, parallel blocks 1.8x, GPU on a flat loop 3.8x faster, GPU on our kernels 3 to 18x slower, copying the matrix per task 10x slower
each line is an experiment in the project notes, with the number measured before the change.

With the numbers in hand, v2 became a new package, bend-ml-tensor-array: the same Mat<r,c> with the shape in the type, but over a flat Array<F32>, row by row (the number at row i, column j sits at position i*c + j).

A detail I like in it: training needs three different products, X·W (forward), H·Wᵀ and Xᵀ·dZ (backward). Transposing a matrix in memory is expensive. Instead, the package has a single generic product, gemm, which takes the strides the indices walk with in each matrix:

def gemm(+n: Nat, +m: Nat, +k: Nat, +arow: U32, +sa: U32, +sb: U32, +bcol: U32,
         a: Array<F32>, b: Array<F32>, c: Array<F32>) -> G:

The element at row i, position p of A sits at i*arow + p*sa. For a plain A, arow = k and sa = 1. To read A transposed without transposing anything, just swap them: arow = 1 and sa = n. The three public products (matmul, matmul_nt, matmul_tn) are the same gemm with different strides, all three with the right shape in the type. All have optional parallel blocks, and GPT-2 runs its own sequentially, because experiment 3 said so.

In the end, the MNIST epoch went from 544 seconds to about 7 seconds, and GPT-2 from 3 seconds to 0.1 seconds per token.

the GPU looked useless

A graphics card is a different beast. A CPU has a few very smart cores; a GPU has thousands of simple cores that do the same math at the same time on different data. It is perfect for ML, and Bend compiles to CUDA, the platform of NVIDIA cards, when you mark a call with !.

I have an RTX 4050. The first problem was installing: Bend asks for CUDA 12, and my Arch package is 13, which would need sudo and break other things. Digging in, I found that Bend honors the CUDA_HOME variable and only needs two pieces of NVIDIA's kit (cudart and nvrtc). I downloaded the standalone archives, unpacked them into ~/.local/cuda12, and it worked without sudo.

Then came the first result: the GPU was slower than the CPU at everything. From 3 to 18 times, with no point where the curve turned. My first note was "the GPU is useless for this". Rereading it, I found the conclusion too quick: if a GPU loses to the CPU at everything, the likeliest explanation is that I am using it wrong.

I was. Bend's guide says a ! call has to fan out into about 16 thousand leaves, one per GPU execution lane, and end in simple loops, numbers only. My tests used 32 to 1024 tasks: the card sat almost entirely idle. With 16,384 leaves and a simple numeric loop, the GPU won, and the advantage grew with the work per leaf:

16,384 leaves, simple F32 loopGPUCPU (22 threads)
20 thousand iterations each0.217 s0.054 s
200 thousand iterations each0.254 s0.300 s
2 million iterations each0.479 s1.808 s

3.8 times faster in the last case. The fixed cost of about 0.2 seconds is compiling the GPU code on the spot.

But the matrix products kept losing, and the reason is beautiful once you see it. Remember the linked list and the array's tree? Each read depends on an address that comes from the previous read. A GPU is great at doing the same math on lots of data that already sits side by side; it is terrible at following addresses. And the memory Bend uses on the GPU is "managed": pages cross the bus between the CPU and the card on first access. Our products are chains of dependent reads, the opposite of what a GPU wants, and that shape of the data is what makes them lose.

The benchmarks stayed on the CPU, with the reason written down. And the first conclusion is still in the notes, corrected instead of erased.

making it stable

Up to that point, everything worked on my machine and on no other. When I asked myself "is this a real v1?", the honest answer was no: a stranger cloning the repository could not run the tests, nothing ran on its own, and the person who wrote the tests was me, the same person who wrote the code. The last part of the work was closing each of those gaps.

To start, make setup && make check on a fresh clone installs Bend pinned to 2.0.35 (checking the binary's SHA256, the file's fingerprint), Lean 4.34.0 for the audited kernel, the Python environment for the references and the MNIST and GPT-2 data, all without sudo. I tested it on a new clone, with an empty HOME, as if I were someone else.

Then CI, GitHub's machine that runs the tests on every change, runs 36 checks on every push, GPT-2 included, on a free runner, in about 11 minutes. Running the full GPT-2 on a free machine was the part I thought would not fit, and it did.

Tests written by the person who wrote the code tend to test what that person already knows works. So randomly generated tests came in: random expression trees for autograd (compared with PyTorch), 80 random corpora for BPE (compared with the Python reference), random shapes for the tensors, and 11 prompts for GPT-2. And a test that imports every package from BendHub, at the version the README states, and applies one of its laws. That one catches a broken publication or a stale README.

For anyone who wants to audit it, a script prints the statement of every law, without the proofs, for anyone who wants to review the specification without reading the implementation. A document lists what is proved, what is tested and what is trusted.

One proof was still missing. The Mat over Array assumed the array had room for r*c numbers. Bend allocates arrays with 2^d slots, and the calculation of d used U32.log2, about which nothing can be proved. I rewrote it with recursion on Nat:

def cap_pair(n: Nat) -> Nat & Nat:
  match n:
    case 0n:
      (0n, 1n)
    case 1n++m:
      cap_step(cap_pair(m), 1n+m)

cap_pair(n) returns the pair (d, 2^d), and each step doubles the power when n no longer fits. With that, the law could be proved:

law cap_ok:
  for +n: Nat
  {True{} == leb(n, pow2(cap_depth(n))) : Bool}

n <= 2^cap_depth(n): the array always has room. What remains trusted is only that the runtime allocates 2^d slots when asked.

Bad inputs now become clear errors. The loss and the hit count gained variants that return None when the number of labels does not match the number of rows, instead of computing garbage. And the tokenizer now validates the rules file before using it: an even number of ids, and no rule referring to a token that does not exist yet:

def ids_ok(ns: List<&2, Nat>, +k: Nat) -> Bool:
  match ns:
    case Nil{}:
      True{}
    case Con{+a, t}:
      match t:
        case Nil{}:
          False{}
        case Con{+b, u}:
          Bool.and(Bool.and(Nat.is_lt(a, Nat.add(256n, k)), Nat.is_lt(b, Nat.add(256n, k))), ids_ok(u, 1n+k))

Rule k may only refer to bytes (below 256) or earlier rules (below 256 + k).

Finally, Unicode. GPT-2's pre-tokenizer treated every byte above 127 as a letter. That worked for accents, but not for em dashes, emoji or Chinese punctuation, which GPT-2's regex treats as "other". The fix: before pre-tokenizing, each byte gets a tag with the class of the character it belongs to, encoded in the number itself:

def tagged(+b: U32, +cls: Nat, +cont: Bool) -> U32:
  U32.add(U32.add(b, U32.mul(U32.from_nat(cls), 256)), U32.mul(U32.from_nat(sel(cont, 1n, 0n)), 2048))

The byte stays in the low 8 bits, the class (letter, number, other, whitespace) in the next bits, and one bit marks the continuation bytes of a multi-byte character. The class comes from a table of 939 Unicode ranges, generated from Python's regex module, with the same categories GPT-2's regex uses. The rest of the pre-tokenizer works on the classes and strips the tags at the end. Result: it matches tiktoken on 79 texts, including 30 strings of random Unicode.

Two stories from this stretch taught me a lot.

The first: the README said that loading GPT-2 used about 10 GB of memory. Measured properly, 10 GB was the address space the process reserved, not memory actually in use; resident memory was 2.3 GB. After reading the weights in 1 MB blocks straight into the array, without building a list in between, it dropped to 1.5 GB, and loading got 5 seconds faster.

The second: after swapping the capacity calculation for the proved version, the MNIST epoch went from 7 to 16 seconds. It looked like the proof had been expensive. I ran the old and new versions side by side: both took 16 seconds. The culprit was my desktop's screensaver, using 134% of a CPU. With the machine idle, it went back to 7. If I had not compared both versions under the same conditions, I would have thrown away a good proof because of a screensaver.

the project's look

The last thing was the README and the video. My first version looked like a generic terminal theme: dark background, saturated blue, cyan and red, macOS-style terminal windows. It was ugly, and I redid it from scratch.

The final version uses Bend's own identity, taken from bend-lang.com's CSS: warm paper (#f2eee7), dark ink, violet as the accent, green for proofs, red for errors. A single monospace font, iA Writer Mono, and the violet cursor blinking after the name, which is their site's signature. The charts follow the site's bar style: vertical, linear scale, Bend in violet and the rest in grey, with the bars that run off the top hatched.

The video at the top of this post is generated too, not recorded. It is an HTML page with a function that draws the frame for a time t; a headless browser takes 3600 screenshots of it and ffmpeg stitches them together, blending frames in pairs for motion blur. Every terminal output and every number in the video is read from the files of the real runs. The 26 cells of the opening matrices travel through the whole film: they become the logo, rearrange in the reshape, become the package boxes, the GPT-2 tokens, the benchmark bars, and return to the logo at the end. make media rebuilds all of it.

where bend got me

A few things I wish I had known on day one:

  • affine variables: use something twice, error. The fix is marking the parameter with + (reuse allowed for Data values), or reshaping the code so you do not need to.
  • structural recursion only: the termination checker requires recursion to visibly shrink something, to guarantee the program ends. When the algorithm does not have that shape, the way out is a "fuel" parameter that shrinks at each step. The pre-tokenizer jumps 1 to 4 bytes at a time (the size of a UTF-8 character), so it takes the number of remaining bytes as fuel, and each step spends one:
def tag(fuel: Nat, +tbl: List<&2, U32>, +xs: List<&2, U32>, +n: Nat, +cls: Nat) -> List<&2, U32>:
  match fuel xs:
    case 1n+p Con{b, t}:
      emit(n, cls, False{}, xs, tag(p, tbl, List.drop(&2, U32, xs, n), ...))
    case _ _:
      Nil{}
  • match only on parameters: you cannot match on an expression; it becomes a helper function that receives the already-computed value as an argument.
  • declaration order: a function can only call another declared before it, and mutual recursion does not pass. More than once the fix was just swapping two functions.
  • nested tuple patterns in match do not infer their type. The fix was replacing tuples with records that have their own type, like MMul.
  • imports are not re-exported: if a package imports a law from another file, whoever imports the package does not see it. That is why each published package is a single file, with law and proof side by side.
  • publishing is forever: a published version cannot be deleted, so every package only went up after all the proofs and the whole test suite passed. The old versions are still there, and the README always points at the current one.

These rules exist for a reason: they are what let Bend manage memory without a garbage collector, guarantee every program terminates and check proofs fast. Still, each one cost me a compile error to understand.

the numbers, no varnish

benchmarks: MNIST epoch, v1 544 s, v2 7 s, PyTorch 0.2 s; GPT-2 per token, v1 3 s, v2 0.1 s, PyTorch 55 ms on 1 thread and 21 ms on 16
linear scale. v1's hatched bars run off the top of the chart; lower is better.
v1 (lists)v2 (array)PyTorch
MNIST, one epoch544 s~7 s0.2 s
GPT-2 small, per token~3 s~0.1 s21 ms (16 threads), 55 ms (1 thread)
GPT-2, loading the weights~10 s~7 s, 1.5 GB~1 s

All measured on the same laptop (Core Ultra 7 155H, 32 GB), with the machine idle; MNIST is the median of three runs. PyTorch is still about 35 times faster on MNIST and 2 to 5 times on GPT-2.

What is left of the gap is in the generated code. Bend 2.0.35 emits scalar code: one operation at a time. Modern processors have SIMD instructions, which do 8 or 16 operations in a single instruction, and PyTorch uses libraries like BLAS, tuned for decades to use every trick the processor has. On the same matrix product, PyTorch does 62 billion multiply-adds per second on one thread; Bend, with an array, 2.3 billion. Closing that gap is compiler work.

what's left

Five packages on BendHub, all MIT, all passing --verdict, with no @unsafe and no ?TODO:

the five packages, who imports whom, and what is proved, tested or trusted
green is proved by the kernel, grey is tested against a reference, orange is trusted.
  • bend-ml-nat-lemmas: 15 laws about Nat and lists
  • bend-ml-bpe-tokenizer: the tokenizer, with roundtrip, vocab_bound, dec_append, train_wf and roundtrip_trained
  • bend-ml-tensor: Mat<r,c> over lists, with a proved reshape
  • bend-ml-tensor-array: the same over a flat array, about 50 times faster, with cap_ok
  • bend-ml-autograd: automatic differentiation, with reverse_eq_forward

24 laws in total. And what I take from the project, more than the packages:

  • types can carry the shapes of a whole training step, including the shape of every gradient, and they cost nothing at run time.
  • proofs cover structure, tests cover numbers, and knowing where the line between them runs is half the work.
  • the trust base fits in a list: the kernel, the F32 primitives and Array.new. It is written down in the audit guide.
  • the data structure matters more than the language: the "1800 times slower" was linked lists, not Bend.
  • measure before you claim. Three of my conclusions were wrong (PyTorch's time per token, the GPU "slow at everything", the 10 GB of memory), and all three were corrected by a measurement, not by an argument. The corrections stayed in the notes, next to the wrong conclusions.

The obvious next step is splitting the weight array between tasks without copying: since the Array is a tree underneath, its branches can be separated and handed one to each task. That would unlock parallelism in matrix times vector, which today loses to the sequential version. From there to a real llama.bend, the path exists.

The code, the video and all the notes are in the repository. Cloning it, running make setup && make check, and checking every number in this post is the best reply I could get.

Writing this post also made me want to write more about AI here, the same way I did with cryptography: starting from scratch and going deep. Something along those lines is coming soon.