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 conv2d are supplied by an explicit Payload.

The evaluator is total on well-formed, well-shaped graphs and returns Except String on malformed graphs or missing payloads.

Softmax and layer norm:

• softmax axis is interpreted as softmax along the given axis, but the spec primitive we have is last-axis softmax (Activation.softmax_spec). We therefore interpret non-last-axis softmax by permuting the requested axis to the last position, applying last-axis softmax, then permuting back. 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/

Dynamic (shape-tagged) values

During evaluation we keep values in a dependent pair Σ s, Tensor α s so we can store a heterogenous table of intermediate tensors while still recovering precise shapes when we need them.

@[reducible, inline]

abbrev NN.IR.DVal (α : Type) [Context α] : Type

Dynamic (shape-tagged) tensor value used by the IR evaluator.

This is a dependent pair Σ s, Tensor α s, which lets us store heterogeneously-shaped intermediate values in one table while still recovering exact shapes when needed.

def NN.IR.DVal.shape {α : Type} [Context α] (v : DVal α) : Spec.Shape

The shape tag carried by a dynamic value.

def NN.IR.DVal.tensor {α : Type} [Context α] (v : DVal α) : Spec.Tensor α v.shape

The underlying tensor, with its shape recovered from the dependent pair.

def NN.IR.DVal.mk {α : Type} [Context α] (s : Spec.Shape) (t : Spec.Tensor α s) : DVal α

Construct a dynamic value from a shape and a tensor of that shape.

Permutation lowering

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

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.permuteDVal {α : Type} [Context α] (v : DVal α) (perm : List ℕ) : Except String (DVal α)

Permute a dynamic tensor value 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 enforces “dynamic shape agrees with declared outShape” at each step.
• evalConst/evalLinear/evalConv2D fetch and apply external payloads keyed by node id.

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

Check a dynamic value has the expected shape and return it as a statically-typed tensor.

def NN.IR.Graph.meanDenom (s : Spec.Shape) : ℕ

Denominator for totalized mean reductions over a dynamic IR shape.

For nonempty shapes this is the real element count. For empty shapes, the mathematical mean is undefined; the IR is total, so it uses denominator 1 and the empty sum contributes 0.

def NN.IR.Graph.mseLossDVal {α : Type} [Context α] [DecidableEq Spec.Shape] (i : ℕ) (yVal tVal : DVal α) : Except String (DVal α)

Evaluate MSE loss on two dynamic values, checking that their runtime shapes agree.

@[simp]

theorem NN.IR.Graph.mseLossDVal_mk {α : Type} [Context α] [DecidableEq Spec.Shape] (i : ℕ) {s : Spec.Shape} (y t : Spec.Tensor α s) : mseLossDVal i (DVal.mk s y) (DVal.mk s t) = Except.ok (DVal.mk Spec.Shape.scalar (Spec.Tensor.scalar (((y.subSpec t).mulSpec (y.subSpec t)).sumSpec / Coe.coe (meanDenom s))))

