Writing tensor programs in Lean

Robert Joseph George · TorchLean Verified Examples

I've wanted to use TorchLean for calculations that don't need a neural-network model around them. If I'm experimenting with an activation or repeating an update in a simulation, I'd like to write that function in Lean and choose CPU or GPU when I run it. I also want to be able to prove something about the function I wrote.

The tensor library already gives us operations on shaped tensors, automatic differentiation (autograd), and FloatLib's numerical formats. Autograd records the operations and intermediate values needed to propagate derivatives backward; I'll call that recording a tape below. I've been working on connecting custom functions to those pieces. We can now compile supported Lean functions for GPU execution and use them through the tensor API.

A GPU is useful when a calculation has lots of independent work. Applying the same update to a million starting values gives us a million computations we can run alongside one another. Within each computation, however, the next update may depend on the previous one. We need a way to express both the independent work and that dependency without maintaining another version of the function in CUDA.

Lean gives us an interesting place to do this. It is a functional programming language as well as a theorem prover, so we can write the calculation, evaluate it, and prove a statement about its definition. I want those activities to feel like parts of the same project. The GPU compiler adds another question we can state precisely: does the lowered program still compute the function we wrote?

Let's try it with a repeated update, a custom loss, and a row sum. Then we'll write matrix products and attention with einsum, compose a model, and look at what takes time when we run it. I also want to show the precision and higher-derivative work, because those features change which calculations we can write, not only how we spell them.

Functional programs on a GPU

Functional programming has a long history of treating a collection of calculations as a program the compiler can transform. In Accelerate, we write array computations inside Haskell using operations such as maps and reductions. Its backends compile those computations for a target device. The distinction between the host language and the embedded array program matters: an ordinary Haskell function doesn't become a GPU program merely because we use Accelerate.

Futhark's introduction gives a motivation I share:

Because it's nicer than writing CUDA or OpenCL by hand!

Here we can also write a theorem about the function before and after lowering it. That connection is what I want to keep as we make the compiler more useful.

Futhark takes a different approach, with a functional array language designed for GPU compilation. Its work on nested parallelism asks how a compiler can expose useful parallel work while retaining efficient loops and updates. A short functional expression can describe many calculations, but the compiler still has to decide how to organize them on the machine.

Dex is particularly relevant to tensor programming. It makes indices part of how we write array computations and studies differentiation that preserves their parallel structure. Indices tell us which entries a calculation reads and which outputs it produces. That is useful information for both a compiler and a proof.

Bend exposes parallel work through functions too. Its parallel-call model lets a programmer split independent calls and run them through a CPU/GPU runtime. It has a broader goal than our current tensor compiler, including recursive data and divide-and-conquer programs. Our custom path starts with tensor entries and ordered numerical loops.

Bend also uses affine ownership: an ordinary value can be used at most once, unless it is explicitly reusable. That restriction helps control sharing and updates. It is a different design choice from adding a tensor compiler to Lean; we aren't introducing an affine language or replacing Lean's runtime here.

Our starting point is the tensor system and certified einsum compiler we already have, which connects a generated computation to reference semantics. We want to extend that approach to functions users write themselves. Mathlib supplies the mathematics we can reuse for proofs, and FloatLib lets us describe the arithmetic rather than assume that a floating-point calculation is an exact real-number one.

That numerical support is also a concrete difference from Bend's current implementation: its documented limitations describe native F32 arithmetic as axiomatic. FloatLib supplies executable arithmetic with correctness and rounding theorems. Connecting those theorems to a foreign GPU implementation is further work, but we already have the numerical definitions to make the question precise. I wouldn't infer a speed advantage from that difference; execution speed needs a matched experiment.

Square, blend, and choose a device

Let's square the values 1, 2, and 3:

import NN.Kernel

open TorchLean

def square {α : Type} [Mul α] (x : α) : α := x * x
def input : Tensor Float32 [3] := [1, 2, 3]

-- The device belongs to the call, so the function stays unchanged.
def onGpu : IO (Tensor Float32 [3]) := square.run input (device := gpu)
#eval square.run input (device := cpu)

[1.000000, 4.000000, 9.000000]

The answer is [1, 4, 9]. Each output needs only its corresponding input, so the three calculations are independent. A tensor with a million entries would expose the same kind of work. The shape stays where it belongs, on the tensor: squaring a matrix replaces each entry in the same row and column. We can use this function on different tensor shapes without naming a new operation for every rank.

The definition also leaves the number type open. In the code, α is the type of an entry, and [Mul α] asks for multiplication on that type. The input chooses Float32 for this example. Later we'll use a wider FloatLib format, keeping the same multiplication expression.

We choose cpu or gpu when we run the function. CPU evaluates the Lean calculation. The custom GPU path recognizes supported arithmetic, produces a checked representation, and generates CUDA. NVIDIA's NVRTC compiles that source at runtime. We can keep the function itself free of a device choice.

A square is a small place to start, but the reason for having this interface is to let the body grow. Local bindings, arithmetic comparisons, and supported conditionals can be part of the function. We should be able to try a calculation before deciding it deserves a permanent name in the library.

An update often needs two inputs: a previous estimate and a new observation, or two fields we want to interpolate. Let's blend corresponding entries with a weight of one quarter:

import NN.Kernel

open TorchLean

def blend (weight left right : Float32) : Float32 :=
  (1 - weight) * left + weight * right

def left : Tensor Float32 [3] := [0, 4, 8]
def right : Tensor Float32 [3] := [8, 4, 0]

#eval (blend 0.25).zip left right
def gpuBlend : IO (Tensor Float32 [3]) :=
  (blend 0.25).zip left right (device := gpu)

[2.000000, 4.000000, 6.000000]

For the first pair, three quarters of zero plus one quarter of eight gives two. The middle pair stays at four, and the last pair gives six. Changing the weight changes that calculation; it doesn't require a new function named after a particular coefficient.

