Affine Propagation

This module contains the plain affine pass for the flat graph engine. It builds one affine form per node with respect to a chosen input node. Forward CROWN and objective-dependent backward CROWN build on these affine forms in their own modules.

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 := (cfg.inH + 2 * cfg.padding - cfg.kH) / cfg.stride + 1; have outW := (cfg.inW + 2 * cfg.padding - cfg.kW) / cfg.stride + 1; 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.

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 an affine pass; requires prior IBP to supply pre-activation bounds for ReLU.