CROWN Graph

Graph-based LiRPA scaffolding (IBP + CROWN-style affine bounds).

This file is the "graph engine" counterpart to NN.MLTheory.CROWN.Core:

• core defines the scalar/box/affine primitives (Box, AffineVec, IBP.linear),
• graph lifts those primitives to arbitrary tensor DAGs (NN.IR.Graph).

What this module is for:

• Representing computation graphs over typed tensors.
• Propagating interval bounds forward (IBP).
• Propagating affine bounds (CROWN/DeepPoly style) w.r.t. a designated input node.
• Optionally running objective-dependent backward passes for tighter bounds.

Design notes:

• We store flattened bounds (FlatBox, FlatAffine*) so we can reuse AffineVec unchanged.
• Sequence models are represented by unrolling recurrent cells into a DAG (shared parameters, repeated ops).
• Transformer-style models rely on matmul/softmax/elementwise ops; relaxations are added incrementally.

References:

• auto_LiRPA: Xu et al., "Automatic Perturbation Analysis for Scalable Certified Robustness and Beyond", NeurIPS 2020, arXiv:2002.12920 (https://arxiv.org/abs/2002.12920).
• CROWN: Zhang et al., "Efficient Neural Network Robustness Certification with General Activation Functions", arXiv:1811.00866 (https://arxiv.org/abs/1811.00866).

PyTorch analogues (conceptual):

• torch.fx graphs as a user-facing DAG representation: https://pytorch.org/docs/stable/fx.html
• auto_LiRPA as a practical LiRPA implementation over PyTorch graphs.

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph : Type

Typed IR graph used by the CROWN/LiRPA engines.

def NN.MLTheory.CROWN.FlatBox.lInfBox {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (center radius : TorchLean.Tensor α s) : FlatBox α

Flatten a shaped center/radius pair into the graph-level interval-box representation.

def NN.MLTheory.CROWN.FlatBox.lInfBall {α : Type} [TorchLean.Storage α] [Context α] {s : Spec.Shape} (center : TorchLean.Tensor α s) (eps : α) : FlatBox α

Uniform ℓ∞ box around a shaped tensor.

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph.OpKind : Type

IR operation kinds used by the graph engine.

@[reducible, inline]

abbrev NN.MLTheory.CROWN.Graph.Node : Type

IR nodes used by the graph engine.

structure NN.MLTheory.CROWN.Graph.FlatAffine (α : Type) [TorchLean.Storage α] [Context α] : Type

Flattened affine form for a node output with respect to a fixed flattened input.

This represents y ≈ A*x + c for a chosen input node x.

• inDim : ℕ

  Flattened input dimension.

• outDim : ℕ

  Flattened output dimension.

• aff : AffineVec α self.inDim self.outDim

  Affine form A,c.

structure NN.MLTheory.CROWN.Graph.FlatAffineBounds (α : Type) [TorchLean.Storage α] [Context α] : Type

Flattened affine lower and upper bounds for a node output w.r.t. a fixed flattened input.

• inDim : ℕ

  Flattened input dimension.

• outDim : ℕ

  Flattened output dimension.

• loAff : AffineVec α self.inDim self.outDim

  Lower affine bound form.

• hiAff : AffineVec α self.inDim self.outDim

  Upper affine bound form.

def NN.MLTheory.CROWN.Graph.FlatAffine.evalOnFlatBox {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] (aff : FlatAffine α) (xB : FlatBox α) (hIn : xB.dim = aff.inDim) : Box α (Spec.Shape.dim aff.outDim Spec.Shape.scalar)

Evaluate a flattened affine form on a flattened input box after checking the input dimension.

def NN.MLTheory.CROWN.Graph.FlatAffine.evalOnFlatBoxAsDim {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] (aff : FlatAffine α) (xB : FlatBox α) (hIn : xB.dim = aff.inDim) {m : ℕ} (hOut : aff.outDim = m) : Box α (Spec.Shape.dim m Spec.Shape.scalar)

Evaluate and view the output at a checked vector dimension.

def NN.MLTheory.CROWN.Graph.FlatAffineBounds.evalOnFlatBox {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] (bounds : FlatAffineBounds α) (xB : FlatBox α) (hIn : xB.dim = bounds.inDim) : Box α (Spec.Shape.dim bounds.outDim Spec.Shape.scalar)

Evaluate lower/upper affine bounds on a flattened input box.

The lower affine form contributes the lower endpoint; the upper affine form contributes the upper endpoint. This is the common CROWN workflow shape.

def NN.MLTheory.CROWN.Graph.FlatAffineBounds.evalOnFlatBoxAsDim {α : Type} [TorchLean.Storage α] [Context α] [BoundOps α] (bounds : FlatAffineBounds α) (xB : FlatBox α) (hIn : xB.dim = bounds.inDim) {m : ℕ} (hOut : bounds.outDim = m) : Box α (Spec.Shape.dim m Spec.Shape.scalar)

Evaluate lower/upper affine bounds and view the output at a checked vector dimension.

structure NN.MLTheory.CROWN.Graph.NodeState (α : Type) [TorchLean.Storage α] [Context α] : Type

Per-node bound state (flattened).

The option fields record which analyses have populated a node: an interval-only pass fills ibp?, while affine CROWN passes additionally fill aff?.

• shape : Spec.Shape

  Original (unflattened) tensor shape of the node output.

• ibp? : Option (FlatBox α)

  Interval bounds (IBP) if available.

• aff? : Option (FlatAffine α)

  Affine form if available.

structure NN.MLTheory.CROWN.Graph.PropState (α : Type) [TorchLean.Storage α] [Context α] : Type

Propagation workspace across the whole graph.

• inputId : ℕ

  Which node id is treated as the designated input for affine bounds.

• inputDim : ℕ

  Flattened input dimension.

• states : Array (NodeState α)

  Per-node bound states.