Both inputs have shape [3]. The type of .zip requires the shapes to agree, so it doesn't truncate one input to the other's length or choose broadcasting we didn't request. The result keeps that shared shape. A tensor of another rank can use the same scalar blend.

The weight is shared by all entries. For GPU compilation, a captured scalar such as this weight becomes a constant in the generated source. Changing it can require recompilation; the runtime reuses identical source through its cache. Frequently changing per-entry values can instead arrive through tensor inputs. This is an execution choice worth understanding before using a small example in a longer loop.

Repeated updates and their closed forms

Now repeat the update \(x\mapsto x^2+1\). Starting at one, the first update gives two, the second gives five, and the third gives twenty-six. Starting at two gives a different history. Neither history needs to read the other.

Move the slider back to zero: the output is the starting value. At two updates, the outputs are five and twenty-six. This is also how we define the recurrence in Lean:

import NN.Kernel

open TorchLean

-- Keep the element type general; the step needs multiplication and addition.
def evolve {α : Type} [One α] [Mul α] [Add α] (x : α) : Nat → α
  | 0 => x
  | n + 1 =>
      let previous := evolve x n
      previous * previous + 1

def starts : Tensor Float32 [2] := [1, 2]
def evolution (steps : UInt64) (device : NN.Backend.Device := cpu) :=
  (fun x => evolve x steps.toNat).run starts (device := device)

#eval evolution 2

-- Recording visits both updates, so backward includes both local derivatives.
#eval do
  let result ← (fun (x : Float32) => evolve x 2).run starts (grad := true)
  IO.println (← result.backward)

[5.000000, 26.000000]

[8.000000, 40.000000]

The zero case returns the input. The successor case computes the previous result, then squares it and adds one. Writing recursion makes the dependency explicit: the new value comes from the old one. We haven't created a separate implementation for GPU.

For a forward-only GPU call, the frontend uses the recurrence's zero and successor equations to derive an ordered loop. It constructs a Lean proof that the loop agrees with those equations. The body isn't copied into the generated program once for every update. Increasing the count gives the loop more work to do.

This also makes the available parallelism clear. A thousand starting values can evolve independently, but ten updates of one value form a dependency chain. A compiler cannot run the second update before it knows the first result. Some recurrences admit more elaborate parallel algorithms; using one would require showing that its reorganized computation has the meaning we intend.

The recognized recursion currently keeps its parameters fixed and recurses on the predecessor of its final natural-number argument. The call passes a UInt64 count converted to a natural number, giving the GPU loop a bounded count without truncating an arbitrary Nat. Tree recursion and dynamically allocated recursive structures remain outside this path. Ordinary CPU Lean can express a much larger range of programs.

Those boundaries are why I find the connection to functional GPU languages useful. Parallelizing a collection of ordered updates is one kind of problem. Scheduling a recursive tree whose branches have different amounts of work is another. The tensor interface gives us a good starting point for the former; it doesn't automatically solve the latter.

The square recurrence grows quickly. Let's try an update with a useful numerical purpose: move part of the remaining distance to a target. If the current value is zero, the target is two, and we move halfway each time, the first three values are one, one and a half, and one and three quarters.

import NN.Kernel
import Mathlib.Tactic.Ring

open TorchLean

-- Move a fraction of the remaining distance to the target at each step.
def relax {α : Type} [Add α] [Sub α] [Mul α] (rate target x : α) : Nat → α
  | 0 => x
  | n + 1 =>
      let previous := relax rate target x n
      previous + rate * (target - previous)

def initial : Tensor Float32 [3] := [0, 2, 4]
def approach (steps : UInt64) (device : NN.Backend.Device := cpu) :
    IO (Tensor Float32 [3]) :=
  (fun x => relax (0.5 : Float32) 2 x steps.toNat).run initial (device := device)

#eval approach 3

-- The real error is multiplied by the same factor at every step.
example (rate target x : ℝ) (steps : Nat) :
    relax rate target x steps = target + (1 - rate) ^ steps * (x - target) := by
  induction steps with
  | zero => simp [relax]
  | succ steps ih =>
      rw [relax, ih, pow_succ]
      ring

[1.750000, 2.000000, 2.250000]

The tensor starts at [0, 2, 4]. After three updates, its entries are [1.75, 2, 2.25]. The entry already at the target stays there; the other two approach it from opposite sides. approach 3 (device := gpu) selects the GPU route for the same recurrence.

We can also understand every step without running all of them. Write the rate as \(r\), the target as \(t\), and the initial value as \(x_0\). Subtracting the target from one update gives

The example proves the second identity for every real rate, target, initial value, and natural-number step count. The proof uses the same generic relax definition as the executable example. Induction handles the repeated updates, and mathlib's ring tactic checks the polynomial identity in the successor case.

That formula tells us how the parameter changes the behavior. A rate of one half halves the remaining distance. A rate of one and a half multiplies the error by −1/2: the value crosses the target each time, while the distance still shrinks. At rate two, the factor is −1 and the error keeps oscillating. Rates greater than two amplify it.

The same identity covers all natural-number step counts and real initial values, including rates that make the update unstable. For the rounded recurrence, we'd need to add a finite-precision error bound using the numerical operations we selected.

Gradients and a custom loss

The recurrence block also records its two updates and runs backward. In real arithmetic, two updates give

At one, the derivative is eight; at two, it is forty. Backward follows both updates in reverse, multiplying their local derivatives. If we differentiated only the final square, we'd miss how the first update depended on the input.

This uses the existing tape. We keep the function definition and request recording with grad := true. Recording retains intermediate values for backward, so it has a different execution cost from the forward-only compiled loop. On GPU, a recorded recurrence currently launches tape operations per step.

We can write a whole-tensor function through the same interface. Let's add the squares:

