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

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

Conv2D parameters with cached spatial dimensions for graph propagation.

• inC : ℕ

  Input channels.

• outC : ℕ

  Output channels.

• kH : ℕ

  Kernel height.

• kW : ℕ

  Kernel width.

• stride : ℕ

  Stride (shared for height/width).

• padding : ℕ

  Zero padding (shared for height/width).

• inH : ℕ

  Input height.

• inW : ℕ

  Input width.

• hIn : self.inC ≠ 0

  Proof that inC is nonzero (required by the typed conv spec).

• hKH : self.kH ≠ 0

  Proof that kH is nonzero.

• hKW : self.kW ≠ 0

  Proof that kW is nonzero.

• spec : Spec.Conv2DSpec self.inC self.outC self.kH self.kW self.stride self.padding α ⋯ ⋯ ⋯

  Typed conv specification payload.

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

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

This is deliberately 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).

@[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.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 α

Basic FlatBox ops

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.

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.

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

@[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.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.propagateIBPNode {α : Type} [Context α] [BoundOps α] (nodes : Array Node) (ps : ParamStore α) (boxes : Array (Option (FlatBox α))) (id : ℕ) : Array (Option (FlatBox α))

IBP propagation for one node using ParamStore.

This executable function expects parents to have already been processed. The proof layer makes that precondition explicit via TopoSorted; callers that execute graphs directly should use graphs whose parents appear before their children.

def NN.MLTheory.CROWN.Graph.runIBP {α : Type} [Context α] [BoundOps α] (g : Graph) (ps : ParamStore α) : Array (Option (FlatBox α))

Run an IBP pass over the whole graph. Caller seeds inputs via ParamStore.inputBoxes.

def NN.MLTheory.CROWN.Graph.runDeriv1D {α : Type} [Context α] [BoundOps α] (g : Graph) (ps : ParamStore α) (ibp : Array (Option (FlatBox α))) : Array (Option (FlatBox α))

Derivative IBP pass for 1D input: computes for each node an interval on dy/dx. Requires value IBP boxes for activations.

def NN.MLTheory.CROWN.Graph.runDerivDirectional {α : Type} [Context α] [BoundOps α] (g : Graph) (ps : ParamStore α) (ibp : Array (Option (FlatBox α))) (seed : FlatBox α) : Array (Option (FlatBox α))

Directional first-derivative pass: like runDeriv1D but seeds the derivative at the input node with a user-provided direction vector (as a FlatBox with lo=hi). This allows extracting partial derivatives for multi-dimensional inputs by choosing e_x, e_y, etc.

def NN.MLTheory.CROWN.Graph.runDeriv2D {α : Type} [Context α] [BoundOps α] (g : Graph) (ps : ParamStore α) (ibp d1 : Array (Option (FlatBox α))) : Array (Option (FlatBox α))

Second-derivative IBP pass for 1D input: computes per node an interval on d²y/dx². Requires value IBP boxes and first-derivative boxes. Covers input, linear/matmul, tanh, add/sub.

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

Basic CROWN-style affine bounds pass (lower + upper)

This pass computes a pair of affine forms per node:

loAff(x) ≤ node(x) ≤ hiAff(x)

where both affines are with respect to a designated flattened input node ctx.inputId.

This pass is a practical CROWN/DeepPoly implementation:

• 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/LayerNorm have conservative last-axis affine relaxations; on unsupported axes or shape mismatches we fall back to constant affine bounds derived from the IBP box.

The pass is structured so a soundness theorem can cite per-op transfer rules for the supported relaxations. Unsupported axes or shape mismatches deliberately fall back to constant affine bounds derived from the IBP box.

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.

Backward/dual CROWN (objective-dependent) propagation

runCROWN above is a forward DeepPoly-style pass: it computes nodewise affine bounds that do not depend on an objective.

Basic CROWN additionally supports backward propagation for a linear objective cᵀ · output, selecting per-neuron relaxations based on the sign of the downstream coefficients.

The implementation below is pragmatic and covers the IR ops we support in runCROWN.

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

Objective-dependent backward CROWN bound for a scalar objective.

Given a linear objective objᵀ * output, this runs a backward pass that propagates the objective coefficients through the graph, selects the relaxation attached to each node, and returns a pair of affine bounds on the objective with respect to ctx.inputId.

The returned FlatAffineBounds always has outDim = 1 (a scalar objective).

def NN.MLTheory.CROWN.Graph.runCROWNBackwardObjectiveLowerWithReluAlpha {α : Type} [Context α] (g : Graph) (ps : ParamStore α) (ctx : AffineCtx) (ibp : Array (Option (FlatBox α))) (outputId : ℕ) (obj : FlatVec α) (reluAlpha : Array (Option (FlatVec α))) : Option (AffineVec α ctx.inputDim 1)

Backward CROWN objective lower bound with externally-provided ReLU alpha slopes.

This is an integration hook for alpha-CROWN style workflows where ReLU slopes are chosen/optimized outside TorchLean and then imported as a per-node vector in reluAlpha.