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 malformed hard-mask shapes or flat lengths. .log rejects an input tensor containing an entry <= 0 (or NaN), because the spec logarithm is undefined there. A custom computation can also reject a missing operand, an out-of-bounds read, or an output count beyond its index range. A pointwise operator fails where its meaning at the scalar type rejects an entry, such as a zero divisor over the reals, and where the scalar type has no meaning for it, such as erf in binary arithmetic. The remaining operations are 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. This equates evaluation with and without the final normalization on accepted graphs; it does not prove that shape checking alone guarantees successful evaluation.

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

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 valid rank-r target permutation.

This is used to implement .permute by repeatedly applying swapAdjacentAtDepth, which is already available in the spec tensor library. Callers must validate the permutation first, as permuteSomeTensor does with Shape.permute?. The errors below catch missing or out-of-range entries encountered during lowering; duplicate axes and extra entries are not rejected.

def NN.IR.Graph.permuteSomeTensor {α : Type} [TorchLean.Storage α] (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 α] (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 shape-erased tensors of equal shape agrees with the shaped formula.

def NN.IR.Graph.castDimScalar {α : Type} [TorchLean.Storage α] {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.matmulFlat {α : Type} [TorchLean.Storage α] [Context α] (dims : OpContracts.MatmulDims) {leftDim rightDim : ℕ} (left : TorchLean.Tensor α [leftDim]) (right : TorchLean.Tensor α [rightDim]) : TorchLean.Tensor α [dims.outShape.size]

Row-major matrix contraction using the shared broadcast and vector-promotion layout.

def NN.IR.Graph.matmulFlatFast {α : Type} [TorchLean.Storage α] [Context α] (dims : OpContracts.MatmulDims) {leftDim rightDim : ℕ} (left : TorchLean.Tensor α [leftDim]) (right : TorchLean.Tensor α [rightDim]) : TorchLean.Tensor α [dims.outShape.size]

matmulFlat with the batch offsets computed once per output entry and the contraction run by Nat.fold. The products are added in the same order, so the results are bit-identical.

@[csimp]

theorem NN.IR.Graph.matmulFlat_eq_matmulFlatFast : @matmulFlat = @matmulFlatFast

Compiled code runs matmulFlatFast in place of matmulFlat.

def NN.IR.Graph.matmulWithDims {α : Type} [TorchLean.Storage α] [Context α] (dims : OpContracts.MatmulDims) (left : TorchLean.Tensor α dims.leftShape) (right : TorchLean.Tensor α dims.rightShape) : TorchLean.Tensor α dims.outShape

Evaluate a checked matmul layout, retaining the typed matrix kernel for equal batch shapes.

def NN.IR.Graph.scaledMatmulWithDims {α : Type} [TorchLean.Storage α] [Context α] (dims : OpContracts.MatmulDims) (left : TorchLean.Tensor α dims.leftShape) (right : TorchLean.Tensor α dims.rightShape) (scale : α) : TorchLean.Tensor α dims.outShape

Evaluate a scaled contraction while retaining the original operands for coefficient replay.

def NN.IR.Graph.evalConst {α : Type} [TorchLean.Storage α] (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.layerNormMatrixValue {α : Type} [TorchLean.Storage α] [Context α] (seqLen embedDim : ℕ) (x : TorchLean.Tensor α [seqLen, embedDim]) (gamma beta : TorchLean.Tensor α [embedDim]) (epsilon : α) (hEmb : embedDim > 0) : TorchLean.Tensor α [seqLen, embedDim]

Normalize every present row; an empty batch has no coordinates to change.

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 with a nonempty normalized dimension and explicit affine parameters.

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 unit scale, zero bias, and default epsilon.

structure NN.IR.Graph.LayerNormAffine (α : Type) [TorchLean.Storage α] (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 α] (i : ℕ) (rest : Spec.Shape) (value : Spec.SomeTensor α) (permuted : Bool := false) : 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. permuted selects the diagnostic for parents moved from another axis; it does not change the shape check or tensor transport.

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

Fold leading-axis concat over dynamic values that already share the same tail shape. The permuted flag is passed only to the decoder's diagnostic.

def NN.IR.Graph.evalConcat {α : Type} [TorchLean.Storage α] (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. Tensor.concatAxisSpec takes any leading prefix, but this evaluator only uses its axis-0 form (leading := .scalar). For another axis it moves that axis to the front, folds Tensor.concatAxisSpec .scalar over the permuted parents, and moves the result back. The IR execution and CROWN concat proofs are stated against this form.

def NN.IR.Graph.normalizeNodeOutput {α : Type} [TorchLean.Storage α] (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 α] (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 α] (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 }

A tensor at a literal node's declared shape passes normalization unchanged.

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.evalTransform {α : Type} [TorchLean.Storage α] (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : Node) : Except String (Spec.SomeTensor α)

Evaluate tensor layout operations without requiring scalar arithmetic.

The same implementation serves homogeneous graphs and typed interpreters, including tensors of Boolean, integer, symbolic, or user-defined elements. Numerical operators are rejected here.

def NN.IR.Graph.evalPointwise {α : Type} [TorchLean.Storage α] (vals : Array (Spec.SomeTensor α)) (i : ℕ) (n : Node) (op : Spec.PointwiseOp) (meaning : Option (Spec.PointwiseFn α op.arity)) : Except String (Spec.SomeTensor α)

Evaluate a pointwise operator from its meaning at the scalar type.

Every parent is read at the node's declared shape, and the meaning is applied to corresponding entries in row-major order. A missing meaning, a wrong parent count, or a meaning that rejects an entry is an error; the first rejected entry stops evaluation.

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.

Recursion decreases the number of unevaluated nodes.

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, or a failed custom-computation read).

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.