import NN.Kernel

open TorchLean

def energy (x : Tensor Float32 [3]) : Float32 := Tensor.sum (x * x)
def trainingInput : Tensor Float32 [3] := [1, 2, 3]

-- The default output seed is one: differentiate the scalar energy.
#eval do
  let result ← energy.run trainingInput (grad := true)
  IO.println result.value
  IO.println (← result.backward)

def gpuEnergy := energy.run trainingInput (device := gpu) (grad := true)

14.000000
[2.000000, 4.000000, 6.000000]

For [1, 2, 3], the energy is fourteen. The gradient is [2, 4, 6]: increasing a coordinate \(x_i\) changes its contribution \(x_i^2\) at rate \(2x_i\). The function takes a tensor and returns one number, so its forward result has scalar shape [], while the gradient has the input's shape [3].

Each multiplication uses the input twice. Backward accumulates a contribution from both uses. The sum passes its output sensitivity to every coordinate. Larger models reuse these same local derivative rules when they call the operations.

The default backward seed differentiates the sum of the output entries. For a scalar energy, that means a seed of one. For a tensor-valued function, result.backward seed lets us specify how much each output contributes. In calculus terminology this is a vector–Jacobian product (VJP); we propagate those output sensitivities back to the input without materializing a full Jacobian.

On GPU, LibTorch supplies supported native tensor values and local gradient operations with its own autograd engine disabled. Configured binary arithmetic keeps its complete format on the encoded tape. Custom functions therefore record through the differentiation system we've already used for the model examples.

Conditionals need care here. Recording follows the chosen branch; an unused branch cannot add an invalid division to that execution. The resulting derivative describes that path. At a branch boundary, a mathematical derivative may not exist. Scalar branch decisions on GPU also synchronize back to Lean, while the supported arithmetic and backward operations stay on the device.

Backward consumes its recording. If we only need the forward value, result.close releases the saved intermediates. That lifetime becomes important when a calculation keeps a large tensor at every step.

Let's use a function of our own as a training objective. We'll compare three predictions with three target values and sum their squared residuals:

import NN.Kernel
import NN.Tactic.Autograd

open TorchLean

def targets : Tensor Float32 [3] := [0, 1, 1]
def predictions : Tensor Float32 [3] := [1, 2, 3]
def loss (prediction : Tensor Float32 [3]) : Float32 :=
  let residual := prediction - targets
  Tensor.sum (residual * residual)

-- Optimize the predictions directly to make this one update easy to inspect.
#eval do
  let result ← loss.run predictions (grad := true)
  let gradient ← result.backward
  let next := predictions - Tensor.full [3] 0.25 * gradient
  IO.println result.value
  IO.println gradient
  IO.println next
  IO.println (loss next)

def gpuLoss := loss.run predictions (device := gpu) (grad := true)

-- This is the real derivative of one loss coordinate, with its target fixed.
example (target x : ℝ) :
    HasDerivAt (fun x : ℝ => (x - target) * (x - target)) (2 * (x - target)) x := by
  autograd

6.000000
[2.000000, 2.000000, 4.000000]
[0.500000, 1.500000, 2.000000]
1.500000

The residuals are [1, 1, 2], so the loss is six and the gradient is [2, 2, 4]. Subtracting one quarter of that gradient gives [0.5, 1.5, 2]. Its residuals are [0.5, 0.5, 1], and the loss falls to one and a half. Every number in this example is exactly representable in binary32.

We're optimizing the predictions directly here to make the update easy to follow. In a network, those predictions depend on parameters through earlier layers; their gradient would have to continue through those layers before the optimizer updated the parameters. The targets are captured constants in this function, so this call differentiates the supplied predictions, not the targets.

The real derivative proof below the run covers one loss coordinate. For a fixed target \(t\), the derivative of \((x-t)^2\) is \(2(x-t)\). It is checked with the existing autograd tactic. The tensor recording also uses the existing operation rules for subtraction, multiplication, and summation. We didn't add a loss-specific tape or a custom native operator.

For this objective, a learning rate \(\eta\) changes each real residual \(e\) to \((1-2\eta)e\). The new loss is therefore \((1-2\eta)^2\) times the old loss. At \(\eta=1/4\), that's a quarter of the original, exactly as the small run shows. Rates between zero and one decrease a nonzero loss in exact arithmetic; a rate greater than one increases it. This calculation concerns independently adjustable predictions and this quadratic objective, not a convergence claim about arbitrary neural-network training.

Changing the objective gives us a new function to inspect and a new derivative question. We can keep using the same runtime to record its gradient and run an update, then prove the mathematical rule behind that calculation.

Writing contractions with einsum

Let's move from a function on individual entries to a calculation that combines them. A linear layer takes several input features, multiplies them by weights, and adds the products. Attention does this twice: first to compare queries with keys, then to combine values using the attention weights. We could write an indexing loop for each calculation, but I'd rather write down which axes meet and which axes survive :)

That's what einsum lets us do. The name comes from Einstein summation notation, although the pattern spells out the output axes rather than making us infer them from a convention. For a matrix product, we give the first tensor a row axis and a feature axis, and the second a feature axis and an output axis. We keep row and output; we multiply and sum over feature:

Here is that calculation using two small matrices. I also use the input twice to compute the squared length of each row:

import NN.Tensor

open TorchLean
open TorchLean.Tensor

namespace ContractionExample

def inputs : Tensor Float32 [2, 3] := [[1, 2, 3], [4, 5, 6]]
def weights : Tensor Float32 [3, 2] := [[1, 2], [3, 4], [5, 6]]

-- Keep row and output; multiply and sum along feature.
def projected := einsum inputs, weights "row feature, feature output -> row output"

-- Both operands read the same row; only feature is summed out.
def squaredNorms := einsum inputs, inputs "row feature, row feature -> row"

