Shared definitions for the graph CROWN engine.

This file contains the flat vector representation, parameter stores, interval boxes, shape permutation helpers, and tensor casts used by the IBP, derivative, affine, CROWN, and backward objective passes.

structure NN.MLTheory.CROWN.Graph.FlatVec (α : Type) [Context α] : Type

Flat vector pack: a tensor paired with its flattened dimension.

This is used for constant payloads and objective coefficient vectors in the flat LiRPA engine.

• n : ℕ

  Vector dimension.

• v : Spec.Tensor α (Spec.Shape.dim self.n Spec.Shape.scalar)

  Vector payload (shape .dim n .scalar).

structure NN.MLTheory.CROWN.Graph.LinParams (α : Type) [Context α] : Type

Parameters for a linear layer y = W*x + b in flattened form.

m is the output dimension and n is the input dimension.

• m : ℕ

  Output dimension.

• n : ℕ

  Input dimension.

• w : Spec.Tensor α (Spec.Shape.dim self.m (Spec.Shape.dim self.n Spec.Shape.scalar))

  Weight matrix W (shape m × n).

• b : Spec.Tensor α (Spec.Shape.dim self.m Spec.Shape.scalar)

  Bias vector b (shape m).

structure NN.MLTheory.CROWN.Graph.MatParams (α : Type) [Context α] : Type

Matrix parameters for bias-free matmul: y = W x.

• m : ℕ

  Output dimension.

• n : ℕ

  Input dimension.

• w : Spec.Tensor α (Spec.Shape.dim self.m (Spec.Shape.dim self.n Spec.Shape.scalar))

  Weight matrix W (shape m × n).

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph.Conv2DParams (α : Type) [Context α] : Type

Conv2D parameters with cached spatial dimensions for graph propagation.

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph.BatchNorm2DNchwEvalParams (α : Type) [Context α] : Type

Eval-mode BatchNorm2d parameters for an N×C×H×W node.

def NN.MLTheory.CROWN.Graph.nchwChannelOfFlat (c h w idx : ℕ) : ℕ

Channel index for a flattened N×C×H×W tensor in row-major order.

def NN.MLTheory.CROWN.Graph.batchNorm2dNchwEvalScale {α : Type} [Context α] (cfg : BatchNorm2DNchwEvalParams α) (ci : Fin cfg.c) : α

Eval BatchNorm scale for one channel.

def NN.MLTheory.CROWN.Graph.batchNorm2dNchwEvalBias {α : Type} [Context α] (cfg : BatchNorm2DNchwEvalParams α) (ci : Fin cfg.c) : α

Eval BatchNorm bias for one channel after folding running statistics into an affine map.

def NN.MLTheory.CROWN.Graph.batchNorm2dNchwEvalLinear? {α : Type} [Context α] (parentShape : Spec.Shape) (cfg : BatchNorm2DNchwEvalParams α) : Option (LinParams α)

Build the exact diagonal affine form for eval-mode BatchNorm2d over an N×C×H×W tensor.

The IR stores the channel parameters in the node payload. The spatial dimensions come from the checked parent shape, so malformed shapes simply do not produce a verifier transfer rule.

structure NN.MLTheory.CROWN.Graph.ParamStore (α : Type) [Context α] : Type

Parameters keyed by node id (weights, biases, constants, and seeded input boxes).

This is kept compact: it is the graph interpreter used to run IBP/CROWN on a pure Graph without pulling in a heavyweight runtime.

• inputBoxes : Std.HashMap ℕ (FlatBox α)

  Seed boxes for designated input nodes (id -> FlatBox).

• constVals : Std.HashMap ℕ (FlatVec α)

  Constants (id -> FlatVec).

• linearWB : Std.HashMap ℕ (LinParams α)

  Linear layer params (id -> (W,b)).

• matmulW : Std.HashMap ℕ (MatParams α)

  Matmul params (id -> W) for bias-free multiplication.

• conv2dCfg : Std.HashMap ℕ (Conv2DParams α)

  Conv2d specs (id -> conv configuration).

• batchNorm2dNchwEval : Std.HashMap ℕ (BatchNorm2DNchwEvalParams α)

  Eval-mode BatchNorm2d parameters (id -> gamma/beta/running stats).

def NN.MLTheory.CROWN.Graph.ParamStore.seedInputBox {α : Type} [Context α] (ps : ParamStore α) (inputId : ℕ) (xB : FlatBox α) : ParamStore α

Insert an input interval box for a graph node.

def NN.MLTheory.CROWN.Graph.ParamStore.seedLInfBall {α : Type} [Context α] {s : Spec.Shape} (ps : ParamStore α) (inputId : ℕ) (center : Spec.Tensor α s) (eps : α) : ParamStore α

Seed a graph input with a uniform ℓ∞ box around a shaped tensor.

