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 ≥ 2, 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=3 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.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 inv[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=4 and axis=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=4 and axis=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.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

Pooling/Conv2D shape arithmetic (CHW-only)

These formulas mirror the spec/runtime conventions (CHW tensors, no dilation, symmetric padding). Centralizing them gives inference, evaluation, verification, and export code a shared convention for convolution and pooling shapes.

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

Output length for a 1D sliding-window op without padding: ⌊(in - k)/stride⌋ + 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.pool2dCHWOutShape (c inH inW kH kW stride : ℕ) : Spec.Shape

Output shape for CHW pooling without padding.

def NN.IR.OpContracts.pool2dCHWOutShapePad (c inH inW kH kW stride padding : ℕ) : Spec.Shape

Output shape for CHW pooling with symmetric padding.

def NN.IR.OpContracts.conv2dCHWOutShape (outC inH inW kH kW stride padding : ℕ) : Spec.Shape

Output shape for CHW conv2d (single-image, no batch dim).

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

Infer the output shape for CHW pooling without padding, from a parent shape.

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

Infer the output shape for CHW pooling with padding, from a parent shape.

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

Infer the output shape for CHW Conv2D, checking the declared inC against the parent shape.