#check projected
#eval projected
#eval squaredNorms

end ContractionExample

ContractionExample.projected : Tensor Float32 (Shape.ofList [2, 2])

[[22.000000, 28.000000], [49.000000, 64.000000]]

[14.000000, 77.000000]

Read the pattern from left to right. The comma separates the operands. Reusing feature connects their feature coordinates; the arrow lists the axes of the result. Because feature is absent after the arrow, the calculation sums it out. row and output stay, in that order. Lean infers a [2, 2] result from the input shapes [2, 3] and [3, 2].

The first output is \(1\times1+2\times3+3\times5=22\). The next one uses the same input row and the second weight column, giving \(1\times2+2\times4+3\times6=28\). The second input row gives 49 and 64. We can follow all four entries without looking inside the generated loop. These small integer results are exactly representable in binary32.

The second pattern, row feature, row feature -> row, keeps row in both operands and in the output. It combines matching entries within each row, then sums their products over feature. The results are \(1^2+2^2+3^2=14\) and \(4^2+5^2+6^2=77\). Repeating an axis name across operands doesn't mean we always sum it: whether it appears in the output is what decides that. This is also how we keep a batch or attention-head axis while contracting another one.

I like having those axis names in the expression. We can read feature output as a weight matrix mapping incoming features to outputs, and see the same feature name in the input. Changing the number of rows or outputs changes the tensor dimensions, not the meaning of the pattern. The elaborator checks their compatibility when it constructs the calculation. A shared non-singleton axis with incompatible lengths is rejected before evaluation; singleton broadcasting is a supported case, so matching axes need not always have identical physical lengths.

The names belong to this pattern. They are not permanent semantic labels stored in the tensor type. If two dimensions happen to be equal, swapping their intended roles can still produce a well-shaped but wrong calculation. The shape check catches incompatible sizes; a theorem relating the result to our intended equation checks a different property.

Moving axes, broadcasting, and reducing

Once the contraction is written, we often need to arrange its input or output for another operation. Let's use the same two-by-three matrix to transpose it, flatten it, add a bias to every row, and sum each row:

import NN.Tensor

open TorchLean
open TorchLean.Tensor

namespace ShapeExample

def matrix : Tensor Float32 [2, 3] := [[1, 2, 3], [4, 5, 6]]
def bias : Tensor Float32 [3] := [10, 20, 30]

-- Move coordinates, rather than regrouping the unchanged flat payload.
def transposed := rearrange matrix "row column -> column row"
def flattened := rearrange matrix "row column -> (row column)"

-- Explicitly give each row the same three bias values.
def shifted := matrix + (expand bias "column -> row column" with row := 2)
def rowSums := reduce matrix "row column -> row" by sum

#eval transposed
#eval flattened
#eval shifted
#eval rowSums

end ShapeExample

[[1.000000, 4.000000], [2.000000, 5.000000], [3.000000, 6.000000]]

[1.000000, 2.000000, 3.000000, 4.000000, 5.000000, 6.000000]

[[11.000000, 22.000000, 33.000000], [14.000000, 25.000000, 36.000000]]

[6.000000, 15.000000]

rearrange matrix "row column -> column row" moves coordinates. The first column [1, 4] becomes the first row, so the result is [[1, 4], [2, 5], [3, 6]]. Merely regrouping the unchanged flat payload into three rows of two would instead give [[1, 2], [3, 4], [5, 6]]. Both have shape [3, 2], but only the former is a transpose. The pattern specifies the coordinate correspondence as well as the output dimensions.

Parentheses combine axes. row column -> (row column) flattens the matrix into [1, 2, 3, 4, 5, 6], retaining that row-then-column ordering. Flattening the transposed tensor would give [1, 4, 2, 5, 3, 6]. This matters when a projection expects features in a particular order: equal element counts don't make two layouts interchangeable.

The bias has shape [3]. expand gives it a row axis of length two, making the same three values available at both rows before addition. The result is [[11, 22, 33], [14, 25, 36]]. We chose the new axis and its size explicitly rather than asking the reader to guess which broadcasting convention applies. Finally, reduce removes the column axis using summation, returning [6, 15].

These operations keep tensors as their inputs and outputs. I don't want readers to learn another numerical container to rearrange a matrix. The pattern says how coordinates are used, and the resulting type records the shape. For reductions, we also have to care about the arithmetic order; a real-valued sum and a rounded sum need different reasoning.

The two contractions inside attention

Now let's keep more axes. Suppose queries and keys each have shape [batch, head, token, feature]. Every head compares each query position with each key position, summing over features while keeping the batch and head separate:

The first equation produces raw dot-product scores. A complete attention layer then scales them, applies its mask, and normalizes them. The second equation combines values using those normalized weights. To make both contractions easy to inspect, I'll compute the raw scores and supply a separate small set of weights for the value calculation. Those supplied weights are illustrative; they aren't the softmax of the displayed scores.

import NN.Tensor

open TorchLean
open TorchLean.Tensor

namespace AttentionExample

-- One batch, two heads, two token positions, two features per head.
def queries : Tensor Float32 [1, 2, 2, 2] :=
  [[[[1, 0], [0, 1]], [[1, 1], [2, 1]]]]
def keys : Tensor Float32 [1, 2, 2, 2] :=
  [[[[1, 0], [0, 1]], [[1, 0], [0, 1]]]]
def scores := einsum queries, keys
  "batch head query feature, batch head key feature -> batch head query key"

-- Supply normalized weights to isolate the value contraction.
-- A complete attention layer computes these using scaling, masking, and softmax.
def weights : Tensor Float32 [1, 2, 2, 2] :=
  [[[[0.75, 0.25], [0.25, 0.75]], [[0.5, 0.5], [1, 0]]]]
