8.4. Backends and Training
Backend selection is a planning step over declared data. A capsule names the operation, provider, device, trust level, reduction policy, and the evidence behind its four contract claims. The contract check turns a plan audit into an accept or reject decision under an assurance policy, and eager execution binds the accepted capsule to a handler with the same identity. Nothing in this group proves a kernel correct; the nodes record what is checked and what is assumed.
For example, selecting CUDA matmul involves a shape obligation, a row-major layout obligation, and numerical obligations for forward and backward. Acceptance explains why this provider is allowed to run. The native result still depends on the implementation satisfying its recorded contract. The training and import nodes below describe how values reach that boundary.
A kernel capsule records an operation, provider, device, trust level (checked or
trustedExternal), forward and VJP support, four contract descriptors (shape, layout, forward
value, VJP), and a numerical policy. Each descriptor pairs a
structured claim with its evidence: a runtime guard, a test suite, a named trusted boundary, or
notApplicable. The capsule states the contract expected from an implementation; it does not
prove that implementation.
Lean code for Definition8.4.1●1 definition
Associated Lean declarations
-
NN.Backend.KernelCapsule[complete]
-
NN.Backend.KernelCapsule[complete]
-
structuredefined in NN/Backend/Capsule.leancomplete
structure NN.Backend.KernelCapsule : Type
structure NN.Backend.KernelCapsule : Type
A contract-carrying fast kernel or reference implementation.
Fields
name : String
Name used in selection reports and runtime errors.
op : NN.Backend.BackendOp
Backend operation implemented by this capsule.
provider : NN.Backend.Provider
Provider responsible for the implementation.
device : NN.Backend.Device
Device on which the implementation runs.
trustLevel : NN.Backend.TrustLevel
Assurance level the planner must accept before selection.
supportsForward : Bool
Whether the capsule supplies forward execution.
vjpMode : NN.Backend.VJPMode
Form of reverse-mode support supplied by the capsule.
shapeContract : NN.Backend.ContractDescriptor
Shape-safety claim and its evidence.
layoutContract : NN.Backend.ContractDescriptor
Tensor-layout claim and its evidence.
valueContract : NN.Backend.ContractDescriptor
Forward-value refinement claim and its evidence.
vjpContract : NN.Backend.ContractDescriptor
Reverse-mode refinement claim and its evidence.
numericalPolicy : NN.Backend.NumericalPolicy
Floating-point behavior consumed by numerical certificates.
The numerical policy of a capsule is its reduction order: fixedLeft for the left fold used by the
canonical tensor semantics, implementationDefined for native and library accumulations, or
notApplicable. Numerical certificates read this field so that a fixed-left range trace is not
reused for an implementation-defined schedule. Rounding mode, subnormal handling, and multiply-add
contraction are not recorded.
The distinction matters even for the same matrix dimensions: changing the accumulation tree changes the sequence of rounded additions. A shape proof cannot substitute for the missing numerical-policy match.
Lean code for Definition8.4.2●1 definition
Associated Lean declarations
-
NN.Backend.NumericalPolicy[complete]
-
NN.Backend.NumericalPolicy[complete]
-
structuredefined in NN/Backend/Capsule.leancomplete
structure NN.Backend.NumericalPolicy : Type
structure NN.Backend.NumericalPolicy : Type
Floating-point choices attached to one kernel capsule that numerical certificates consume.
Fields
reduction : NN.Backend.ReductionPolicy
The order a reduction may use, which fixes whether summation is reproducible.
The contract check reads every obligation of a plan audit and rejects those whose evidence the
assurance policy does not accept. The checked policy admits runtime guards, test suites, and
notApplicable; only the external policy admits a trusted boundary. The result is a
ContractCheck: either accepted or rejected with the failing obligation reports.
Lean code for Definition8.4.3●1 definition
Associated Lean declarations
-
defdefined in NN/Backend/ContractCheck.leancomplete
def NN.Backend.KernelPlanAudit.checkContracts (policy : NN.Backend.AssurancePolicy) (a : NN.Backend.KernelPlanAudit) : NN.Backend.ContractCheck
def NN.Backend.KernelPlanAudit.checkContracts (policy : NN.Backend.AssurancePolicy) (a : NN.Backend.KernelPlanAudit) : NN.Backend.ContractCheck
Check every obligation of a plan audit against an assurance policy.
An accepted kernel is a planned operation whose
capsule passed the complete kernel-policy gate, including the
evidence check. Its proof field states
PlannedKernel.acceptable policy = true: operation identity, forward support, contract alignment,
trust, provider, device, and VJP mode must also agree. AcceptedGraphKernelPlan carries the
corresponding check for groups derived from its stored graph plan.
Lean code for Definition8.4.4●1 definition
Associated Lean declarations
-
NN.Backend.AcceptedKernel[complete]
-
NN.Backend.AcceptedKernel[complete]
-
structuredefined in NN/Backend/ContractCheck.leancomplete
structure NN.Backend.AcceptedKernel : Type
structure NN.Backend.AcceptedKernel : Type
One planned operation whose capsule has passed the complete policy and evidence gate.
Fields
op : NN.Backend.BackendOp
The operation the kernel implements.
capsule : NN.Backend.KernelCapsule
The capsule that carries the kernel's contract claims and evidence.
policy : NN.Backend.KernelPolicy
The policy the capsule was checked against.
accepted : NN.Backend.PlannedKernel.acceptable self.policy { op := self.op, capsule := self.capsule } = true
Evidence that the complete policy and evidence gate accepted this kernel.
Cuda.Buffer is an opaque handle to a contiguous float32 buffer. A CUDA build stores device
memory behind the handle; the default stub keeps parity storage on the host. Lean code cannot
inspect either representation directly.
A typed shape supplies a logical element count; runtime validation compares it with the handle's reported length. This checks an observable interface condition without exposing the storage as a Lean array or deriving its contents from the type.
Lean code for Definition8.4.5●1 definition
Associated Lean declarations
-
Runtime.Autograd.Cuda.Buffer[complete]
-
Runtime.Autograd.Cuda.Buffer[complete]
-
defdefined in NN/Runtime/Autograd/Engine/Cuda/Trusted.leancomplete
def Runtime.Autograd.Cuda.Buffer : Type
def Runtime.Autograd.Cuda.Buffer : Type
Runtime representation used for native CUDA buffer handles. The `NonemptyType` wrapper is Lean's standard representation for external resources: it gives extern declarations a nonempty result type while preserving reference-counting information in compiled code. The underlying value is still created only by the native buffer constructors.
A backend profile stores a name, a kernel policy (device, provider preference, assurance policy, and VJP mode), the devices and providers declared available, and the capsule modules that form its planning registry. Capsule modules are validated for duplicate names when a graph is planned.
Lean code for Definition8.4.6●1 definition
Associated Lean declarations
-
NN.Backend.BackendProfile[complete]
-
NN.Backend.BackendProfile[complete]
-
structuredefined in NN/Backend/Profile.leancomplete
structure NN.Backend.BackendProfile : Type
structure NN.Backend.BackendProfile : Type
A named kernel-selection profile for one device.
Fields
name : String
Human-readable profile name used in diagnostics and reports.
policy : NN.Backend.KernelPolicy
Device, provider preference, assurance policy, and VJP ownership.
availability : NN.Backend.Availability
Devices and providers the build declares available to planning.
capsuleModules : Array NN.Backend.Registry.CapsuleModule
Capsule modules used to construct and validate the profile's planning registry.
The maintained registry collects contract capsules contributed by the attention, native CUDA, and reference modules. It contains planning metadata, not executable handlers. Profiles add the separate LibTorch module when requested.
Lean code for Definition8.4.7●1 definition
Associated Lean declarations
-
NN.Backend.Registry.maintainedModules[complete]
-
NN.Backend.Registry.maintainedModules[complete]
-
defdefined in NN/Backend/Registry.leancomplete
def NN.Backend.Registry.maintainedModules : Array NN.Backend.Registry.CapsuleModule
def NN.Backend.Registry.maintainedModules : Array NN.Backend.Registry.CapsuleModule
Maintained operation/provider modules. A new architecture does not modify this list; only a new primitive implementation or provider does.
The checked CPU profile instantiates the profile record with the maintained capsule modules, CPU-only availability, and the checked assurance policy.
Lean code for Definition8.4.8●1 definition
Associated Lean declarations
-
NN.Backend.BackendProfile.checkedCpu[complete]
-
NN.Backend.BackendProfile.checkedCpu[complete]
-
defdefined in NN/Backend/Profile.leancomplete
def NN.Backend.BackendProfile.checkedCpu : NN.Backend.BackendProfile
def NN.Backend.BackendProfile.checkedCpu : NN.Backend.BackendProfile
Maintained portable CPU/reference profile with runtime guards and regression evidence.
Binding a selected capsule to a handler checks that their operation, provider, and device agree. The resulting executable kernel carries those identity equalities; binding does not strengthen the capsule's evidence.
A CUDA capsule paired with a CPU handler is rejected even when both declare matmul. Agreement on
those labels permits dispatch; it does not inspect the handler's IO body.
Lean code for Definition8.4.9●1 definition
Associated Lean declarations
-
NN.Backend.KernelCapsule.bind[complete]
-
NN.Backend.KernelCapsule.bind[complete]
-
defdefined in NN/Backend/Capsule.leancomplete
def NN.Backend.KernelCapsule.bind {β : Type} (c : NN.Backend.KernelCapsule) (handler : NN.Backend.KernelHandler β) : Except String (NN.Backend.ExecutableKernel β)
def NN.Backend.KernelCapsule.bind {β : Type} (c : NN.Backend.KernelCapsule) (handler : NN.Backend.KernelHandler β) : Except String (NN.Backend.ExecutableKernel β)
Pair a selected contract with the runtime handler that will execute it. The returned equalities prevent an executor for one operation or provider from being presented as another merely because both happen to share a Lean result type.
A backend profile validates its configured capsule modules, requires a structurally well-formed IR graph, and selects a capsule for each runtime-relevant node. It then groups those choices and runs the contract check under the profile's assurance policy, returning an accepted graph plan or the rejected obligations.
Lean code for Definition8.4.10●1 definition
Associated Lean declarations
-
NN.Backend.BackendProfile.acceptGraph[complete]
-
NN.Backend.BackendProfile.acceptGraph[complete]
-
defdefined in NN/Backend/Profile.leancomplete
def NN.Backend.BackendProfile.acceptGraph (p : NN.Backend.BackendProfile) (g : NN.IR.Graph) : Except String NN.Backend.GraphKernelPlanResult
def NN.Backend.BackendProfile.acceptGraph (p : NN.Backend.BackendProfile) (g : NN.IR.Graph) : Except String NN.Backend.GraphKernelPlanResult
Plan, group, and contract-check a graph under the profile.
For one eager operation, the session selects and caches an admitted capsule, finds a handler with the same operation, provider, and device, and runs it through the checked binding. This dispatcher does not append an autograd node; each operation implementation records its own forward value and VJP.
Caching a selection reuses the admitted provider choice for this session, whose profile stays fixed. It does not cache tensor results or cotangents: later calls still execute with their current inputs, and the operation remains responsible for recording dependencies on those inputs.
Lean code for Definition8.4.11●1 definition
Associated Lean declarations
-
defdefined in NN/Runtime/Autograd/Torch/Core/Session/Backend.leancomplete
def Runtime.Autograd.Torch.Internal.EagerSession.executeSelected {α β : Type} [TorchLean.Storage α] (s : Runtime.Autograd.Torch.Internal.EagerSession α) (op : NN.Backend.BackendOp) (handlers : Array (NN.Backend.KernelHandler β)) : IO β
def Runtime.Autograd.Torch.Internal.EagerSession.executeSelected {α β : Type} [TorchLean.Storage α] (s : Runtime.Autograd.Torch.Internal.EagerSession α) (op : NN.Backend.BackendOp) (handlers : Array (NN.Backend.KernelHandler β)) : IO β
Execute an operation through a handler that matches the selected capsule. Fail if the chosen provider is not linked into this build. `ExecutableKernel` certifies dispatch identity; the selected capsule states the numerical evidence.
The CUDA tape stores device buffers, parent ids, and local VJP
closures in evaluation order. requireValue, requireGrad, and backward accumulation check shape
tags and native buffer lengths. Dense backward returns one buffer per node; sparse backward retains
owned buffers only for selected node ids and requires the caller to release them.
Lean code for Definition8.4.12●1 definition
Associated Lean declarations
-
Runtime.Autograd.Cuda.Tape[complete]
-
Runtime.Autograd.Cuda.Tape[complete]
-
structuredefined in NN/Runtime/Autograd/Engine/Cuda/Tape.leancomplete
structure Runtime.Autograd.Cuda.Tape : Type
structure Runtime.Autograd.Cuda.Tape : Type
CUDA autograd tape: a grow-only array of nodes. Node ids are array indices.
Fields
nodes : Array Runtime.Autograd.Cuda.Node
Tape nodes in evaluation order (id = index).
Given the stated native bit-agreement hypothesis and a finite native result, decoded native scalar
addition equals ExecFloat.add. Both the hypothesis and the
finiteness side condition remain visible in the theorem type.
Lean code for Theorem8.4.13●1 theorem
Associated Lean declarations
-
theoremdefined in NN/Runtime/Autograd/Engine/Cuda/Float32Contract.leancomplete
theorem Runtime.Autograd.Cuda.Float32Contract.native_add_eq_ieee32_of_isFinite {native : Runtime.Autograd.Cuda.Float32Contract.NativePrimitiveBits} (h : Runtime.Autograd.Cuda.Float32Contract.NativePrimitiveAgreement native) (x y : Runtime.Autograd.Cuda.Float32Contract.RefScalar) (hfin : ExecFloat.Binary.isFinite (Runtime.Autograd.Cuda.Float32Contract.fromNativeBits (native.addBits x y)) = true) : Runtime.Autograd.Cuda.Float32Contract.fromNativeBits (native.addBits x y) = ExecFloat.add x y
theorem Runtime.Autograd.Cuda.Float32Contract.native_add_eq_ieee32_of_isFinite {native : Runtime.Autograd.Cuda.Float32Contract.NativePrimitiveBits} (h : Runtime.Autograd.Cuda.Float32Contract.NativePrimitiveAgreement native) (x y : Runtime.Autograd.Cuda.Float32Contract.RefScalar) (hfin : ExecFloat.Binary.isFinite (Runtime.Autograd.Cuda.Float32Contract.fromNativeBits (native.addBits x y)) = true) : Runtime.Autograd.Cuda.Float32Contract.fromNativeBits (native.addBits x y) = ExecFloat.add x y
Native addition is the reference value when its result is finite and the contract holds.
The proof rewrites the supplied native result bits to
the executable binary32 result, using finiteness to rule out the
NaN case the agreement hypothesis tolerates. It does not prove the external kernel implementation
from source.
The bit-agreement premise permits different NaN encodings. Finiteness removes that alternative, leaving exact bit equality for the scalar result. Applying a real-valued error bound then needs the separate theorem relating the executable binary32 operation to real arithmetic.
Objective wraps a scalar trainer, its runtime options, and the selected host/device tensor
conversion. The trainer owns a shape-indexed mutable parameter pack and runs its scalar-loss
program through eager execution or
typed graph execution. Gradients are returned to callers;
generic optimizer state is passed to and returned from update methods rather than stored here.
State shapes include persistent model buffers as well as trainable parameters. Differentiability flags determine which entries receive gradients, so possessing a state tensor does not by itself make that tensor an optimization variable.
Lean code for Definition8.4.14●1 definition
Associated Lean declarations
-
Runtime.Autograd.Model.Module.Objective[complete]
-
Runtime.Autograd.Model.Module.Objective[complete]
-
structuredefined in NN/Runtime/Autograd/Model/Module/Objective.leancomplete
structure Runtime.Autograd.Model.Module.Objective (α β : Type) [TorchLean.Storage α] [TorchLean.Storage β] [Context α] (stateShapes inputShapes : List Spec.Shape) (dataInputShapes : List Spec.Shape := []) : Type
structure Runtime.Autograd.Model.Module.Objective (α β : Type) [TorchLean.Storage α] [TorchLean.Storage β] [Context α] (stateShapes inputShapes : List Spec.Shape) (dataInputShapes : List Spec.Shape := []) : Type
Runtime state for a model together with a scalar objective. This is lower level than PyTorch's loss classes: it owns model state as well as the objective. It wraps `Torch.ScalarTrainer` and exposes objective evaluation, explicit gradients, and updates.
Fields
trainer : Runtime.Autograd.Torch.ScalarTrainer α β stateShapes inputShapes dataInputShapes
Trainer that owns trainable parameters and persistent buffers.
runtime : Runtime.Autograd.Torch.Config
Runtime configuration used to instantiate the module.
tensorTransfer : Runtime.Autograd.Torch.TensorTransfer α
Concrete host/device tensor conversion selected when the module was instantiated.
Session owns a supervised update action and step counter behind the public step, stepBatch,
and steps operations, together with evaluation-mode prediction and loss and the finish
operation that packages the live state as a trained result.
A batch update and an evaluation may use the same parameters with different mode-dependent behavior, such as dropout. The step counter tracks updates; it is not a count of every forward call made for prediction or loss inspection.
Lean code for Definition8.4.15●1 definition
Associated Lean declarations
-
TorchLean.Trainer.Session[complete]
-
TorchLean.Trainer.Session[complete]
-
structuredefined in NN/API/Trainer/Session.leancomplete
structure TorchLean.Trainer.Session {σ τ : Spec.Shape} (trainer : TorchLean.Trainer σ τ) : Type
structure TorchLean.Trainer.Session {σ τ : Spec.Shape} (trainer : TorchLean.Trainer σ τ) : Type
A trainer with its model instantiated and the optimizer loop owned by the caller. Obtain one with `trainer.open`. The session type is indexed by the trainer it runs, so the state layout `nn.stateShapes trainer.model` is available to `state` and `load`. Mode handling is implicit: updates run stateful layers in training mode, while `predict`, `loss`, and `eval` run them in evaluation mode.
Opening a trainer builds the stateful training loop around a scalar module. It selects the trainer's optimizer, applies an optional learning-rate schedule, and refreshes mode-dependent model buffers before each update.
Lean code for Definition8.4.16●1 definition
Associated Lean declarations
-
TorchLean.Trainer.open[complete]
-
TorchLean.Trainer.open[complete]
-
defdefined in NN/API/Trainer/Session.leancomplete
def TorchLean.Trainer.open {σ τ : Spec.Shape} (trainer : TorchLean.Trainer σ τ) (scheduler : Option TorchLean.Trainer.Scheduler.Config := none) : IO trainer.Session
def TorchLean.Trainer.open {σ τ : Spec.Shape} (trainer : TorchLean.Trainer σ τ) (scheduler : Option TorchLean.Trainer.Scheduler.Config := none) : IO trainer.Session
Instantiate the model under the trainer's runtime settings and hand the optimizer loop to the caller. The trainer's optimizer is bound immediately; `scheduler` optionally adjusts its learning rate by completed step. Optimizer and scheduler settings must remain valid after binary32 conversion and are checked before the model is instantiated. CUDA execution requires `.native` arithmetic, and `.complex` arithmetic is not supported by supervised training. Example: ```lean -- A session holds the instantiated model and the bound optimizer. `trainer.train` is exactly this -- call, a loop of `step`, and a `finish`, so opening a session is how you take that loop over. def stepOnce (trainer : TorchLean.Trainer [2] [1]) : IO Float := do let session ← trainer.open session.step { input := [1.0, 0.0], target := [1.0] } ```
NN.IR.OpTag identifies each IR operation constructor. Wire.opTag gives its fixed v1 kind
string, and Wire.parseOpTag? reads that string back into a tag.
Lean code for Definition8.4.17●1 definition
Associated Lean declarations
-
Interop.PyTorch.Wire.parseOpTag?[complete]
-
Interop.PyTorch.Wire.parseOpTag?[complete]
-
defdefined in NN/Runtime/PyTorch/Wire.leancomplete
def Interop.PyTorch.Wire.parseOpTag? (s : String) : Option NN.IR.OpTag
def Interop.PyTorch.Wire.parseOpTag? (s : String) : Option NN.IR.OpTag
Parse a v1 constructor spelling into its semantic identity.
Every operation tag parses back from its own wire string, so the
tag table is complete and collision free. parse_op_kind_tag states the same round trip starting
from an OpKind with payload.
The conclusion recovers the constructor tag, not all its attributes. Reconstructing an axis, stride, shape, or tensor payload needs the corresponding parser and validation beyond this tag identity.
Lean code for Theorem8.4.18●1 theorem
Associated Lean declarations
-
Interop.PyTorch.Wire.parse_op_tag[complete]
-
Interop.PyTorch.Wire.parse_op_tag[complete]
-
theoremdefined in NN/Runtime/PyTorch/Wire.leancomplete
theorem Interop.PyTorch.Wire.parse_op_tag (tag : NN.IR.OpTag) : Interop.PyTorch.Wire.parseOpTag? (Interop.PyTorch.Wire.opTag tag) = some tag
theorem Interop.PyTorch.Wire.parse_op_tag (tag : NN.IR.OpTag) : Interop.PyTorch.Wire.parseOpTag? (Interop.PyTorch.Wire.opTag tag) = some tag
Every v1 tag parses back to its identity: the codec is complete and collision free.
Case analysis over the tag constructors; each case is rfl.
The torch.export adapter parses TorchLean's captured graph schema, lowers its supported values to
the shared IR, runs shape validation, and checks that the named input
and output nodes exist.
Lean code for Definition8.4.19●1 definition
Associated Lean declarations
-
Import.PyTorch.TorchExport.parseGraph[complete]
-
Import.PyTorch.TorchExport.parseGraph[complete]
-
defdefined in NN/Runtime/PyTorch/Import/TorchExport.leancomplete
def Import.PyTorch.TorchExport.parseGraph (j : Lean.Json) : Except String Import.PyTorch.TorchExport.CapturedGraph
def Import.PyTorch.TorchExport.parseGraph (j : Lean.Json) : Except String Import.PyTorch.TorchExport.CapturedGraph
Parse and validate a captured PyTorch graph. Success means: - the JSON uses the TorchLean graph-artifact schema, - every op is in the supported TorchLean IR subset, - node ids are disciplined and topologically ordered, - arities are valid, and - declared output shapes match `NN.IR.Infer`.
Every graph returned successfully by the torch.export parser
satisfies TorchLean's executable shape check.
This permits downstream code to rely on the accepted graph's shape equations. Numerical agreement with the captured Python model still requires correct operation lowering and matching parameter payloads; two different activations can satisfy the same shape check.
Lean code for Theorem8.4.20●1 theorem
Associated Lean declarations
-
theoremdefined in NN/Runtime/PyTorch/Import/TorchExport.leancomplete
theorem Import.PyTorch.TorchExport.parseGraph_wellShaped {j : Lean.Json} {cg : Import.PyTorch.TorchExport.CapturedGraph} (h : Import.PyTorch.TorchExport.parseGraph j = Except.ok cg) : cg.graph.WellShaped
theorem Import.PyTorch.TorchExport.parseGraph_wellShaped {j : Lean.Json} {cg : Import.PyTorch.TorchExport.CapturedGraph} (h : Import.PyTorch.TorchExport.parseGraph j = Except.ok cg) : cg.graph.WellShaped
Guarantee exposed by the parser: a successfully parsed graph is well-shaped. This theorem is compact but important. It is the theorem downstream verification/export code can quote when it receives a graph artifact through this importer.
The proof unfolds the parser, rules out each rejected branch, and returns the successful shape check.