Forward CROWN Bounds

This module computes lower and upper affine bounds for each graph node. The pass is DeepPoly/CROWN-style: linear nodes compose exactly, nonlinear nodes attach local relaxations, and unsupported cases fall back to constant affine bounds derived from the already-computed IBP box.

For a chosen flattened input node ctx.inputId, the pass computes a pair of affine forms loAff(x) ≤ node(x) ≤ hiAff(x) for each supported node. The transfer rules below use the usual CROWN/DeepPoly ingredients:

• Linear layers use sign-splitting (W⁺/W⁻) to combine parent bounds.
• ReLU uses the standard triangle upper bound and a simple evidence-based lower choice (0 vs x).
• Exp/log use secant/tangent bounds (convex/concave).
• Softmax and LayerNorm use conservative last-axis relaxations.

Unsupported axes or shape mismatches fall back to constant affine bounds derived from the IBP box.

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

Exact lower and upper affine bounds for the identity node.

def NN.MLTheory.CROWN.Graph.boundsConst {α : Type} [Context α] (inputDim outDim : ℕ) (lo hi : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)) : FlatAffineBounds α

Constant affine bounds with zero coefficient matrix and explicit lower/upper offsets.

def NN.MLTheory.CROWN.Graph.matPos {α : Type} [Context α] {m n : ℕ} (W : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Positive part of a matrix, used for sign-splitting affine bounds through linear layers.

def NN.MLTheory.CROWN.Graph.matNeg {α : Type} [Context α] {m n : ℕ} (W : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Negative part of a matrix, used with matPos to propagate lower and upper affine forms.

def NN.MLTheory.CROWN.Graph.materializeAffineVec {α : Type} {inDim outDim : ℕ} (a : AffineVec α inDim outDim) : AffineVec α inDim outDim

Materialize an affine vector so later graph passes do not accumulate deep tensor closures.

def NN.MLTheory.CROWN.Graph.propagateLinearBounds {α : Type} [Context α] {n m : ℕ} (W : Spec.Tensor α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (b : Spec.Tensor α (Spec.Shape.dim m Spec.Shape.scalar)) (xB : FlatAffineBounds α) (hout : xB.outDim = n) : FlatAffineBounds α

Propagate affine lower/upper bounds through an affine layer W*x + b.

def NN.MLTheory.CROWN.Graph.affApplyDiag {α : Type} [Context α] {inDim outDim : ℕ} (slopes bias : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)) (aff : AffineVec α inDim outDim) : AffineVec α inDim outDim

Compose an affine form with a diagonal affine relaxation slope * x + bias.

def NN.MLTheory.CROWN.Graph.affApplyDiagSignedUpper {α : Type} [Context α] {inDim outDim : ℕ} (slopes bias : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)) (xLo xHi : AffineVec α inDim outDim) : AffineVec α inDim outDim

Apply a diagonal relaxation for an upper bound, selecting parent rows by slope sign.

def NN.MLTheory.CROWN.Graph.affApplyDiagSignedLower {α : Type} [Context α] {inDim outDim : ℕ} (slopes bias : Spec.Tensor α (Spec.Shape.dim outDim Spec.Shape.scalar)) (xLo xHi : AffineVec α inDim outDim) : AffineVec α inDim outDim

Apply a diagonal relaxation for a lower bound, selecting parent rows by slope sign.

def NN.MLTheory.CROWN.Graph.sigmoidLineBounds {α : Type} [Context α] (l u : α) : α × α × α × α

One-dimensional sigmoid linear bounds on an interval [l,u].

Returns (a_lo, b_lo, a_hi, b_hi) for lower and upper affine lines.

def NN.MLTheory.CROWN.Graph.propagateReluBounds {α : Type} [Context α] (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Propagate CROWN bounds through ReLU using standard per-neuron triangle relaxations.

def NN.MLTheory.CROWN.Graph.propagateReluBoundsWithAlpha {α : Type} [Context α] (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) (alpha : Spec.Tensor α (Spec.Shape.dim preB.dim Spec.Shape.scalar)) : FlatAffineBounds α

Propagate ReLU bounds with externally supplied α slopes for crossing neurons.

def NN.MLTheory.CROWN.Graph.propagateExpBounds {α : Type} [Context α] (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Propagate CROWN bounds through exp using tangent/secant relaxations.

def NN.MLTheory.CROWN.Graph.propagateLogBounds {α : Type} [Context α] (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Propagate CROWN bounds through log on the positive-domain convention used by the verifier.

def NN.MLTheory.CROWN.Graph.propagateSigmoidBounds {α : Type} [Context α] (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Propagate CROWN bounds through sigmoid with convex/concave one-dimensional relaxations.

def NN.MLTheory.CROWN.Graph.propagateTanhBounds {α : Type} [Context α] (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Propagate CROWN bounds through tanh with convex/concave one-dimensional relaxations.

def NN.MLTheory.CROWN.Graph.permuteAffineOut {α : Type} {inDim outDim : ℕ} (perm : Fin outDim → Fin outDim) (aff : AffineVec α inDim outDim) : AffineVec α inDim outDim

Permute the output coordinates of an affine bound when the output shape permutation is valid.

def NN.MLTheory.CROWN.Graph.propagateSoftmaxBoundsLastAxis {α : Type} [Context α] (s : Spec.Shape) (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Conservative CROWN-style affine bounds for softmax along the last tensor axis.

def NN.MLTheory.CROWN.Graph.propagateLayernormBoundsLastAxis {α : Type} [Context α] (s : Spec.Shape) (preB : FlatBox α) (xB : FlatAffineBounds α) (hout : xB.outDim = preB.dim) : FlatAffineBounds α

Conservative affine bounds for layer normalization over the last tensor axis.

def NN.MLTheory.CROWN.Graph.propagateMulElemBounds {α : Type} [Context α] (Bx By : FlatBox α) (xB yB : FlatAffineBounds α) (houtX : xB.outDim = Bx.dim) (houtY : yB.outDim = By.dim) : Option (FlatAffineBounds α)

Propagate affine bounds through componentwise multiplication using per-coordinate product planes.

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

Propagate a single node’s affine bounds (lower/upper) given parent bounds.

This is the CROWN/DeepPoly-style transfer step used by runCROWN. For node kinds without a dedicated rule, we fall back to the IBP enclosure (turned into a constant affine bound).

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

Run the basic CROWN affine-bounds pass; requires prior IBP for per-node intervals.

def NN.MLTheory.CROWN.Graph.evalCROWNOutputBox? {α : Type} [Context α] [BoundOps α] (bounds : Array (Option (FlatAffineBounds α))) (xB : FlatBox α) (outputId inputDim : ℕ) : Except String (FlatBox α)

Evaluate already-computed CROWN output affine bounds on an input box.

def NN.MLTheory.CROWN.Graph.outputBoxCROWN? {α : Type} [Context α] [BoundOps α] (g : Graph) (ps : ParamStore α) (xB : FlatBox α) (inputId outputId inputDim : ℕ) : Except String (FlatBox α)

Run IBP, run forward CROWN, and evaluate the output affine bounds on the selected input box.

This is the common "forward CROWN output box" workflow. It keeps callers from open-coding the same output-array lookup and input-dimension proof checks around runCROWN.

def NN.MLTheory.CROWN.Graph.ParamStore.outputBoxCROWN? {α : Type} [Context α] [BoundOps α] (ps : ParamStore α) (g : Graph) (xB : FlatBox α) (inputId outputId inputDim : ℕ) : Except String (FlatBox α)

Run outputBoxCROWN? from an input-seeded parameter store.

This method form reads naturally at call sites that already thread a ParamStore.