Semantics

Denotational semantics for NN.IR.Graph.

This file defines an evaluator for the current IR fragment:

• it evaluates nodes in SSA/topological order,
• each node applies the corresponding spec-layer tensor operation to its parents,
• parameter payloads for const, linear, and conv are supplied by an explicit Payload.

Partiality

The evaluator returns Except String. It fails on malformed graphs (Graph.checkWellFormed), on nodes whose parents or declared outShape violate the shape rules, on missing or mismatched payloads (const flat length, linear dimensions, conv geometry, batchNormEval channels, layernorm normalized suffix), and on exactly one data-dependent condition: .log rejects an input tensor containing an entry <= 0 (or NaN), because the spec logarithm is undefined there. Every other operation is total on well-shaped inputs; reduceMean and layernorm divide by extents that the shape rules already require to be positive, and mseLoss divides by the totalized TorchLean.Tensor.meanDenominator rather than by the raw element count.

denoteAll runs the structural check but not Graph.checkShapes; each node instead checks its parents' shapes locally and normalizeNodeOutput compares the computed shape with the declared one. NN.IR.ShapeSoundness proves that on a graph accepted by checkShapes this final comparison never fails, so the two validation routes agree.

Softmax and layer norm:

• softmax axis normalizes independently along the zero-based tensor dimension named by axis. The canonical specification handles every in-bounds dimension and preserves the tensor shape. This matches the meaning of torch.softmax(x, dim=axis) in PyTorch.
• layernorm axis matches PyTorch's F.layer_norm(x, normalized_shape=x.shape[axis:]) convention: axis selects the start of the normalized suffix. We implement this by reshaping the tensor into a 2D view (seqLen, embedDim), applying the spec 2D LayerNorm (Spec.layerNorm), then reshaping back.

How this relates to PyTorch:

• Graph.nodes is analogous to a topologically-sorted IR like FX/TorchScript.
• Payload is analogous to “parameters / buffers / constants” that live outside the pure graph structure.
• The evaluator is a pure, denotational model of running the graph. It is designed for clarity and for connecting to proofs and verification passes (not for performance).

References / related systems:

• PyTorch FX: https://pytorch.org/docs/stable/fx.html
• TorchScript: https://pytorch.org/docs/stable/jit.html
• ONNX (graph + initializers): https://onnx.ai/

Except helpers

The evaluator returns Except String, so almost every proof about it has to unfold a throw at some point. That one rfl fact lives here, at the definition of the semantics, instead of being restated by each consumer: NN.IR.ShapeSoundness and the runtime correctness proofs in NN.Runtime.Autograd.IRExec both simp with it.

theorem NN.IR.throw_eq_error {β : Type} (msg : String) : throw msg = Except.error msg

throw msg in the Except String monad is .error msg.

Permutation lowering

def NN.IR.Graph.swapDepthsForPerm (perm : Array ℕ) (r : ℕ) : Except String (Array ℕ)

Compute a sequence of adjacent swaps that realizes a target permutation.

This is used to implement .permute by repeatedly applying swapAdjacentAtDepth, which is already available in the spec tensor library. If the permutation is ill-formed, this returns an error explaining what went wrong.

def NN.IR.Graph.permuteSomeTensor {α : Type} [TorchLean.Storage α] [Context α] (v : Spec.SomeTensor α) (perm : Array ℕ) : Except String (Spec.SomeTensor α)

Permute a shape-tagged tensor according to perm.

This checks that perm is a valid permutation for the input shape (using Shape.permute?), then lowers it to a sequence of adjacent swaps and applies them to the tensor.

Evaluation helpers

The evaluator itself (evalAt / denoteAll) is a fold over nodes. These helpers keep the fold readable:

• expectShape checks a value's stored shape against the node's declared outShape.
• evalConst, evalLinear, and evalConv fetch and apply external payloads keyed by node id.

def NN.IR.Graph.expectShape {α : Type} [TorchLean.Storage α] [Context α] (expected : Spec.Shape) (v : Spec.SomeTensor α) : Except String (TorchLean.Tensor α expected)

Check that a shape-erased tensor has the expected shape and recover its statically typed tensor.

