Affine Propagation

This module contains the one-sided affine pass for the flat graph engine. It builds an upper affine form for each node with respect to a chosen input node. Linear and structural operations preserve affine dependence. A nonlinear operation uses the upper endpoint of its checked IBP box as a constant affine bound unless a directed affine rule is available.

structure NN.MLTheory.CROWN.Graph.AffineCtx : Type

Context for affine (CROWN/DeepPoly) propagation.

Affine bounds are computed with respect to a single designated input node, whose flattened dimension is inputDim.

• inputId : ℕ

  Node id treated as the input variable for affine bounds.

• inputDim : ℕ

  Flattened input dimension.

def NN.MLTheory.CROWN.Graph.affIdentity {α : Type} [Context α] (n : ℕ) : AffineVec α n n

Identity affine map on a flattened vector of length n.

def NN.MLTheory.CROWN.Graph.affAdd {α : Type} [Context α] {n m : ℕ} (a1 a2 : AffineVec α n m) : AffineVec α n m

Pointwise addition of two affine maps with the same input and output dimensions.

def NN.MLTheory.CROWN.Graph.affSub {α : Type} [Context α] {n m : ℕ} (a1 a2 : AffineVec α n m) : AffineVec α n m

Pointwise subtraction of two affine maps with the same input and output dimensions.

def NN.MLTheory.CROWN.Graph.affOfConv2d {α : Type} [Context α] (cfg : Conv2DParams α) : have outH := Spec.Shape.slidingWindowOutDim cfg.inH cfg.kH cfg.stride cfg.padding; have outW := Spec.Shape.slidingWindowOutDim cfg.inW cfg.kW cfg.stride cfg.padding; AffineVec α (cfg.inC * cfg.inH * cfg.inW) (cfg.outC * outH * outW)

Flatten a typed convolution into the affine map it denotes.

The CROWN pass uses this when a convolution is linear in the selected input. Keeping the conversion here lets convolution share the same affine machinery as linear and matmul nodes.

Precondition: cfg.stride ≠ 0. Engine call sites check this before calling the converter.

def NN.MLTheory.CROWN.Graph.propagateAffineNode {α : Type} [Context α] (nodes : Array Node) (ps : ParamStore α) (ibp : Array (Option (FlatBox α))) (affs : Array (Option (FlatAffine α))) (ctx : AffineCtx) (id : ℕ) : Array (Option (FlatAffine α))

Propagate a single node’s affine form (CROWN/DeepPoly style) given parent affine forms.

This updates the affs array at index id when the node kind admits an affine transfer rule. For non-affine nodes (or missing parents/params), the array is left unchanged so downstream code can fall back to IBP boxes.

def NN.MLTheory.CROWN.Graph.runAffine {α : Type} [Context α] (g : Graph) (ps : ParamStore α) (ctx : AffineCtx) (ibp : Array (Option (FlatBox α))) : Array (Option (FlatAffine α))

Run the one-sided affine pass.

Linear nodes keep their affine dependence on the selected input. Nonlinear nodes use their checked IBP upper endpoint as a constant affine bound unless this pass has a separately justified rule.