Running Lean Files🔗

Create a file named Primer.lean at the TorchLean repository root and run it:

# Check the saved primer file in the repository environment.
cd TorchLean
lake env lean Primer.lean

lake env runs the command in the project's environment, including the search paths for its Lean libraries. The checkout pins its toolchain in lean-toolchain. Use this command from the repository root so that NN.API resolves against the project and its dependencies.

When every command in the file succeeds, Lean prints nothing except the output you asked for with #check, #eval, and friends, and exits with status zero. With the imports and namespace opening shown in the next section, a file ending in

-- Name a length-two Float tensor, then evaluate its
-- entries.
def point : Tensor Float [2] := [0.25, -0.75]
#eval point

prints exactly one line:

[0.250000, -0.750000]

Repository Lake targets enable warnings-as-errors. A direct lake env lean Primer.lean invocation sets the search environment but does not automatically apply every target-specific compiler option; inspect any warnings as well as errors.

Imports And Namespaces🔗

Begin the file with:

-- Import the application definitions and make their short
-- names available.
import NN.API

open TorchLean

import NN.API loads the application interface: tensors, model builders, datasets, optimizers, and the high-level trainer. Proofs and backend internals have additional imports, introduced where the examples need them.

Most application-facing names live in the TorchLean namespace. open TorchLean lets us write Tensor, nn.linear, and Trainer.new instead of prefixing each one with TorchLean..

Lower-case nn and optim are namespaces. In Python, import torch.nn as nn binds a module object to a variable, and nn.Linear is an attribute lookup performed while the program runs (Paszke et al., 2019). Here, nn.linear is resolved by the elaborator before execution, and there is no nn value you could pass to a function or reassign.

Lean can show you the type of any name it knows:

TorchLean.Tensor (α : Type) (shape : Shape) [Storage α] : Type

So a tensor type takes two arguments, an element type and a shape, plus one instance argument describing how the elements are stored. Now ask about a layer:

nn.linear 2 4 : nn.Builder (nn.Sequential (Shape.appendDim [] 2) (Shape.appendDim [] 4))

The widths 2 and 4 appear in the type. A 2 → 4 linear layer and a 4 → 2 linear layer therefore have different types, allowing subsequent operations to require the appropriate widths.

Lean prints the input shape as Shape.appendDim [] 2: the layer's definition constructs it by appending an innermost axis to the empty shape. appendDim is marked @[reducible], so Lean can unfold it during type checking and accept [2] in its place:

Printed types can retain expressions such as appendDim rather than reducing every dimension. When two shapes look different, check whether their expressions compute the same list.

One more, this time a definition with implicit and instance arguments and a default value:

TorchLean.Trainer.new.{u} {Model : Type u} {σ τ : Shape} [Trainer.ToModel Model σ τ] (model : Model)
  (config : Trainer.Config σ τ := { }) : Trainer σ τ

Start with the explicit argument (model : Model) and the result Trainer σ τ. Lean infers the curly-braced arguments, including the shapes, and searches for the square-bracketed ToModel instance that connects this model representation to the trainer. The printed .{u} is a universe parameter describing how general Model's type may be.

The last argument has a default: := { } fills the configuration with its default field values. Thus Trainer.new model is a complete call, with an objective and runtime settings even though we have not written them out.

Definitions🔗

def gives a name to a value or a function:

144

The first Nat is the argument type and the second is the result type. Lean can infer many result types. An explicit annotation at an API boundary also checks the implementation against the intended interface, locating a mismatch at the definition.

Local values are introduced with let:

16

The let binds displacement to x - center for the rest of this expression. Later uses refer to that value, so we can substitute its definition when reasoning about the result.

Tensors🔗

The type of a TorchLean tensor has two arguments:

-- Read alpha as the element type and s as the list of axis
-- lengths.
Tensor α s

α is the element type and s is the shape; a storage instance supplies their representation. Because the type depends on the value s, it is a dependent type. This lets an operation express relationships between dimensions in its argument and result types.

point : Tensor Float [2]

[0.250000, -0.750000]

#check elaborates an expression and reports its type; #eval executes it and displays the result. The first line tells us that point has two Float entries, and the second tells us which entries they are. The distinction also matters for effectful expressions: checking a prediction method's type leaves the model unexecuted.

The bracket literal is not a List Float that happens to be used as a tensor. A custom elaborator reads it against the expected type, checks each axis, and builds the tensor's buffer directly, which is also why its error messages talk about dimensions rather than about lists.

