Operation Contracts

Shared operation contracts for NN.IR.Graph.

Several IR passes need to agree on the same small set of “shape contracts”:

• NN.IR.Infer: recompute output shapes from op parameters + parent shapes.
• NN.IR.Check: expose the documented Graph.checkShapes wrapper.
• NN.IR.Semantics: evaluate nodes and reject ill-shaped graphs with readable error messages.

The point of this file is to keep shape arithmetic out of individual passes. If an op has nontrivial shape behavior (concat, matmul, pooling, convolution, LayerNorm flattening, axis moves), define the contract here first and call it from inference/semantics instead of copying the formula.

Small shape utilities

These helpers are used by multiple IR passes, especially Infer and Semantics.

def NN.IR.ShapeUtil.flattenOutShape (s : Spec.Shape) : Spec.Shape

The output shape of flattening a tensor of shape s to a 1D vector.

def NN.IR.ShapeUtil.swapFirstTwoShape? : Spec.Shape → Option Spec.Shape

If s has rank at least two, return the shape obtained by swapping its first two axes.

Example: (a, b, rest) becomes (b, a, rest).

def NN.IR.ShapeUtil.transpose3dLastTwoShape? : Spec.Shape → Option Spec.Shape

If s has shape (a, b, c) (rank three with scalar base), return (a, c, b).

This is the common “transpose the last two axes” pattern for batched matrices.

Generic contract helpers

These functions live outside any particular pass (Infer/Check/Semantics) so they can be reused without introducing import cycles.

def NN.IR.OpContracts.checkAxisValid (axis : ℕ) (s : Spec.Shape) : Except String Unit

Check that an axis is in-bounds for a given shape.

def NN.IR.OpContracts.checkPositive (tag param : String) (n : ℕ) : Except String Unit

Check that a natural-number op parameter is nonzero.

def NN.IR.OpContracts.mkCanBroadcastTo? (s₁ s₂ : Spec.Shape) : Option (s₁.CanBroadcastTo s₂)

Reconstruct the proof object required by the typed tensor broadcast primitive.

IR nodes store dynamic shapes, so every pass that accepts .broadcastTo must rebuild this witness instead of trusting that the declared input and output shapes are compatible.

def NN.IR.OpContracts.layerNorm2DParams (axis : ℕ) (s : Spec.Shape) : Except String (ℕ × ℕ)

Compute the (seqLen, embedDim) pair used to interpret layernorm axis.

TorchLean’s IR stores LayerNorm as an axis : Nat instead of a full normalized_shape tuple. We interpret this in the same way the PyTorch exporter does:

normalized_shape = dims.drop axis

That is, we normalize over the suffix of dimensions starting at axis. To reuse the current spec primitive (Spec.layerNorm), we flatten the input shape s into a 2D view:

• seqLen = product of dimensions before axis (dims.take axis)
• embedDim = product of dimensions from axis onward (dims.drop axis)

Then we run 2D last-axis LayerNorm on a (seqLen × embedDim) tensor and reshape back.

def NN.IR.OpContracts.checkLastAxis (tag : String) (axis : ℕ) (s : Spec.Shape) : Except String Unit

Check that axis refers to the last axis of s.

This is a convenience predicate for passes/backends that restrict an op to last-axis behavior. For example, some verification bounds are implemented only for last-axis softmax/layernorm and use this check to fail fast with a readable error.

def NN.IR.OpContracts.inversePerm (perm : List ℕ) : Except String (List ℕ)

Compute the inverse of a permutation list.

If perm is a permutation of [0,1,...,r-1] (where r = perm.length), then the inverse inv satisfies $\mathrm{inv}[\mathrm{perm}[i]]=i$.

def NN.IR.OpContracts.inversePerm.setOnce (perm : List ℕ) (r : ℕ) (xs : List (Option ℕ)) (axis j val : ℕ) : Except String (List (Option ℕ))

def NN.IR.OpContracts.inversePerm.getOpt : List (Option ℕ) → ℕ → Option (Option ℕ)

def NN.IR.OpContracts.permMoveAxisToLast (axis : ℕ) (s : Spec.Shape) : Except String (List ℕ)

Permutation (0-based axes) that moves axis to the last position, preserving the relative order of the other axes.

Example: rank four with axis set to 1 yields [0,2,3,1].

def NN.IR.OpContracts.permMoveAxisToFront (axis : ℕ) (s : Spec.Shape) : Except String (List ℕ)

Permutation (0-based axes) that moves axis to the first position, preserving the relative order of the other axes.

Example: rank four with axis set to 2 yields [2,0,1,3].

def NN.IR.OpContracts.inferMatmulOutShape (a b : Spec.Shape) : Except String Spec.Shape

Infer the output shape for matmul from the two parent shapes.