def values : Tensor Float32 [1, 2, 2, 1] := [[[[10], [20]], [[30], [50]]]]
def attended := einsum weights, values
  "batch head query key, batch head key channel -> batch head query channel"
def joined := rearrange attended "batch head query channel -> batch query (head channel)"

#check scores
#eval scores
#eval attended
#eval joined

end AttentionExample

AttentionExample.scores : Tensor Float32 (Shape.ofList [1, 2, 2, 2])

[[[[1.000000, 0.000000], [0.000000, 1.000000]], [[1.000000, 1.000000], [2.000000, 1.000000]]]]

[[[[12.500000], [17.500000]], [[40.000000], [30.000000]]]]

[[[12.500000, 40.000000], [17.500000, 30.000000]]]

There is one batch and two heads. Within each head we have two query positions, two key positions, and two features. The first head's query and key vectors are the coordinate vectors, so its score matrix is [[1, 0], [0, 1]]. The second head's keys are the same coordinate vectors, while its queries are [1, 1] and [2, 1], giving [[1, 1], [2, 1]]. The pattern keeps every head's calculation separate. It never forms a dot product between a query from the first head and a key from the second.

For the value calculation, the first head reads 10 and 20. Its first query puts three quarters of the weight on 10 and one quarter on 20, producing 12.5. The second query reverses those weights and produces 17.5. The other head reads 30 and 50: equal weights give 40, and selecting the first value gives 30. We now have one output per query and head.

The final rearrange moves the head axis after query and combines head with channel. The first query's output row is [12.5, 40]; the second is [17.5, 30]. Keeping the old flat order and changing its shape would have grouped the values differently. The pattern states which head outputs belong beside one another before the next projection uses them.

A matrix product, a row norm, and attention use different patterns with the same contraction machinery. I like being able to change the axes in the expression without writing another model-specific helper. These examples run as CPU tensor calculations. An arbitrary einsum expression doesn't automatically acquire a GPU implementation or a backward rule merely because its pattern was accepted; recorded model contractions use their supported runtime operations.

einsum is a term elaborator, rather than a proof tactic. It reads the literal pattern when Lean elaborates the definition, checks it against the tensor shapes, and constructs the computation and its certificates. We aren't parsing those axis names again for every output entry at runtime. Proof tactics such as einops have a separate job: helping us prove statements about tensor transformations.

The compiler connection is worth looking at too. The executable einsum path can read its operands directly and accumulate each output without constructing a full tensor of every intermediate product. Generated code carries certificates identifying its indexing and ordered arithmetic with the reference evaluator. The lowering theorem then connects that evaluator to the independent contraction definition. Its general ordered-fold theorem needs addition and multiplication operations, not a claim that they obey exact-real algebraic laws. Reordering a contraction requires additional laws; the planner only takes that route when the scalar instances supply them.

We can therefore reuse the contraction interface with native floats or a configured FloatLib scalar while keeping its numerical interpretation visible. A proof about the ordered evaluator preserves that particular calculation. It does not make every grouping of its floating-point sum equal, or prove that an external GPU kernel followed that order. The tensor chapter has more patterns and explains the implementation choices behind them.

What fusion saves

For the matrix product above, a literal implementation could first build every product \(X_{rf}W_{fo}\), then reduce the feature axis. That temporary has three axes: row, feature, and output. With 256 rows, 1024 features, and 256 outputs, it would contain 67,108,864 binary32 values, or 256 MiB of scalar payload. The final [256, 256] result needs only 256 KiB. Those are element counts, not allocator measurements, but they show why the temporary is worth avoiding.

The checked lowering reads the operands and accumulates one output directly. It compiles the named axes into a stride plan, so a read becomes coordinate arithmetic followed by an access to the operand's storage. Singleton broadcasting and diagonal selection are part of that plan. We don't need to materialize a broadcast tensor or a tensor of all the products. The reference product tensor remains available in the mathematical semantics for the proof.

The theorem einsumTensorKernel_correct connects the generated result to that reference. It needs storage, addition, multiplication, zero, and one. It does not require floating addition to be associative: the operand products and contracted additions keep their specified order. This is especially useful for FloatLib scalars, where changing that order can change the answer. A faster regrouping needs the appropriate algebraic law or a numerical comparison, not an assumed real-number identity.

Avoiding the temporary does not turn this into a tuned GPU matrix library. These pattern examples use the compiled Lean tensor path; GPU model matrix products use the LibTorch backend. I want both routes to have clear semantics, but their schedules and performance need to be measured separately.

From tensor expressions to a model

The contraction examples make the axes visible. When I'm building a network, I also want the layer composition to be easy to read. Here is the public API used by the MLP quickstart:

import NN.API
open TorchLean

def model :=
  nn.Sequential![
    nn.linear 2 8,
    nn.relu,
    nn.linear 8 1
  ]

def trainer := Trainer.new model
  { objective := Trainer.Objective.mse
    optimizer := optim.adam { learningRate := 0.03 }
    device := cpu }

The model accepts two features, expands to eight, applies ReLU, and returns one value. Lean checks the shapes at the composition points. The trainer supplies the objective, optimizer, and execution settings; it doesn't change the model definition. The quickstart then gives it tensor data and runs the training loop. For the supported native GPU path, we can select gpu in the settings.

This separation also matters when we choose precision. The builders have seeded initializers, while the typed state records the scalar type of the parameters and buffers. If extra digits matter, we can construct that state directly from exact rationals rather than widen values already rounded by an initializer. The binary128 affine-training example later in this essay does exactly that.

The larger examples use these same ideas. A transformer composes projections, attention, normalization, residuals, and feed-forward layers. A Fourier neural operator adds spectral transforms and frequency mixing. A diffusion model repeats a learned denoising update. Those models have different state and training objectives, but they use the shared tensor, layer, and differentiation APIs. The model examples include those calculations alongside convolutional, recurrent, reinforcement-learning, and self-supervised models.

