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, and check declared shapes with Graph.checkShapes.
• NN.IR.Check: state WellFormed and WellShaped as propositions over those checkers.
• 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 normalized suffix must be nonempty; the leading batch may be empty.

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

Geometry of a matrix product after broadcasting batch axes and promoting vectors.

A left vector is viewed as a one-row matrix and a right vector as a one-column matrix. Their inserted axes are removed from the output; the original operand shapes are retained for validation and flat indexing.

• leading : Spec.Shape

  Broadcast batch shape of the output.

• rows : ℕ

  Row count of the left operand.

• inner : ℕ

  Contracted extent shared by both operands.

• cols : ℕ

  Column count of the right operand.

• leftShape : Spec.Shape

  Original left operand shape.

• rightShape : Spec.Shape

  Original right operand shape.

• leftLeading : Spec.Shape

  Batch axes of the left operand, before broadcasting.

• rightLeading : Spec.Shape

  Batch axes of the right operand, before broadcasting.

• leftVector : Bool

  Whether the left operand was promoted from a vector.

• rightVector : Bool

  Whether the right operand was promoted from a vector.

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

Shape of the matmul result.

def NN.IR.OpContracts.MatmulDims.batchIndex (source target : Spec.Shape) (index : ℕ) : ℕ

Project a flat output batch coordinate to one operand's unbroadcast batch coordinate.

def NN.IR.OpContracts.MatmulDims.leftIndex (dims : MatmulDims) (output inner : ℕ) : ℕ

Left scalar used by one output coordinate and one contraction coordinate.

def NN.IR.OpContracts.MatmulDims.rightIndex (dims : MatmulDims) (output inner : ℕ) : ℕ

Right scalar used by one output coordinate and one contraction coordinate.

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

Broadcast reversed batch dimensions, aligning from the trailing axis. Singleton axes also broadcast to zero; using max here would incorrectly turn zero into one.

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

Infer vector promotion and trailing-axis batch broadcasting for non-scalar operands.

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