def NN.MLTheory.CROWN.Graph.outputBox? {α : Type} [Context α] (boxes : Array (Option (FlatBox α))) (outId : ℕ) : Except String (FlatBox α)

Read a node's interval box from an IBP-style result array.

@[implicit_reducible]

instance NN.MLTheory.CROWN.Graph.instInhabitedFlatBox {α : Type} [Context α] : Inhabited (FlatBox α)

Default inhabitant for FlatBox (a 0-dimensional box at 0).

def NN.MLTheory.CROWN.Graph.box_mul_elem {α : Type} [Context α] [BoundOps α] (B1 B2 : FlatBox α) : Option (FlatBox α)

Elementwise product of two FlatBoxes (interval product per component). Requires equal dims.

def NN.MLTheory.CROWN.Graph.derivBoxExp {α : Type} [Context α] (zB : FlatBox α) : FlatBox α

Derivative range for exp over a value box; exp' = exp is monotone.

def NN.MLTheory.CROWN.Graph.derivBoxLog {α : Type} [Context α] (zB : FlatBox α) : FlatBox α

Derivative range for the positive-domain log rule used by derivative propagation.

def NN.MLTheory.CROWN.Graph.chainMul {α : Type} [Context α] [BoundOps α] (dZ dF : FlatBox α) : Option (FlatBox α)

Chain-rule multiplication for derivative intervals. Returns none on dimension mismatch.

def NN.MLTheory.CROWN.Graph.toFlatBox {α : Type} [Context α] (n : ℕ) (B : Box α (Spec.Shape.dim n Spec.Shape.scalar)) : FlatBox α

Convert a dependent Box of shape .dim n .scalar into a FlatBox with dim := n.

def NN.MLTheory.CROWN.Graph.ofFlatBox {α : Type} [Context α] (B : FlatBox α) : Box α (Spec.Shape.dim B.dim Spec.Shape.scalar)

Convert a FlatBox to a dependent Box at shape .dim B.dim .scalar.

def NN.MLTheory.CROWN.Graph.box_add {α : Type} [Context α] [BoundOps α] (B1 B2 : FlatBox α) : FlatBox α

Add two flat interval boxes coordinatewise; dimension mismatches preserve the left box.

def NN.MLTheory.CROWN.Graph.box_sub {α : Type} [Context α] [BoundOps α] (B1 B2 : FlatBox α) : FlatBox α

Interval subtraction on FlatBox endpoints (sound enclosure).

def NN.MLTheory.CROWN.Graph.box_relu {α : Type} [Context α] (B : FlatBox α) : FlatBox α

Apply ReLU to both endpoints of a FlatBox (monotone activation, so endpoints suffice).

def NN.MLTheory.CROWN.Graph.boxAbs {α : Type} [Context α] (B : FlatBox α) : FlatBox α

Componentwise absolute value bounds. Soundly encloses abs over each interval component.

def NN.MLTheory.CROWN.Graph.boxSqrt {α : Type} [Context α] (B : FlatBox α) : FlatBox α

Componentwise sqrt bounds. Uses the spec semantics sqrt(max(x,0)), which is monotone.

def NN.MLTheory.CROWN.Graph.boxInv {α : Type} [Context α] (B : FlatBox α) : FlatBox α

Componentwise reciprocal bounds using operators.arithmetic.ibp_reciprocal.

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph.FlatDVal (α : Type) [Context α] : Type

Dynamic tensor value used while reshaping and permuting flattened boxes.

def NN.MLTheory.CROWN.Graph.flatDValShape {α : Type} [Context α] (v : FlatDVal α) : Spec.Shape

Shape projection for FlatDVal.

def NN.MLTheory.CROWN.Graph.flatDValTensor {α : Type} [Context α] (v : FlatDVal α) : Spec.Tensor α (flatDValShape v)

Tensor projection for FlatDVal, preserving the dependent shape stored beside it.

def NN.MLTheory.CROWN.Graph.findIndex? (xs : List ℕ) (x : ℕ) : Option ℕ

Return the position of x in a list, if present.

def NN.MLTheory.CROWN.Graph.swapAt (xs : List ℕ) (d : ℕ) : List ℕ

Swap adjacent entries d and d+1 in a list of axis ids.

def NN.MLTheory.CROWN.Graph.swapDepthsForPerm? (perm : List ℕ) (r : ℕ) : Option (List ℕ)

Decompose an axis permutation into adjacent swaps, rejecting invalid permutations.

def NN.MLTheory.CROWN.Graph.applySwapDepth {α : Type} [Context α] (v : FlatDVal α) (d : ℕ) : FlatDVal α

Apply one adjacent-axis swap to a dynamic tensor value.

def NN.MLTheory.CROWN.Graph.permuteDVal? {α : Type} [Context α] (v : FlatDVal α) (perm : List ℕ) : Option (FlatDVal α)