Where the time goes

I don't want “runs on GPU” to stand in for a speed result. There are several costs we can inspect: generating and compiling a custom body, launching operations, moving values, saving intermediates, and doing the arithmetic itself.

For a forward-only custom recurrence, the frontend can generate one ordered loop. With gradient recording, the existing tape records the steps and retains what backward needs. On GPU that recorded path launches operations per step. Both execute the recurrence, but their launch count and memory use can be very different. The faster route for one forward call is not automatically the route we can use for training.

The runtime also caches identical generated CUDA source. A first call can include NVRTC compilation; a warm call can reuse the compiled program. Capturing a different scalar constant can produce different source, so a benchmark should say whether it is timing compilation, warm execution, or both. Asking for a scalar branch decision on GPU can also synchronize with Lean. That is a cost to consider if we put the branch inside a long loop.

For standard model operations, we moved to LibTorch because maintaining our own numerical kernels duplicated infrastructure that PyTorch already provides. TorchLean still owns the model and differentiation traversal; ATen, PyTorch's C++ tensor library, computes the supported tensor values and local gradients, with LibTorch autograd disabled. Matrix backward can pass transpose flags to the native matrix operation instead of first creating transposed copies. These are concrete implementation choices, not a promise that every model is now faster than PyTorch.

We do have a retained end-to-end comparison from the GPT training experiment:

The comparison record contains the exact numbers. PyTorch was about 3.43 times faster on that setup, while the largest training-loss difference was about 0.00001193. Those measurements used the earlier FP32 implementation, so they are a baseline for a new comparison, not timings for the current LibTorch adapter or the custom compiler.

The cached-generation experiment in the same article asks a different performance question. For its archived checkpoint, prompt, and eleven-token continuation, three warm runs took about 168–172 ms with the cache versus 9.43–9.74 s when recomputing the full prefix. The roughly 56-fold difference comes from reusing earlier work on that workload. Its cache theorems describe what the stored keys and values mean; the timings show what that reuse saved in those runs. I'd like to keep making that kind of comparison as we improve the runtime, with the configurations and results beside the code.

Indexed reads and row sums

The earlier square reads one entry per output. A row sum reads several entries and combines them. Let's make that access pattern explicit:

import NN.Kernel

open TorchLean NN.Kernel
open scoped NN.Kernel

def rowSum (width : UInt64) : Program Float32 :=
  Program.of (fun (read : Reader Float32) (row : UInt64) =>
    iterate (fun column acc => do
      let x ← read 0 (row * width + column)
      pure (acc + x)) width.toNat 0 0)

def matrix : Tensor Float32 [2, 3] := [[1, 2, 3], [4, 5, 6]]
#eval (rowSum 3).run (Arguments.empty.push matrix) [2]

def gpuRows : IO (Tensor Float32 [2]) :=
  (rowSum 3).run (Arguments.empty.push matrix) [2] (device := gpu)

#check Program.eval_eq_reference
#check Program.lower_correct

[6.000000, 15.000000]

NN.Kernel.Program.eval_eq_reference {α : Type} [Storage α] {shapes : List Shape} (program : Program α)
  (inputs : Arguments α shapes) (shape : Shape) :
  program.eval inputs shape =
    if shape.size ≤ UInt64.size then
      Tensor.generateFlatM shape fun index => program.reference (Reader.ofArguments inputs) (↑index).toUInt64
    else Except.error (Error.size shape.size)

NN.Kernel.Program.lower_correct {α : Type} (program : Program α) (read : Reader α) (index : UInt64) :
  Target.Block.eval program.arithmetic read (Env.output index) (Target.lower program.expression) =
    program.reference read index

This gives [6, 15] again, as the earlier reduce did. The first row reads 1, 2, 3; the second reads 4, 5, 6. The rows can run independently, while additions within a row follow the order we wrote.

The body chooses a row-major address, row * width + column, and adds that value to its accumulator. read 0 names the first input tensor. Arguments.empty.push matrix supplies that tensor, and output shape [2] asks for two row results. An out-of-range read produces an error; the output shape alone can't establish that every address in the body is valid.

For a row sum, I'd use the reduce pattern above. This indexed interface is useful when we need a particular sequence of reads or an unusual accumulation. The result is still a tensor, rather than a new collection type readers have to learn.

This is also where numerical semantics affect the compiler. Over real numbers, several addition trees compute the same sum. In floating point, they can differ. The first article gave a three-number example in which moving the parentheses changes one into zero. The row-sum lowering preserves its ordered additions; it doesn't assume associativity.

An indexed program also needs more information before we can differentiate or certify it. Compiling its reads and arithmetic proves a source correspondence. It doesn't automatically supply the derivative of an arbitrary body or an enclosure over an input region. The recording examples above work because their operations already have supported local rules.

Beyond hardware precision

GPU hardware has a fixed collection of number formats. A mathematical definition doesn't have to be limited to that collection, but keeping extra digits requires an implementation of their arithmetic.

FloatLib lets us configure the representation used by a tensor. Here we'll use binary128, with fifteen exponent bits and 112 fraction bits. Choose \(a=1+2^{-100}\): the extra term is too small to survive storage in binary64.

import NN.Kernel

open TorchLean
open FloatLib.Floats

abbrev Wide := ExecFloat.Binary 15 112
def squareWide (x : Wide) : Wide := x * x

-- Construct the extra digit from a rational, without passing through binary64.
def a : Wide := Rat.cast (1 + 1 / (2 ^ 100 : Nat) : Rat)
def wideInput : Tensor Wide [1] := [a]
def gpuWide : IO (Tensor Wide [1]) := squareWide.run wideInput (device := gpu)

#eval do
  let output ← squareWide.run wideInput
  pure (ExecFloat.Binary.toRat? output[0] == some (1 + 1 / (2 ^ 99 : Nat) : Rat))