The shape index is itself a Lean value:

[2, 3] : List ℕ

A shape is a list of natural numbers, outermost axis first. Rank and element count are computed from that list:

2

6

1

The rank of [2, 3] is two because there are two axes; its size is six because there are two rows with three entries each. The empty shape has rank zero but size one, following the empty-product convention. This is how a scalar fits into the same tensor interface as a matrix. An axis of length zero would describe a different object with no entries. When a signature uses [] for a loss, it means one scalar value, not an empty batch or a missing result.

Both TorchLean and PyTorch track shape information. TorchLean's tensor type makes it available when checking the surrounding program:

The Basic Operations🔗

With shapes in their types, tensor operations can require equal dimensions or describe how an output shape is computed. Two vectors illustrate the equal-shape case:

[4.250000, 7.250000]

[1.000000, -6.000000]

+ works because tensors of a fixed shape form an additive structure, and Tensor.mul is the elementwise product, PyTorch's *. Both require equal shapes; neither one broadcasts silently.

A dot product is a product followed by a sum, and you can spell it either way:

-5.000000

-5.000000

The two expressions have the same specification: dotSpec is defined as sumSpec (mulSpec a b) in Spec/Core/.../Reductions.lean. The expanded form shows which products are summed; the named operation gives proofs a reusable definition to refer to.

Reductions collapse a tensor to a scalar:

-0.500000

-0.250000

0.312500

The sum is 0.25 - 0.75 = -0.5, and dividing by two gives the mean -0.25. For mean squared error, we square the deviations first: 0.0625 and 0.5625 average to 0.3125. Squaring before reducing prevents the positive and negative deviations from cancelling.

Tensor.matmul uses the same type-level dimension for the left matrix's columns and the right matrix's rows. The arguments must agree on that dimension:

[[4.000000, 5.000000], [10.000000, 11.000000]]

Matrix times vector and vector times matrix are separate operations rather than one overloaded product with a promotion rule:

[6.000000, 15.000000]

[5.000000, 7.000000, 9.000000]

In PyTorch, a @ b covers vector times vector, matrix times vector, matrix times matrix, and batched products, selecting the operation from the ranks. TorchLean uses separate names with checked rank contracts, so matvec states the intended operation at the call site. Dependent typing also permits overloaded operations when their promotion rules are specified.

ReLU preserves the shape, while reshape changes its axes without changing the element count:

[0.000000, 2.000000]

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

The reshaped matrix keeps the flat order 1, 2, 3, 4, 5, 6, grouping it into three rows of two. A transpose would exchange row and column coordinates and rearrange that order. Both produce shape [3, 2] from [2, 3], so their shapes alone cannot distinguish them. ReLU changes the negative entries while preserving both shape and coordinate layout.

reshape takes its target shape as a named argument and checks that the two shapes have the same size while elaborating, so there is no runtime reshape failure to handle and no -1 to infer.

To add a vector to every row of a matrix, first repeat it along the leading axis. The addition then receives two tensors with the same shape:

[[10.000000, 20.000000, 30.000000], [10.000000, 20.000000, 30.000000]]

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

In PyTorch the same result is left + row, using broadcasting. Broadcasting can also combine torch.ones(2, 1) * torch.ones(1, 3) into a 2 by 3 matrix. That outer product is useful when intended, but it would conceal a mistaken axis in an intended elementwise product. The explicit repetition in TorchLean identifies which axis is being added.

The corresponding PyTorch expressions give the following results:

>>> # Compare the scalar dot product, matrix contraction,
>>> # and a broadcast result's shape.
>>> import torch
>>> a = torch.tensor([0.25, -0.75]); b = torch.tensor([4.0, 8.0])
>>> torch.dot(a, b).item()
-5.0
>>> (a * b).sum().item()
-5.0
>>> left = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
>>> right = torch.tensor([[1., 0.], [0., 1.], [1., 1.]])
>>> left @ right
tensor([[ 4.,  5.],
        [10., 11.]])
>>> (torch.ones(2, 3) + torch.ones(3)).shape
torch.Size([2, 3])

Tensor Indexing And Updates🔗

A coordinate into a shape has its own type, Shape.Coord, built from the shape:

point.at : Spec.Shape.Coord [2] → Float

0.250000

6.000000

A coordinate for [2] is a Fin 2 paired with the unit value, and a coordinate for [2, 3] is a Fin 2 paired with a coordinate for [3], which is why the second read is written (1, (2, ())). The coordinate type unfolds to this nested pair:

Fin n pairs a natural number with evidence that it is below n. The PUnit terminates the coordinate tuple; a matrix coordinate has a bounded row and column index before that terminator. Here rfl unfolds the coordinate type. When we subsequently look up an entry, the bounds travel with its indices and need no further check.

A numeral at a Fin type requires care, however:

-0.750000

There is no entry 5 in a length-two vector. What happened is that the numeral 5 was elaborated at type Fin 2, and Lean's OfNat instance for Fin reduces numerals modulo the bound, so 5 became 1 and we read the last entry. The access is in range, but it need not be the coordinate the programmer intended.

A Fin 2 value carries evidence that its stored natural number is below two; it does not record whether that number came from a wrapped numeral. For coordinates supplied as runtime naturals, the checked accessor preserves an out-of-range result as a failure:

some 6.000000

none

The first array names a valid row and column and returns some 6. Row five in the second array exceeds the matrix's two rows, so it returns none. Use this interface for coordinates read from a file or supplied by a user when an out-of-range coordinate should fail.

PyTorch raises IndexError: index 5 is out of bounds for dimension 0 with size 2 for the same index. The checked Lean accessor represents the outcome with Option Float: some contains a value and none indicates an invalid coordinate. A caller must account for both to obtain a Float.

Shape Errors🔗

The following examples exercise three shape checks. Each named block is required to fail during the guide build, and its diagnostic is checked against the output shown.

The first is a literal whose length does not match its shape:

tensor literal has leading dimension 3, but the expected leading dimension is 2

The tensor literal elaborator checks the dimensions against the expected tensor type and reports the mismatch directly. Its implementation is in Internal/Elab/TensorLiteral.lean.

The second is two tensors that must agree and do not. The helper takes one shape variable for both arguments, which is how you say "the same shape" in a type:

Application type mismatch: The argument
  rowVector
has type
  Tensor Float [3]
but is expected to have type
  Tensor Float [2]
in the application
  requireSameShape point rowVector

The same requirement extends to composed layers. This model connects matching widths:

nn.build 2026 primerModel : nn.Sequential [2] [1]

The first linear layer maps a length-two vector to length eight, nn.relu preserves that shape, and the second consumes length eight and returns length one. Now change one number, 8 to 7, in the final layer:

Application type mismatch: The argument
  bc✝
has type
  nn.Sequential (Shape.appendDim [] 7) (Shape.appendDim [] 1)
but is expected to have type
  ?m.67 a✝ __r✝¹ bc✝ __r✝ (Shape.appendDim [] 8) [1]
in the application
  nn.compose a✝ bc✝

The names ending in ✝ are generated by the nn.Sequential! macro, and ?m.67 is a metavariable Lean had not yet solved. The dimension mismatch identifies the problem: a sequence starting at width 7 was supplied where one starting at width 8 was required. To change widths here, the model needs an operation whose type connects them.

>>> # Construction succeeds; the first forward call exposes
>>> # the incompatible hidden widths.
>>> import torch, torch.nn as nn
>>> net = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(7, 1))
>>> net
Sequential(
  (0): Linear(in_features=2, out_features=8, bias=True)
  (1): ReLU()
  (2): Linear(in_features=7, out_features=1, bias=True)
)
>>> net(torch.zeros(1, 2))
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x8 and 7x1)

>>> # Exercise both branches: only the matrix product needs
>>> # the incompatible extent seven.
>>> def scale(x, flag):
...     if flag:
...         return x @ torch.zeros(7, 1)
...     return x * 2.0
...
>>> scale(torch.zeros(1, 2), False).shape
torch.Size([1, 2])
>>> scale(torch.zeros(1, 2), True)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x2 and 7x1)

Structures🔗

A structure groups named fields:

2.500000

deriving Repr asks Lean to generate the printing code, which is why #eval on an Affine value would show its fields. Note that Affine.forward is called as p.forward x as well: any function whose first explicit argument has type Affine can be written with dot notation, and TorchLean's whole session and trainer API is built on that convention.

The same record syntax lets us configure a trainer:

-- Set the loss, Adam learning rate, and initialization seed
-- in a configuration record.
{ objective := .meanSquaredError
  optimizer := optim.adam { learningRate := 0.03 }
  seed := 2026 }

Lean recovers this literal's type from the expected type at the call site. That also tells it which type .meanSquaredError belongs to.

Inductive Types🔗

