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.

Shape inference and the reference semantics both use this definition. It is a simp lemma so that proofs about either pass see the concrete one-axis shape.

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.checkReductionAxis (axis : ℕ) (s : Spec.Shape) : Except String (PLift (Spec.Shape.NonemptyAxis axis s))

Validate the axis of an axis-dropping reduction (reduceSum, reduceMean).

The typed reductions Tensor.reduceSum and Tensor.reduceMean require the reduced axis to have positive extent, so the shared IR rule is Spec.Shape.nonemptyAxis? rather than a bare rank check: an in-bounds axis of extent zero is rejected by both inference and evaluation. The result carries the evidence needed by the typed operation.

This is a simp definition so proofs that already know the nonemptyAxis? verdict can reduce it.

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.transposePerm (rank axis₁ axis₂ : ℕ) : Except String (Array ℕ)

Axis permutation that swaps axis₁ and axis₂ and leaves every other axis fixed.

def NN.IR.OpContracts.inferTransposeOutShape (axis₁ axis₂ : ℕ) (parent : Spec.Shape) : Except String Spec.Shape

Infer the shape obtained by swapping two arbitrary axes.

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

Compute the matrix dimensions 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 matrix:

• the row count is the product of dimensions before axis (dims.take axis);
• the column count is the product of dimensions from axis onward (dims.drop axis).

LayerNorm then runs over each row and the result is reshaped to s. This construction works for every nonempty tensor rank; the matrix is an evaluation view, not a restriction on the input shape.

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

Compute the inverse of a permutation array.

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.permMoveAxisToLast (axis : ℕ) (s : Spec.Shape) : Except String (Array ℕ)

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 (Array ℕ)

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].

structure NN.IR.OpContracts.MatmulDims : Type

Decomposition of a pair of matmul operand shapes into a shared leading shape and the three matrix extents.

The left operand has shape leading ++ [rows, inner] and the right operand has shape leading ++ [inner, cols]. Shape inference reads the output shape from this record and the reference semantics uses the same record to recover the typed operands, so there is one matmul shape rule.

• leading : Spec.Shape

  Batch axes shared by both operands.

• rows : ℕ

  Row count of the left operand.

• inner : ℕ

  Contracted extent shared by both operands.

• cols : ℕ

  Column count of the right operand.

def NN.IR.OpContracts.MatmulDims.leftShape (dims : MatmulDims) : Spec.Shape

Shape of the left matmul operand.

def NN.IR.OpContracts.MatmulDims.rightShape (dims : MatmulDims) : Spec.Shape

Shape of the right matmul operand.

def NN.IR.OpContracts.MatmulDims.outShape (dims : MatmulDims) : Spec.Shape

Shape of the matmul result.

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

Decompose two matmul operand shapes.

Both inputs must have rank at least two and exactly the same leading shape. The final two axes follow the usual matrix rule: (...×m×n) · (...×n×p) → (...×m×p).

This is a simp definition so that proofs about concrete operand shapes reduce the match.

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

Infer the output shape for matmul from the two parent shapes (see matmulDims).

theorem NN.IR.OpContracts.matmulDims_shapes {a b : Spec.Shape} {dims : MatmulDims} (h : matmulDims a b = Except.ok dims) : a = dims.leftShape ∧ b = dims.rightShape

A successful decomposition determines the operand shapes it was computed from.

def NN.IR.OpContracts.mergeConcatDims (axis : ℕ) : ℕ → List ℕ → List ℕ → Except String (List ℕ)

Merge one concat input into the accumulated dimensions.

def NN.IR.OpContracts.foldConcatDims (axis : ℕ) : List ℕ → List Spec.Shape → Except String (List ℕ)

Fold compatible concat inputs into the accumulated output dimensions.

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

Infer the output shape for concat from an array of parent shapes.

Every parent must have the same rank, the selected axis must exist, and all dimensions other than that axis must agree. The selected dimensions are summed. Both the axis and the tensor rank are arbitrary.

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.effectiveKernel (kernel dilation : ℕ) : ℕ