true

Constructing \(a\) from a rational keeps that digit before any execution begins. If we first converted the value to binary64 and then widened it, we'd get one. A wider destination can't recover information the earlier conversion lost. The exact square is \(1+2^{-99}+2^{-200}\). Binary128 nearest-even rounding keeps \(1+2^{-99}\) and loses the much smaller final term. The example compares the CPU result's exact rational value with that rounded answer, so true means the retained digit is there. Printing through a binary64 conversion would be a poor way to inspect it.

For the custom GPU path, wide values travel as complete integer words. Software arithmetic on the device operates on those words rather than narrowing them to a native float. We retain the chosen precision, at the cost of more instructions than a hardware multiply. This example concerns the arithmetic result; it gives no throughput measurement.

Standard configured binary formats can use compatible hardware operations. Selection depends on the complete format, operation, and device capability. When those don't match, the configured backend keeps its software path. Separate products and sums also remain separate, so the generated arithmetic must preserve that choice rather than silently replacing the pair with one fused multiply-add.

Proving the function and its lowering

Once we've written a calculation, we can ask a mathematical question about it. For the square, we can prove the real derivative using the existing calculus tactic:

import NN.Kernel
import NN.Tactic.Autograd

open TorchLean

-- A real derivative statement uses the existing calculus tactic.
example (x : ℝ) : HasDerivAt (fun x : ℝ => x * x) (2 * x) x := by
  autograd

The autograd tactic uses derivative rules to construct a proof, which Lean's kernel checks. Its conclusion holds for every real input. A rounded floating-point function isn't a smooth real function; numerical error bounds are what connect its values to the real calculation.

For compilation, we need to show that the generated loop returns the result of the original recurrence, including what happens when an input read fails. When we name intermediate results, later instructions must read the intended value, rather than another value with the same type.

Program.eval_eq_reference identifies the checked evaluator with the source calculation. Program.lower_correct carries that result into the structured loop representation. Their types appear with the row-sum example above. The native FP32/FP64 emitted-source theorem has explicit arithmetic and input-buffer contracts; it doesn't verify NVRTC or the GPU. The configured-format CUDA arithmetic is checked against FloatLib in runtime tests and remains outside the kernel-verified part.

This distinction helps me decide what to prove next. A function can have correct lowering while lacking a numerical error theorem. A derivative rule can be correct over the reals while the computed gradient rounds. A tensor can have the right shape while its values are wrong. Each property has a statement we can inspect and improve.

For the examples on this page, Lean checks the displayed code and records the CPU outputs. The GPU calls are typechecked, and the scalar examples have accepted frontend compilation results; this article doesn't present a fresh GPU run. The custom computation guide has the device setup and the full API.

The relaxation example connects three things that are often maintained separately. Its definition accepts the operations it needs from an element type. We instantiate it with Float32 to run a tensor computation, with real numbers to prove the closed form, and with a GPU-supported format to obtain a lowered body with a source-correspondence proof. We can inspect those connections in one development.

The square, blend, and loss are familiar calculations; the additional work is connecting execution to statements about their meaning. The relaxation theorem covers every real input and step count, rather than only the starting values displayed above. The compiler theorem preserves the specified read and arithmetic order. The numerical definitions let us describe rounding when a real-valued equality isn't the right description of a rounded result.

PyTorch's autograd records tensor operations for differentiation, and torch.compile provides a compiler interface for Python/PyTorch computations. We reuse LibTorch's numerical infrastructure because those execution capabilities are valuable. Here we additionally develop Lean statements about the source functions and transformations, reusing mathlib's calculus and algebra instead of maintaining a second informal argument beside the code.

Formal numerical reasoning also exists elsewhere. Flocq provides multi-radix, multi-precision floating-point theory and computations in Rocq. Other languages and proof assistants can support related workflows. In this library, we've brought the tensor API, differentiation runtime, compiler proofs, mathlib, and FloatLib together. The examples show how those pieces work together; they don't establish that other languages couldn't express these programs.

Other features we've been working on

I started with a square because we can understand its value and derivative immediately. We've also been working on the calculations around it: arranging axes, composing layers, taking derivatives, and choosing the arithmetic. The contraction examples above are one part of that work. I'd like the model and proof APIs to be as convenient to use as those patterns.

The differentiation work goes beyond one backward pass too. Nested dual numbers represent higher-order directional derivatives, and the operation theorems connect them to mathlib's iteratedFDeriv at every finite order. We can combine supported arithmetic and smooth functions, including local rules for division, logarithms, and square roots when their domain conditions hold. A new named function can expose its definition or register a proved rule for the autograd tactic to reuse. I like this because the proof can explain why a derivative calculation works, including its assumptions, instead of leaving us with only a printed Hessian. The tactic examples show the scalar and tensor statements; the scientific-model example uses differentiation in models for numerical problems.

Differentiating through a differential equation

The physics-informed neural network (PINN) example makes the higher-derivative API concrete. We ask a smooth network to represent a function \(u_\theta(x)\) satisfying \(u''(x)=-2\) on \([-1,1]\), with zero values at both endpoints. Its loss penalizes the equation residual at interior points and the boundary values. We don't give training samples of the known solution \(1-x^2\).

The input direction has the same shape as an input tensor. Stacking that direction twice asks autograd.model.derivative for a second directional derivative; using two different directions asks for a mixed derivative. The order is the length of that leading axis, not a separate function named after every derivative order.

Training on the residual also needs a parameter derivative. The second input derivative still depends on the network's weights, so autograd.model.derivativeVjp pulls an output cotangent (the supplied output sensitivities) through that calculation to the state. The implementation nests dual scalars around the existing model evaluator and reverse pass. We aren't approximating \(u''\) by subtracting nearby predictions.