Apply a full axis permutation to a dynamic tensor value when the permutation is valid.

def NN.MLTheory.CROWN.Graph.boxMaxElem {α : Type} [Context α] (B1 B2 : FlatBox α) : FlatBox α

Componentwise max bounds: max(x,y) over interval boxes.

def NN.MLTheory.CROWN.Graph.boxMinElem {α : Type} [Context α] (B1 B2 : FlatBox α) : FlatBox α

Componentwise min bounds: min(x,y) over interval boxes.

def NN.MLTheory.CROWN.Graph.boxSquare {α : Type} [Context α] (B : FlatBox α) : FlatBox α

Componentwise square of an interval box: for each component [l,u] produce [min (l^2,u^2), max (l^2,u^2)], with 0 as the minimum when the interval crosses 0.

The body is exposed because the proof-facing theorem module unfolds this executable rule when proving dimension preservation and pointwise enclosure.

def NN.MLTheory.CROWN.Graph.intervalMul {α : Type} [Context α] (aLo aHi bLo bHi : α) : α × α

Interval multiplication for scalar endpoints: given [aLo,aHi] and [bLo,bHi], return bounds on the product.

def NN.MLTheory.CROWN.Graph.lastDimLen : Spec.Shape → ℕ

Length of the last axis of a shape; scalars are treated as length one.

def NN.MLTheory.CROWN.Graph.mkValidAxis? (axis : ℕ) (s : Spec.Shape) : Option (PLift (Spec.Shape.valid_axis axis s))

Runtime witness that an axis is valid for a shape.

def NN.MLTheory.CROWN.Graph.mkCanBroadcastTo? (s₁ s₂ : Spec.Shape) : Option (s₁.CanBroadcastTo s₂)

Runtime witness that one shape can broadcast to another.

def NN.MLTheory.CROWN.Graph.ibpUnflatten {α : Type} [Context α] {s : Spec.Shape} (dim : ℕ) (t : Spec.Tensor α (Spec.Shape.dim dim Spec.Shape.scalar)) (h : dim = s.size) : Spec.Tensor α s

Reinterpret a flattened tensor as shape s when the element counts agree.

def NN.MLTheory.CROWN.Graph.ibpBroadcastTo {α : Type} [Context α] (s₁ s₂ : Spec.Shape) (Xin : FlatBox α) : Option (FlatBox α)

IBP rule for broadcasting a flattened input box to a target shape.

def NN.MLTheory.CROWN.Graph.ibpReduceSumAxis {α : Type} [Context α] (axis : ℕ) (Xin : FlatBox α) (s : Spec.Shape) : Option (FlatBox α)

IBP rule for reducing a shaped box by summing along one axis.

def NN.MLTheory.CROWN.Graph.ibpReduceMeanAxis {α : Type} [Context α] (axis : ℕ) (Xin : FlatBox α) (s : Spec.Shape) : Option (FlatBox α)

IBP rule for reducing a shaped box by averaging along one axis.

Softmax IBP (last axis)

For a 1D vector x with interval bounds l <= x <= u, a standard componentwise enclosure for softmax is:

• softmax_i(x) = exp(x_i) / sum_j exp(x_j)
• Lower bound (worst-case denominator): exp(l_i) / (exp(l_i) + sum_{j != i} exp(u_j))
• Upper bound (best-case denominator): exp(u_i) / (exp(u_i) + sum_{j != i} exp(l_j))

This uses monotonicity of exp and the fact that all terms in the denominator are nonnegative. The implementation below applies the 1D rule on the last tensor axis and recurses over leading batch dimensions.

References:

• CROWN / DeepPoly context: Zhang et al., 2018 (CROWN): https://arxiv.org/abs/1811.00866
• auto_LiRPA: Xu et al., 2020: https://arxiv.org/abs/2002.12920