An inductive type enumerates the possible forms of a value:

"mean"

The definition of reductionName matches one branch to each constructor. Omitting a case leaves the function undefined for that alternative, which Lean reports:

Missing cases:
Reduction.sum

TorchLean's execution modes and graph operations are inductive types too. Adding a constructor exposes matches that listed only the old cases. A wildcard branch remains exhaustive, so its fallback behavior needs review when the alternatives change.

Exhaustiveness checks that every case has a branch. Whether an unsupported combination returns an error depends on the branch's implementation; the type checker does not choose that policy.

Implicit Arguments And Instances🔗

We already used an implicit argument above. In

the shape s in curly braces is written once and used twice. Callers can usually omit it: Lean infers it from the first tensor and requires the second to match. Reusing a shape variable expresses equality of dimensions without asking callers to pass and check separate integers.

Square brackets introduce instance arguments, resolved by search rather than by inference:

"7, 7"

Lean infers Nat from the argument seven, then finds a ToString Nat implementation. The Storage α argument in the chapter's first #check works similarly: it supplies the operations for storing tensor entries. TorchLean also requests arithmetic, decidable shape equality, and printing through instances.

This lets a definition serve Float tensors at runtime and real-valued tensors in proofs. The operations alone do not supply algebraic laws. A theorem using associativity must obtain that law from its scalar type or hypotheses; replacing reals with floats can preserve the program's structure while changing which theorems hold.

Scalar Types🔗

A declaration such as

can be instantiated with Float, with Rat, or with FloatLib binary32. Within any one instantiation, the input and the output use the same α. Nothing in the type lets a run start in one element type and finish in another, and Tensors And Shapes develops what that means for a whole training run.

Programs, Propositions, And Proofs🔗

Lean uses the same language for programs and claims about programs. Their types distinguish computing a result, stating a proposition, and proving it.

A definition that returns data computes:

1

A definition that returns Prop states a claim, and states it without proving anything:

HasOneOutput names the proposition. To establish it, we supply a term of that type:

rfl closes the goal because both sides reduce to the same natural number. When Lean accepts that declaration, the kernel has checked a proof of exactly outputWidth = 1, and you can ask what the proof relied on:

'hasOneOutput' does not depend on any axioms

#print axioms lists the axioms on which a theorem depends, including dependencies reached through other declarations. A proof using sorry depends on a placeholder axiom that appears in this list. Inspecting it is one part of auditing the result.

Mathematical results can also depend on the standard logical foundations:

'TorchLean.Tensor.mean_eq_meanSpec' depends on axioms: [propext, Classical.choice, Quot.sound]

The command follows theorem dependencies down to their axioms. Here propext identifies logically equivalent propositions, Classical.choice supplies choices from existence, and Quot.sound supports quotient constructions. These are standard foundations used throughout Mathlib (Community, 2020); individual theorems may depend on fewer of them.

To audit this result, read its conclusion and hypotheses alongside that list, checking for unexpected axioms or unfinished proofs. The report concerns the named theorem, so it gives no guarantee about a compiler or foreign kernel outside that theorem's statement.

The displayed theorem connects a public operation to its specification:

Unfolding the public Tensor.mean wrapper gives Tensor.meanSpec, so rfl closes the equality without a tolerance or a numerical evaluation. The program and proof use the same definition. Relating that computation to an ideal real mean, or to a native reduction with another evaluation order, requires a further theorem. The Specification Layer explains which connections follow by unfolding and which need additional proofs.

The same separation appears throughout verification. An interval pass computes a candidate output box: that is a program, and running it produces data. A predicate states that the box contains every semantic output: that is a proposition. A soundness theorem proves the predicate under explicit assumptions: that is a proof. Returning a candidate box therefore does not by itself establish containment. The theorem and its hypotheses supply that connection.

Effects And IO🔗

A pure function returns its value directly:

Printing is an effect, so the result type says so:

next = 5

IO Unit describes an action whose return value carries no useful data. The text next = 5 is printed as an effect; printIncrement does not return that string. A training script likewise prints losses separately from returning the trained model that a later prediction call needs.

The do block sequences effectful actions, and s!"..." is string interpolation. Training sessions, checkpoint access, subprocess calls, and native runtime entry points all live in IO; pure shape calculations and mathematical specifications return values directly. External data needed by a pure specification must be supplied as an argument, where it can also appear in the statement of a theorem.

In an executable module, an entry point such as main : List String → IO UInt32 receives the command-line arguments and performs the run, which is how lake exe torchlean reaches the examples.