The example uses tanh because the equation needs a smooth field. ReLU's branch derivative does not give a classical second derivative at its kink. More generally, the real higher-derivative theorems keep their smoothness and domain hypotheses. A low sampled residual is useful training evidence, while a guarantee over the whole interval needs a uniform residual bound and a theorem connecting it to solution error.

For a small scalar proof, the existing tactic can check a proposed derivative:

import NN.Tactic.Autograd

example (x : ℝ) : HasDerivAt (fun y : ℝ => Real.exp (y * y))
    (2 * x * Real.exp (x * x)) x := by
  autograd

This is the chain rule for an exponential composed with a square. We can inspect the mathlib derivative statement independently of a floating-point gradient run. For a new named function, unfolding its definition or registering a proved rule lets the same tactic work with it. It doesn't need a new differentiation language.

The Burgers FNO example uses a different objective: learn a map from an initial field to a later solution. Fourier transforms and frequency mixing become layers in that model. Both examples use the tensor and training APIs, but the statement we would prove about a trained PDE surrogate depends on the equation, boundary conditions, and error notion we choose.

Keeping precision through training and saving

The binary128 square earlier only follows one value. Typed training keeps the selected scalar in the samples, model state, predictions, gradients, and losses too. trainer.openTyped opens that CPU session, and its checkpoint encoding preserves the scalar words when saving and loading. The precision API describes these operations and the conversion boundaries.

There are still choices to make at the beginning of the calculation. Widening seeded Float parameters cannot create random bits that weren't there. The trainer's optimizer settings also start as Float values before conversion. When the learning rate itself needs extra precision, the typed nn.sgdStep call accepts it in the selected scalar type. That's why the affine example constructs both the parameter and learning rate from rationals. Saving all the resulting words is useful only if we kept those digits on the way into the computation.

Complex values use the same idea with two components. Complex α stores the real and imaginary parts in α; a tensor of complex values keeps its shape as usual. The complex regression example minimizes a real squared-residual loss, differentiating both component coordinates with the forward-mode reference method. This is useful for complex parameters without identifying a real loss gradient with a holomorphic complex derivative.

Starting from JAX or PyTorch

Sometimes the model we want to inspect already exists in Python. The Jaxpr example captures supported operations into TorchLean's tensor intermediate representation (IR). Jaxpr is JAX's traced representation of the calculation, rather than the Python source itself. The importer checks shapes and precision metadata, then gives that graph a Lean semantics. PyTorch has a supported graph-import path as well. We are still extending the coverage; these adapters aren't translators for arbitrary Python programs.

The small JAX example captures sum(x * x). Its theorem first identifies the graph's real-valued denotation with the tensor sum of squares, then proves that permuting a vector preserves the graph's answer. The permutation result covers every admitted vector length and permutation, rather than only the input used to compare the JAX and Lean runs. Floating-point permutation invariance is another question because reordering a reduction can change its rounded answer, as we saw in the first essay.

This is the direction I'd like to develop further: keep the frontend someone already uses, bring its supported computation into a common representation, and work on a theorem about that graph. The importer, numerical execution, and graph property each need their own evidence. A successful sample comparison is a good start, but it doesn't prove the translation correct for every input.

Bounds and model verification

On the proof side, the updated FloatLib interval tactic handles more real expressions, including positive-base powers, logarithms to a chosen base, and inverse hyperbolic functions. We can register another function's enclosure together with its containment proof, and reuse that scalar theory when bounding tensor calculations. This matters when an equality over real numbers isn't enough: a rounding bound or an enclosure can describe the behavior we need to protect. Those proofs remain distinct from assumptions about external GPU execution. The floating-point chapter has examples of the arithmetic and interval proofs.

For classifiers, the verification API lets us check a trained model over an input region, asking whether a label survives the allowed perturbations. A supported certificate and its soundness hypotheses supply the guarantee; switching on gradients alone doesn't. For a ReLU classifier we might want a fixed label on an input box. For a smooth neural field we might want a residual bound. Those are different statements, even when both start from tensors and the same model APIs.

Convergence proofs also need a statement about the update. The converges tactic can apply registered contraction and gradient-descent results, retaining their step-size, smoothness, and other hypotheses. It doesn't infer convergence from a decreasing loss curve. I want to keep these tools close to the runnable examples so we can work on the property a calculation needs, with mathlib's definitions available underneath.

Conclusion

Try changing the update in the relaxation example, or the objective in the custom loss. The execution API stays the same; the closed-form and derivative proofs need to match your new function. For a starting value away from the target, try a rate of one half, then two: the real recurrence switches from approaching the target to oscillating around it. Keeping the theorem and the run together lets us ask which behavior we meant to implement :)

References

• TorchLean. Custom tensor computations and compiler source. The execution API and lowering theorems used here.
• FloatLib. Arithmetic and proofs. Configured formats and numerical theory.
• David Thrane Christiansen. Functional Programming in Lean. Lean's programming and proof facilities.
• Accelerate. The Accelerate language. Embedded array computations.
• Troels Henriksen and colleagues. Futhark: Purely Functional GPU-Programming with Nested Parallelism and In-Place Array Updates, 2017.
• Adam Paszke and colleagues. Getting to the Point: Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programming, 2021. Dex.
• Higher Order Company. Bend and its language guide. Dependent types, ownership, and parallel calls.
• NVIDIA. NVRTC documentation.
• PyTorch. Autograd mechanics and torch.compile.
• Sylvie Boldo and Guillaume Melquiond. Flocq. Formal floating-point arithmetic in Rocq.

@misc{george2026programming,
  author = {Robert Joseph George},
  title = {{Writing tensor programs in Lean}},
  year = {2026},
  howpublished = {TorchLean Verified Examples},
  url = {https://robertboy18.github.io/TorchLean-Verified-Examples/torchlean-tensor-programming/}
}