Supported cases:

• 2D: (m×n) · (n×p) → (m×p)
• limited 3D “batched matmul”: (b×m×n) · (b×n×p) → (b×m×p)

def NN.IR.OpContracts.checkConcatRanks (expected : ℕ) : List Spec.Shape → Except String Unit

Check that every concat parent has the expected rank.

def NN.IR.OpContracts.inferConcatLeadingAxis (expectedTail : Spec.Shape) : ℕ → List Spec.Shape → Except String Spec.Shape

Infer leading-axis concat after the first input has fixed the common tail and initial dimension.

def NN.IR.OpContracts.inferConcatOutShape (axis : ℕ) (parents : List Spec.Shape) : Except String Spec.Shape

Infer the output shape for concat from the parent shapes.

All parents must:

• have the same rank,
• agree on every dimension except axis, and
• have axis in bounds.

The output shape matches the parents except at axis, where the dimension is the sum of the input dimensions.

PyTorch analogy: torch.cat(xs, dim=axis) for a list xs of tensors.

def NN.IR.OpContracts.inferConcatOutShape.go (axis : ℕ) (shs : List Spec.Shape) : Except String Spec.Shape

Sliding-window shape arithmetic

Convolution and pooling preserve a leading channel axis, but their admissible padding domains are not identical. The contracts below share validation and traversal while retaining the correct output formula for each operation family.

def NN.IR.OpContracts.slideOut (inLen k stride : ℕ) : ℕ

Output length for a 1D sliding-window op without padding: $\left\lfloor(\mathrm{in}-k)/\mathrm{stride}\right\rfloor+1$.

def NN.IR.OpContracts.slideOutPad (inLen k stride padding : ℕ) : ℕ

Output length for a 1D sliding-window op with symmetric padding: ⌊(in + 2*pad - k)/stride⌋ + 1.

def NN.IR.OpContracts.checkWindowFits (tag axis : String) (inLen k padding : ℕ) : Except String Unit

Reject sliding-window shapes where the kernel has no valid placement.

Lean Nat subtraction saturates at zero, so $\mathrm{in}+2\,\mathrm{pad}-k$ would otherwise turn an invalid window into a plausible one-element output.

def NN.IR.OpContracts.inferSlidingWindowDims (tag : String) (axisNames : List String) (inputs kernels strides paddings : List ℕ) : Except String (List ℕ)

Infer the output lengths of a channel-first sliding-window operation.

The four lists describe the input length, kernel width, stride, and symmetric padding on each spatial axis. Their lengths must agree. Invalid kernels, strides, and windows are rejected before Nat subtraction can hide the error by saturating at zero.

def NN.IR.OpContracts.inferSlidingWindowDims.go (tag : String) : List String → List ℕ → List ℕ → List ℕ → List ℕ → Except String (List ℕ)

def NN.IR.OpContracts.inferPoolingDims (tag : String) (axisNames : List String) (inputs kernels strides paddings : List ℕ) : Except String (List ℕ)

Infer pooling output lengths while enforcing the same basic window checks as graph validation.

Unlike convolution, pooling uses poolOutDim, which assigns an empty output to empty input axes and to padding outside the pooling domain.

def NN.IR.OpContracts.inferPoolingDims.go (tag : String) : List String → List ℕ → List ℕ → List ℕ → List ℕ → Except String (List ℕ)

def NN.IR.OpContracts.inferPoolOutShape (tag : String) (axisNames : List String) (kernels strides paddings : List ℕ) (parent : Spec.Shape) : Except String Spec.Shape

Infer a channel-first pooling shape for an arbitrary number of spatial dimensions.

def NN.IR.OpContracts.inferConvOutShape (tag : String) (inChannels outChannels : ℕ) (axisNames : List String) (kernels strides paddings : List ℕ) (parent : Spec.Shape) : Except String Spec.Shape

Infer a channel-first convolution shape for an arbitrary number of spatial dimensions.

The declared input-channel count is checked against the leading input axis. The output-channel count becomes the leading output axis; all remaining axes are inferred by inferSlidingWindowDims.

def NN.IR.OpContracts.inferPool2dOutShape (tag : String) (kH kW stride padding : ℕ) (parent : Spec.Shape) : Except String Spec.Shape

Infer the output of the current two-dimensional pooling IR operators.

def NN.IR.OpContracts.inferConv2dOutShape (inC outC kH kW stride padding : ℕ) (parent : Spec.Shape) : Except String Spec.Shape

Infer the output of the current two-dimensional convolution IR operator.

def NN.IR.OpContracts.inferBatchNorm2dNchwEvalOutShape (channels : ℕ) (parent : Spec.Shape) : Except String Spec.Shape

Output shape for eval-mode BatchNorm2d on NCHW tensors.