def NN.IR.Graph.castDimScalar {α : Type} [Context α] {n n' : ℕ} (h : n = n') (t : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim n' Spec.Shape.scalar)

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

def NN.IR.Graph.evalConst {α : Type} [Context α] (payload : Payload α) (id : ℕ) (s : Spec.Shape) : Except String (Spec.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} [Context α] [DecidableEq Spec.Shape] (payload : Payload α) (id : ℕ) (x : DVal α) (outShape : Spec.Shape) : Except String (DVal α)

Evaluate a linear node from the external payload.

We enforce:

• the input dynamic value has shape (inDim), and
• the node's declared outShape matches (outDim).

The actual math is the usual affine map: $y=Wx+b$.

def NN.IR.Graph.evalConv2D {α : Type} [Context α] [DecidableEq Spec.Shape] (payload : Payload α) (id : ℕ) (x : DVal α) : Except String (DVal α)

Evaluate a conv2d node from the external payload.

The output shape is computed with the standard (no dilation) formula $$ \mathrm{out}=\left\lfloor \frac{\mathrm{in}+2\,\mathrm{pad}-k}{\mathrm{stride}} \right\rfloor+1 $$ for each spatial dimension.

def NN.IR.Graph.batchNorm2dEvalTensor {α : Type} [Context α] (cfg : BatchNorm2DNchwEvalParams α) {n h w : ℕ} (x : Spec.Tensor α (Spec.Shape.dim n (Spec.Shape.dim cfg.c (Spec.Shape.dim h (Spec.Shape.dim w Spec.Shape.scalar))))) : Spec.Tensor α (Spec.Shape.dim n (Spec.Shape.dim cfg.c (Spec.Shape.dim h (Spec.Shape.dim w Spec.Shape.scalar))))

Apply fixed-statistics BatchNorm2d to a batched channel-first tensor.

The input shape records the batch, channel, height, and width axes. Naming the shared tensor operation independently of that layout notation keeps compiler and evaluator code concise while the type retains the exact contract.

def NN.IR.Graph.evalBatchNorm2DNchwEval {α : Type} [Context α] [DecidableEq Spec.Shape] (payload : Payload α) (id : ℕ) (x : DVal α) : Except String (DVal α)

Evaluate eval-mode BatchNorm2d over a batched channel-first tensor.

def NN.IR.Graph.layernormPure {α : Type} [Context α] (seqLen embedDim : ℕ) (x : Spec.Tensor α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar))) : Except String (Spec.Tensor α (Spec.Shape.dim seqLen (Spec.Shape.dim embedDim Spec.Shape.scalar)))

Deterministic LayerNorm used by the IR evaluator ($\gamma=1$, $\beta=0$).

def NN.IR.Graph.expectLeadingAxisInput {α : Type} [Context α] (i : ℕ) (rest : Spec.Shape) (value : DVal α) : Except String ((size : ℕ) × Spec.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 list-indexed concat evaluation and its proofs.

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

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

def NN.IR.Graph.evalConcat {α : Type} [Context α] [DecidableEq Spec.Shape] (i : ℕ) (n : Node) (axis : ℕ) (parents : List (DVal α)) : Except String (DVal α)

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.concatLeadingAxisSpec over the permuted parents, and moving the result back.

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

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

@[simp]

theorem NN.IR.Graph.normalizeNodeOutput_declared {α : Type} [Context α] [DecidableEq Spec.Shape] (i : ℕ) (n : Node) (t : Spec.Tensor α n.outShape) : normalizeNodeOutput i n (DVal.mk n.outShape t) = Except.ok (DVal.mk n.outShape t)

@[simp]

theorem NN.IR.Graph.normalizeNodeOutput_nodeShape {α : Type} [Context α] [DecidableEq Spec.Shape] (i id : ℕ) (parents : List ℕ) (kind : OpKind) (s : Spec.Shape) (t : Spec.Tensor α s) : normalizeNodeOutput i { id := id, parents := parents, kind := kind, outShape := s } ⟨s, t⟩ = Except.ok ⟨s, t⟩

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

Evaluate a known node from its already computed parent values.

Keeping operator dispatch separate from graph lookup lets local correctness proofs reduce only the selected OpKind branch. The caller remains responsible for the graph's topological invariant.

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

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} [Context α] [DecidableEq Spec.Shape] (g : Graph) (payload : Payload α) (input : DVal α) (i : ℕ) (vals : Array (DVal α)) : Except String (Array (DVal α))

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} [Context α] [DecidableEq Spec.Shape] (g : Graph) (payload : Payload α) (input : DVal α) : Except String (Array (DVal α))

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 compiler-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 is total in the sense that it always returns either:

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

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} [Context α] [DecidableEq Spec.Shape] (g : Graph) (payload : Payload α) (input : DVal α) (outputId : ℕ) : Except String (DVal α)

Evaluate the graph and return the value at outputId.