Effective kernel width for a dilated window.

def NN.IR.OpContracts.slideOutDilated (input kernel stride dilation paddingBefore paddingAfter : ℕ) : ℕ

Output length for a dilated window with independent low/high padding.

def NN.IR.OpContracts.inferConvDims (tag : String) (axisNames : List String) (inputs kernels strides dilations paddingBefore paddingAfter : List ℕ) : Except String (List ℕ)

Infer dilated convolution dimensions from one parameter per spatial axis.

def NN.IR.OpContracts.inferConvDims.go (tag : String) : List String → List ℕ → List ℕ → 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.

Pooling uses poolOutDim, which assigns an empty output to an empty input axis, an oversized window, or padding outside the pooling domain. Kernels and strides must still be positive.

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

structure NN.IR.OpContracts.PoolPlan (config : WindowConfig) (parent : Spec.Shape) : Type

Validated shape information for pooling over a tensor suffix.

The plan is the common boundary between shape inference, denotational evaluation, and executable lowering. It carries the exact split and positivity evidence required by the typed pooling operators, so those passes cannot silently disagree about which axes are spatial.

• leading : Spec.Shape

  Axes preserved before the pooled suffix.

• spatial : TorchLean.Tensor ℕ [config.spatialRank]

  Input extents along the pooled axes.

• concat_eq : self.leading.concat (Spec.Shape.ofList (self.spatial.to (List ℕ))) = parent

  The prefix and spatial suffix reconstruct the parent shape.

• kernelNonzero (axis : Fin config.spatialRank) : config.kernel.getScalar axis ≠ 0

  Every kernel extent is nonzero.

• strideNonzero (axis : Fin config.spatialRank) : config.stride.getScalar axis ≠ 0

  Every stride is nonzero.

def NN.IR.OpContracts.PoolPlan.outShape {config : WindowConfig} {parent : Spec.Shape} (plan : PoolPlan config parent) : Spec.Shape

Output shape computed by a validated pooling plan.

def NN.IR.OpContracts.planPool (tag : String) (config : WindowConfig) (parent : Spec.Shape) : Except String (PoolPlan config parent)

Validate and plan pooling over a spatial suffix of arbitrary rank.

Every preceding axis is preserved, so one operation handles unbatched tensors, ordinary batches, and tensors with several leading batch dimensions.

def NN.IR.OpContracts.inferWindowOutShape (tag : String) (config : WindowConfig) (parent : Spec.Shape) : Except String Spec.Shape

Infer the output shape of an arbitrary-rank pooling operation from its window geometry.

This is the shape rule shared by Infer.nodeOutShape and the reference semantics: both call planPool on the same configuration, so the evaluator's pooled tensor has exactly this shape.

def NN.IR.OpContracts.inferPoolOutShape (tag : String) {spatialRank : ℕ} (kernels strides paddings : TorchLean.Tensor ℕ [spatialRank]) (parent : Spec.Shape) : Except String Spec.Shape

Infer the output shape of an arbitrary-rank pooling operation.

def NN.IR.OpContracts.inferConvConfigOutShape (tag : String) (config : ConvConfig) (parent : Spec.Shape) : Except String Spec.Shape

Infer the output shape for the full parameterized convolution configuration.

def NN.IR.OpContracts.inferConvOutShape (tag : String) (channelAxis inChannels outChannels : ℕ) {spatialRank : ℕ} (kernels strides paddings : TorchLean.Tensor ℕ [spatialRank]) (parent : Spec.Shape) : Except String Spec.Shape

Infer dense, unit-dilation convolution output geometry for an arbitrary spatial rank.

This is the ordinary convolution specialization of inferConvConfigOutShape; keeping one checker prevents the default and parameterized APIs from assigning different shapes to the same operation.

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

Check eval-mode BatchNorm metadata against an arbitrary channel axis.