Immutability And Storage Reuse🔗

Tensor updates in this chapter return values while preserving any earlier values still in use. That semantic guarantee does not require copying the entire buffer on every update.

Unique ownership can avoid a full payload copy for operations such as scalar updates. Allocation, index calculation, and reference-counting overhead still depend on the operation. Lean's compiler can update uniquely referenced objects in place, as described in Counting Immutable Beans (Ullrich and de Moura, 2019). The following loops expose two ownership patterns by repeatedly updating entry zero. The first keeps no reference to the previous value, so its tensor is uniquely owned; the second holds original alive across every update, so it is not:

Both functions return zero for zero updates and the last assigned index as a Float otherwise. Their return values hide the storage difference: one loop replaces its current tensor, while the other repeatedly updates from an original it must preserve.

To measure that difference, time both with IO.monoNanosNow, repeat with a smaller tensor, and report the execution mode and update count. #eval and a compiled executable can have different costs. These loops provide an experiment, not a timing result; changing the tensor size helps reveal whether a measured cost grows with the amount of data that must be copied.

A retained value also matters to differentiation: the backward computation may need a value from the forward pass. PyTorch tracks mutation constraints and rejects, for example, this update:

>>> # Updating a view of this leaf tensor is rejected while
>>> # gradient tracking is enabled.
>>> w = torch.zeros(4, requires_grad=True)
>>> w[0] = 1.0
RuntimeError: a view of a leaf Variable that requires grad is being used in an
in-place operation.

PyTorch rejects this tracked in-place update on a view of a leaf requiring gradients, even before a forward tape has recorded it. With TorchLean's immutable tensor interface, a reference retained by the tape keeps the old value available. Storage reuse must respect that sharing, so an update cannot overwrite a buffer while another reference still needs its contents.

Reading Types🔗

Long TorchLean types get easier when you read from the result backwards. Take one from earlier in this chapter:

-- Read the result from the inside out: a [2]-to-[4]
-- sequence produced by a seeded builder.
nn.linear 2 4 :
  nn.Builder (nn.Sequential [2] [4])

Read it as three statements:

1. the result, eventually, is an nn.Sequential;

2. that sequence maps a length-two vector to a length-four vector;

3. producing it lives in nn.Builder, because initialization consumes deterministic seeds and the builder is what threads them.

Similarly, Trainer.new model : Trainer σ τ says that construction returns a trainer whose model input shape is σ and whose output shape is τ, and every later method call preserves those two indices. Those indices connect the dataset and prediction methods to the model's interface.

Elaboration Errors🔗

Elaboration messages can be noisy, because they expose generated names, metavariables, and qualified internal constants, as the model composition error above shows. A reliable reading order:

• find has type: that is the value you actually supplied;

• find but is expected to have type: that is the interface it has to satisfy;

• compare the two from the outside in, and stop at the first shape, element type, or namespace where they differ;

• ignore everything with a ✝ or a ?m. in it until the end, since those are generated names and unsolved holes rather than things you wrote.

For model composition, work left to right and compare each layer's output with the next layer's input. For tensors, compare the element type as well as the dimension list, since Tensor Float [2] and Tensor Rat [2] fail differently than two mismatched shapes do. For training, check that the dataset target shape agrees with the model output shape and with the objective.

When an expression is too large to read, cut it up. Give a subexpression a name with def, ask Lean for its type with #check, and compose the pieces one at a time. Hover information in the editor is the same service in a faster interface.

Further Reading🔗

• Functional Programming in Lean is an introduction to Lean as a programming language, including do notation and IO.

• The Lean language reference gives precise syntax and semantics when you need the exact rule rather than the idea.

• The language paper de Moura and Ullrich (2021) explains the design, including why the elaborator is extensible enough for nn.Sequential! and the tensor literal to exist.

• Mathematics in Lean is the standard introduction to writing proofs with Mathlib (Community, 2020).

• The reference-counting paper Ullrich and de Moura (2019) is the background for the ownership examples above.

Sources:

• NN/Tensor/Operations.lean for the elementwise, matrix, and indexing operations used here;

• NN/Tensor/Reductions.lean for sum, mean, and the mean squared error;

• NN/Spec/Core/Shape.lean for shapes, coordinates, and appendDim;

• Internal/Elab/TensorLiteral.lean for the tensor literal elaborator and its error messages;

• NN/API.lean for what import NN.API brings into scope.