def NN.MLTheory.CROWN.Graph.ibpSoftmaxLastTensor {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s → Spec.Tensor α s × Spec.Tensor α s

Interval bound propagation for softmax, applied on the last axis and lifted over leading dims.

LayerNorm IBP (last axis)

Layer normalization (Ba et al.) computes, per vector, something like:

y = (x - mean(x)) / sqrt(var(x) + eps).

We implement a conservative enclosure by:

1. Bounding mean using sums of endpoints.
2. Bounding variance using a max-deviation upper bound.
3. Bounding the per-component ratio by checking endpoint combinations against a positive denominator interval.

This is intended as a simple checker-side transfer rule. It is conservative and is not an optimized relaxation.

References:

• Ba, Kiros, Hinton, "Layer Normalization", 2016: https://arxiv.org/abs/1607.06450
• Bound propagation context: Xu et al., 2020 (auto_LiRPA): https://arxiv.org/abs/2002.12920

def NN.MLTheory.CROWN.Graph.layerNormVarianceUpper {α : Type} [Context α] {n : ℕ} (lo hi : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) (muLo muHi : α) : α

Upper bound on the variance term used by the LayerNorm interval rules.

Given endpoint bounds for a vector and bounds on its mean, each coordinate is at most max |x_i - μ| away from the bounded mean interval. Squaring and summing those coordinate radii gives the conservative variance upper bound reused by IBP, affine propagation, and derivative interval passes.

def NN.MLTheory.CROWN.Graph.layerNormMeanBounds {α : Type} [Context α] {n : ℕ} (lo hi : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) : α × α

Mean bounds for a vector whose coordinates are bounded by endpoint tensors.

def NN.MLTheory.CROWN.Graph.layerNormCenteredBounds {α : Type} [Context α] {n : ℕ} (lo hi : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) (muLo muHi : α) : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar) × Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)

Bounds for x - μ when x is bounded coordinatewise and μ is bounded by an interval.

LayerNorm transfer rules repeatedly need this centered interval for the input, first derivative, and second derivative streams. Keeping it here avoids duplicating the same endpoint arithmetic in IBP and derivative propagation.

def NN.MLTheory.CROWN.Graph.layerNormInvStdBounds {α : Type} [Context α] (varHi : α) : α × α

Bounds for the reciprocal LayerNorm denominator from an upper variance bound.

def NN.MLTheory.CROWN.Graph.ibpLayernormLastTensor {α : Type} [Context α] {s : Spec.Shape} : Spec.Tensor α s → Spec.Tensor α s → Spec.Tensor α s × Spec.Tensor α s

Interval bound propagation for layernorm, applied on the last axis and lifted over leading dims.

def NN.MLTheory.CROWN.Graph.getDimScalarFn {α : Type} {n : ℕ} (t : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) : Fin n → Spec.Tensor α Spec.Shape.scalar

For tensors known to have shape .dim n .scalar, extract the underlying function.

def NN.MLTheory.CROWN.Graph.castBoxDim {α : Type} {n n' : ℕ} (h : n = n') (B : Box α (Spec.Shape.dim n Spec.Shape.scalar)) : Box α (Spec.Shape.dim n' Spec.Shape.scalar)

Cast a 1D Box along an equality of dimensions.

def NN.MLTheory.CROWN.Graph.castRelax {α : Type} {n n' : ℕ} (h : n = n') (r : Spec.Tensor (Runtime.Ops.ReLURelax α) (Spec.Shape.dim n Spec.Shape.scalar)) : Spec.Tensor (Runtime.Ops.ReLURelax α) (Spec.Shape.dim n' Spec.Shape.scalar)

Cast a ReLU relaxation vector across a proven-equal hidden dimension.

def NN.MLTheory.CROWN.Graph.castAffineIn {α : Type} {n n' m : ℕ} (h : n = n') (a : AffineVec α n m) : AffineVec α n' m

Cast the input dimension of an affine map across a proven equality.

def NN.MLTheory.CROWN.Graph.castAffineOut {α : Type} {n m m' : ℕ} (h : m = m') (a : AffineVec α n m) : AffineVec α n m'

Cast the output dimension of an affine map across a proven equality.

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph.castDimScalar {α : Type} {n n' : ℕ} (h : n = n') (t : Spec.Tensor α (Spec.Shape.dim n Spec.Shape.scalar)) : Spec.Tensor α (Spec.Shape.dim n' Spec.Shape.scalar)

Cast a dim-scalar tensor across an equality of dimensions.

We keep this as an abbrev so it unfolds aggressively in simp-based soundness proofs.

def NN.MLTheory.CROWN.Graph.ibpLinearParams {α : Type} [Context α] [BoundOps α] (p : LinParams α) (Xin : FlatBox α) : Option (FlatBox α)

IBP propagation through explicit linear parameters.

def NN.MLTheory.CROWN.Graph.ibp_linear {α : Type} [Context α] [BoundOps α] (id : ℕ) (ps : ParamStore α) (Xin : FlatBox α) : Option (FlatBox α)

IBP propagation for a .linear node using ParamStore.linearWB.

def NN.MLTheory.CROWN.Graph.ibp_matmul {α : Type} [Context α] [BoundOps α] (id : ℕ) (ps : ParamStore α) (Xin : FlatBox α) : Option (FlatBox α)

IBP propagation for a .matmul node (bias-free) using ParamStore.matmulW.

def NN.MLTheory.CROWN.Graph.ibpConv2dNode {α : Type} [Context α] (id : ℕ) (ps : ParamStore α) (Xin : FlatBox α) : Option (FlatBox α)

IBP transfer for a convolution node whose parameters are stored in ParamStore.conv2dCfg.