def NN.IR.Graph.mseLossSomeTensor {α : Type} [TorchLean.Storage α] [Context α] (i : ℕ) (yVal tVal : Spec.SomeTensor α) : Except String (Spec.SomeTensor α)

Evaluate MSE loss on two shape-erased tensors after checking that their stored shapes agree.

@[simp]

theorem NN.IR.Graph.mseLossSomeTensor_mk {α : Type} [TorchLean.Storage α] [Context α] (i : ℕ) {s : Spec.Shape} (y t : TorchLean.Tensor α s) : mseLossSomeTensor i { shape := s, tensor := y } { shape := s, tensor := t } = Except.ok { shape := Spec.Shape.scalar, tensor := TorchLean.Tensor.scalar (((y.subSpec t).mulSpec (y.subSpec t)).sumSpec / ↑(TorchLean.Tensor.meanDenominator s)) }

MSE on two shape-erased tensors of the same shape unfolds to the shaped formula.

The IR hides tensor shapes, so the loss first has to re-discover that its two operands agree before it can subtract them. This equation says nothing else happens on the way.

def NN.IR.Graph.castDimScalar {α : Type} [TorchLean.Storage α] [Context α] {n n' : ℕ} (h : n = n') (t : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n']

Transport a Tensor α (dim n scalar) across an equality n = n' (helper for payload casts).

def NN.IR.Graph.linearLeading {α : Type} [TorchLean.Storage α] [Context α] (leading : Spec.Shape) {inDim outDim : ℕ} (weights : TorchLean.Tensor α [outDim, inDim]) (bias : TorchLean.Tensor α [outDim]) (input : TorchLean.Tensor α (leading.concat [inDim])) : TorchLean.Tensor α (leading.concat [outDim])

Apply one affine map independently at every index of an arbitrary leading shape.

def NN.IR.Graph.matmulLeading {α : Type} [TorchLean.Storage α] [Context α] (leading : Spec.Shape) {m n p : ℕ} (left : TorchLean.Tensor α (leading.concat [m, n])) (right : TorchLean.Tensor α (leading.concat [n, p])) : TorchLean.Tensor α (leading.concat [m, p])

Multiply matrices independently at every index of a shared leading shape.

def NN.IR.Graph.evalConst {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (id : ℕ) (s : Spec.Shape) : Except String (TorchLean.Tensor α s)

Evaluate a const node from the external payload.

Constants are stored “flat” (1D) for convenience, so we check the flattened length matches Spec.Shape.size s and then unflatten to the requested shape.

def NN.IR.Graph.evalLinear {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (id : ℕ) (x : Spec.SomeTensor α) (outShape : Spec.Shape) : Except String (Spec.SomeTensor α)

Evaluate a linear node from the external payload.

The final input axis must match inDim; every leading axis is preserved while the same affine map $y=Wx+b$ is applied independently.

def NN.IR.Graph.evalMaxPool {α : Type} [TorchLean.Storage α] [Context α] (config : WindowConfig) (x : Spec.SomeTensor α) : Except String (Spec.SomeTensor α)

Evaluate arbitrary-rank max pooling over the spatial suffix of a shape-erased tensor.

def NN.IR.Graph.evalAvgPool {α : Type} [TorchLean.Storage α] [Context α] (config : WindowConfig) (x : Spec.SomeTensor α) : Except String (Spec.SomeTensor α)

Evaluate arbitrary-rank average pooling over the spatial suffix of a shape-erased tensor.

def NN.IR.Graph.evalConv {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (id : ℕ) (config : ConvConfig) (x : Spec.SomeTensor α) : Except String (Spec.SomeTensor α)

Evaluate an arbitrary-rank convolution independently over every leading index.

def NN.IR.Graph.evalBatchNorm {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (id channelAxis channels : ℕ) (x : Spec.SomeTensor α) : Except String (Spec.SomeTensor α)

Evaluate fixed-statistics BatchNorm along an arbitrary channel axis.

def NN.IR.Graph.layerNormMatrix {α : Type} [TorchLean.Storage α] [Context α] (seqLen embedDim : ℕ) (x : TorchLean.Tensor α [seqLen, embedDim]) (gamma beta : TorchLean.Tensor α [embedDim]) (epsilon : α) : Except String (TorchLean.Tensor α [seqLen, embedDim])

Layer normalization of a matrix with explicit scale, bias, and epsilon.

def NN.IR.Graph.layerNormWithoutAffine {α : Type} [TorchLean.Storage α] [Context α] (seqLen embedDim : ℕ) (x : TorchLean.Tensor α [seqLen, embedDim]) : Except String (TorchLean.Tensor α [seqLen, embedDim])

Layer normalization with the historical unit affine transform and default epsilon.

structure NN.IR.Graph.LayerNormAffine (α : Type) [TorchLean.Storage α] [Context α] (embedDim : ℕ) : Type

Affine data for a LayerNorm matrix view after validating its normalized suffix.

• gamma : TorchLean.Tensor α [embedDim]
• beta : TorchLean.Tensor α [embedDim]
• epsilon : α

def NN.IR.Graph.resolveLayerNormAffine {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (id axis : ℕ) (outShape : Spec.Shape) (embedDim : ℕ) : Except String (LayerNormAffine α embedDim)

Resolve optional LayerNorm payload data into the vector shape consumed by the matrix semantics.

An absent payload means the standard unit scale, zero bias, and default normalization epsilon. Learned parameters are accepted only when their declared suffix is exactly the suffix normalized by the IR node.

def NN.IR.Graph.expectLeadingAxisInput {α : Type} [TorchLean.Storage α] [Context α] (i : ℕ) (rest : Spec.Shape) (value : Spec.SomeTensor α) : Except String ((size : ℕ) × TorchLean.Tensor α (Spec.Shape.dim size rest))

Decode a dynamic concat parent as a tensor with an existential leading dimension and the requested tail shape. This is the checked boundary shared by concat evaluation and its proofs.

def NN.IR.Graph.evalConcatLeadingAxisFold {α : Type} [TorchLean.Storage α] [Context α] (i nOut : ℕ) (rest : Spec.Shape) (parents : Array (Spec.SomeTensor α)) : Except String (Spec.SomeTensor α)

Fold leading-axis concat over dynamic values that already share the same tail shape.

def NN.IR.Graph.evalConcat {α : Type} [TorchLean.Storage α] [Context α] (i : ℕ) (n : Node) (axis : ℕ) (parents : Array (Spec.SomeTensor α)) : Except String (Spec.SomeTensor α)

Evaluate a concat node from already evaluated parent values.

The IR concat operation accepts any valid axis. The tensor primitive concatenates along axis 0, so the evaluator implements the generic case by moving the requested axis to the front, folding Tensor.concatAxisSpec .scalar over the permuted parents, and moving the result back.

def NN.IR.Graph.normalizeNodeOutput {α : Type} [TorchLean.Storage α] [Context α] (i : ℕ) (n : Node) (v : Spec.SomeTensor α) : Except String (Spec.SomeTensor α)

Normalize a node result to the node's declared shape, rejecting inconsistent implementations.

@[simp]

theorem NN.IR.Graph.normalizeNodeOutput_declared {α : Type} [TorchLean.Storage α] [Context α] (i : ℕ) (n : Node) (t : TorchLean.Tensor α n.outShape) : normalizeNodeOutput i n { shape := n.outShape, tensor := t } = Except.ok { shape := n.outShape, tensor := t }

A result whose shape matches the node declaration passes normalization unchanged.

@[simp]

theorem NN.IR.Graph.normalizeNodeOutput_nodeShape {α : Type} [TorchLean.Storage α] [Context α] (i id : ℕ) (parents : Array ℕ) (kind : OpKind) (s : Spec.Shape) (t : TorchLean.Tensor α s) : normalizeNodeOutput i { id := id, parents := parents, kind := kind, outShape := s } { shape := s, tensor := t } = Except.ok { shape := s, tensor := t }

The same statement with the node written out as a literal record.

Both spellings show up in proofs depending on whether the node came from a builder or was written inline, and simp will not see through the record projection on its own.

def NN.IR.Graph.getParentValue {α : Type} [TorchLean.Storage α] (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : Node) (pid : ℕ) : Except String (Spec.SomeTensor α)

Read an already-evaluated parent, reporting the node and available prefix on failure.

def NN.IR.Graph.unaryParentId (i : ℕ) (n : Node) : Except String ℕ

Decode the sole parent of a unary node.

def NN.IR.Graph.binaryParentIds (i : ℕ) (n : Node) : Except String (ℕ × ℕ)

Decode the two parents of a binary node.

def NN.IR.Graph.evalNodeRaw {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : Node) : Except String (Spec.SomeTensor α)

Evaluate a known node from its already computed parent values, without the final declared-shape normalization.

This is the operator dispatch of the evaluator: each OpKind branch validates its parents, applies the spec-layer operation, and returns a shape-erased value whose shape is computed by the shared shape rules (NN.IR.OpContracts). evalNode wraps it with normalizeNodeOutput. NN.IR.ShapeSoundness proves that the raw value already has the shape inferred by Infer.nodeOutShape. The definition is a simp lemma so proofs about evalNode reduce the selected branch exactly as before the split.

def NN.IR.Graph.evalNode {α : Type} [TorchLean.Storage α] [Context α] (payload : Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : Node) : Except String (Spec.SomeTensor α)

Evaluate a known node from its already computed parent values.

Keeping operator dispatch (evalNodeRaw) separate from graph lookup lets local correctness proofs reduce only the selected OpKind branch. The caller remains responsible for the graph's topological invariant. The result is normalized to the node's declared outShape; on graphs accepted by Graph.checkShapes that normalization is provably the identity.

def NN.IR.Graph.evalAt {α : Type} [TorchLean.Storage α] [Context α] (g : Graph) (payload : Payload α) (input : Spec.SomeTensor α) (vals : Array (Spec.SomeTensor α)) (i : ℕ) : Except String (Spec.SomeTensor α)

Evaluate node i after checking the graph's id discipline and retrieving the corresponding node.

denoteAll checks the full graph structure before repeatedly calling this one-step evaluator.

@[irreducible]

def NN.IR.Graph.denoteAllFrom {α : Type} [TorchLean.Storage α] [Context α] (g : Graph) (payload : Payload α) (input : Spec.SomeTensor α) (i : ℕ) (vals : Array (Spec.SomeTensor α)) : Except String (Array (Spec.SomeTensor α))

Evaluate nodes i, i+1, ... given already computed prefix values vals.

This is written as a structurally recursive function so it is easy to reason about in proofs (evaluation is “a simple loop over node ids”).

def NN.IR.Graph.denoteAll {α : Type} [TorchLean.Storage α] [Context α] (g : Graph) (payload : Payload α) (input : Spec.SomeTensor α) : Except String (Array (Spec.SomeTensor α))

Evaluate a graph to a table of node values.

This returns an array vals of length g.size where vals[i] is the value of node i.

We do a structural well-formedness check once up front (ids/arity/topology). For lowering-produced graphs, the boolean Graph.wellFormed check is a fast path; if it fails we fall back to the exception-producing Graph.checkWellFormed so callers get a readable error message.

The evaluator always returns either:

• .ok vals (all nodes evaluated successfully), or
• .error msg describing the first failure (malformed IR, missing payload, a local shape error, or a .log of a nonpositive entry).

Graph.checkShapes is not run here: the per-node checks reject the same ill-shaped graphs, and the existing lowering-correctness proofs unfold denoteAll with only the structural check in place. NN.IR.ShapeSoundness.denoteAllRaw_eq_denoteAll shows that on a checkShapes-accepted graph the per-node declared-shape normalization is redundant.

Scoped notation

def IR.«term_⟦_,_⟧» : Lean.TrailingParserDescr

Scoped notation for evaluating a graph to all node values.

Use with:

open scoped IR
g⟦payload, input⟧

def IR.«term_[[_,_]]» : Lean.TrailingParserDescr

ASCII alternative to g⟦payload, input⟧.

def NN.IR.Graph.denote {α : Type} [TorchLean.Storage α] [Context α] (g : Graph) (payload : Payload α) (input : Spec.SomeTensor α) (outputId : ℕ) : Except String (Spec.SomeTensor α)

Evaluate the graph and return the value at outputId.