BoundLab User Guide#

Mental Model#

BoundLab computes sound bounds on neural network outputs under input perturbations. The workflow has three steps:

Build a symbolic expression representing the uncertain input.
Transform it through the network (manually or via zono.interpret).
Concretize the result into numeric upper/lower bounds with ub, lb, or ublb.

\[Z = c + G\epsilon, \quad \epsilon \in [-1,1]^m\]

Expr nodes preserve dependency structure so shared symbols are not lost across operations.

Core Building Blocks#

Expressions (`boundlab.expr`)#

The central abstraction is Expr. Most workflows use:

ConstVal(tensor): deterministic value.
LpEpsilon(shape, p="inf"): perturbation symbol with bounded norm.
Affine expression composition using operators like +, -, *, @, reshape/index ops.
Cat and Stack for concatenation/stacking symbolic tensors.

Typical uncertain input construction:

import torch
import boundlab.expr as expr

x_center = torch.randn(8)
x = expr.ConstVal(x_center) + 0.1 * expr.LpEpsilon([8])

Bound Propagation (`boundlab.prop`)#

Given any expression e, BoundLab can compute:

e.ub() or boundlab.prop.ub(e): upper bound.
e.lb() or boundlab.prop.lb(e): lower bound.
e.ublb() or boundlab.prop.ublb(e): both at once (more efficient).
e.center(): midpoint of the bound interval.
e.bound_width(): ub - lb.

ub, lb = x.ublb()
print("max width:", (ub - lb).max().item())

Zonotope Interpreter (`boundlab.zono.interpret`)#

zono.interpret(program_or_graph_module) returns a callable that maps input expressions to output expressions. Interpreters accept torch.export.ExportedProgram or torch.fx.GraphModule inputs, not raw nn.Module.

import torch
from torch import nn
import boundlab.zono as zono

model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 4))
exported = torch.export.export(model, (torch.randn(8),))
op = zono.interpret(exported)
y = op(x)
ub, lb = y.ublb()

Under the hood, BoundLab dispatches each operation in the exported/FX graph to registered handlers.

Working Patterns#

Pattern 1: Exported-program interpretation (recommended)#

Use zono.interpret(exported_program) for end-to-end model flows. This is the simplest way to get started.

import torch
from torch import nn
import boundlab.expr as expr
import boundlab.zono as zono

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
x = expr.ConstVal(torch.randn(4)) + 0.1 * expr.LpEpsilon([4])

exported = torch.export.export(model, (torch.randn(4),))
op = zono.interpret(exported)
y = op(x)
ub, lb = y.ublb()

Pattern 2: Manual expression construction#

Use this when you need direct control over the symbolic structure.

import torch
import boundlab.expr as expr
import boundlab.zono as zono

x = expr.ConstVal(torch.randn(8)) + expr.LpEpsilon([8])
W = torch.randn(3, 8)
b = torch.randn(3)

y = W @ x + expr.ConstVal(b)
y = zono.interpret["relu"](y)
ub, lb = y.ublb()

Pattern 3: Soundness spot checks with sampling#

A practical validation loop:

Sample concrete perturbations around your center input.
Evaluate the concrete model.
Assert all outputs lie in [lb, ub].

This is used extensively in BoundLab tests and is a good regression guard for custom handlers.

Inspecting Expressions and Operators#

BoundLab provides several ways to inspect expressions and linear operators for debugging and exploration.

Printing an expression graph#

str(expr) renders the full expression DAG in SSA (static single assignment) form. Each node is assigned a %N name, and shared subexpressions appear once. Purely affine sub-graphs are eagerly flattened by AffineSum, so the SSA form is most informative after nonlinear operations (e.g. ReLU).

import torch
import boundlab.expr as expr
import boundlab.zono as zono

x = expr.ConstVal(torch.tensor([1.0, -0.5])) + expr.LpEpsilon([2])
y = zono.interpret["relu"](x)
print(y)

bl.Expr {
    %0 = 𝜀_<6>
    %1 = 𝜀_<1>
    %2 = <einsum [2]: [0] -> [0]>(%1) + (set_indices([1] -> [2]) ∘ <einsum [1]: [0] -> [0]>)(%0)
}

Here %1 is the original noise symbol, %0 is a fresh error symbol introduced by the ReLU linearizer, and %2 combines them via EinsumOp weights.

You can also call expr.expr_pretty_print(e) directly to get the SSA body without the wrapper.

Key `Expr` attributes#

Every expression node exposes:

e.shape — output tensor shape.
e.children — tuple of child expressions.
e.flags — ExprFlags (e.g. IS_AFFINE, SYMMETRIC_TO_0, IS_CONST).
e.id — unique time-ordered integer (used for topological ordering).
repr(e) — compact one-line summary: AffineSum(id=5, flags=ExprFlags.IS_AFFINE).

For AffineSum nodes specifically:

e.children_dict — dict[Expr, LinearOp] mapping each child to its operator.
e.constant — optional constant tensor (None if absent).

for child, op in x.children_dict.items():
    print(f"{repr(child)}  ->  {op}  ({op.input_shape} -> {op.output_shape})")

Inspecting `LinearOp`#

str(op) gives a human-readable description of the operator. Composed operators display their chain with ∘:

from boundlab.linearop import ScalarOp, ReshapeOp

a = ScalarOp(2.0, torch.Size([2, 3]))
b = ReshapeOp(torch.Size([2, 3]), (6,))
print(b @ a)  # (reshape([6]) ∘ 2.0)

Key attributes:

op.input_shape, op.output_shape — tensor shapes.
op.flags — LinearOpFlags (e.g. IS_NON_NEGATIVE).

You can materialize the full Jacobian for debugging:

op.jacobian() — efficient closed-form Jacobian if the subclass implements it (returns NotImplemented otherwise).
op.force_jacobian() — always works by applying the operator to a basis; returns a tensor of shape (*output_shape, *input_shape).

j = (b @ a).force_jacobian()
print(j.shape)  # torch.Size([6, 2, 3])

You can also test operators on concrete tensors:

op.forward(x) — apply the linear map.
op.backward(grad) — apply the transpose.
op.vforward(x) / op.vbackward(grad) — batched variants over trailing / leading dimensions.

Supported Operations#

Affine and shape operations#

BoundLab includes handlers for many affine operations, including:

Arithmetic: add, sub, neg, scalar/tensor mul, division.
Linear layers: nn.Linear, torch.nn.functional.linear.
Shape transforms: reshape, view, flatten, permute, transpose, unsqueeze, squeeze, contiguous.
Indexing and scatter/gather primitives through expression methods.

Nonlinear zonotope linearizers#

The zonotope domain currently registers linearizers/handlers for:

relu
exp
reciprocal (for positive-domain inputs)
tanh
softmax (implemented as composition of exp, sum, reciprocal, element-wise product)
bilinear matmul (McCormick-style relaxation when both operands are symbolic)

Internals#

This section covers implementation details for users who want to understand or extend BoundLab.

Linear Operators (`boundlab.linearop`)#

boundlab.linearop provides linear maps used by expression backpropagation and affine rewrites. Rather than materializing full Jacobian matrices, LinearOp stores structured/sparse representations.

Most important types:

EinsumOp: general tensor-linear map (Einstein notation style).
ComposedOp: composition of linear maps (outer ∘ inner).
SumOp: sum of linear maps with matching input/output shapes.

Supporting families include:

Scalar and structural maps (ScalarOp, ZeroOp).
Shape transforms (ReshapeOp, PermuteOp, FlattenOp, TransposeOp, …).
Indexing transforms (GetSliceOp, SetSliceOp, GatherOp, ScatterOp, …).

In bound propagation, each Expr.backward(...) returns child weights as LinearOp objects. These are composed and merged using @ and +.

How `ub`, `lb`, and `ublb` work#

Concretization uses backward traversal over the expression DAG:

Start from identity weight at the query expression.
Pop expressions in reverse-topological order.
Call local backward(...) rules.
Accumulate bias terms and push child weights.
Finish when all reachable leaves are processed.

ublb optimizes this by jointly propagating upper and lower directions and reusing exact paths when available.

Important implementation details:

Symmetric leaves (for example LpEpsilon) can reuse one side by sign flip.
Shared subexpressions are merged by weight accumulation (dependency-aware).
Tuple outputs are handled with a dedicated tuple-weight map.
Results are cached (_UB_CACHE, _LB_CACHE) for repeated queries.

`TupleExpr` and multi-output flows#

Some traced operations naturally produce multiple outputs. BoundLab supports this via TupleExpr-style nodes (for example MakeTuple and GetTupleItem in boundlab.expr._tuple).

Key points:

TupleExpr.shape is a tuple of torch.Size objects.
TupleExpr.backward(*weights, direction=...) receives one weight per output slot.
GetTupleItem routes propagation to the parent tuple expression and selected index.

Linearizers: how nonlinear ops become tractable#

Nonlinear functions are handled by linearizers that return a zonotope relaxation (ZonoBounds).

A linearizer provides:

bias: constant term.
input_weights: slope-like weights applied to each input expression.
error_coeffs: coefficients for fresh error symbols.

Conceptually:

\[y \approx \sum_i w_i \odot x_i + b + E\eta\]

where η is a fresh bounded noise symbol.

Use @zono._register_linearizer("op_name") on a function returning ZonoBounds. The decorator installs a dispatcher handler that applies input weights, adds bias, and introduces a fresh LpEpsilon for approximation error.

Limitations and Notes#

BoundLab is currently an early-stage (0.1.0) project; interfaces may evolve.
Unsupported operators in a model trace will raise dispatch errors (typically KeyError in interpreter dispatch).
softmax handler currently supports 2D inputs along the last dimension.
reciprocal relaxation assumes positive-domain inputs and clamps lower bounds for numerical safety.
Some tuple APIs are currently internal (boundlab.expr._tuple) and may evolve.
For repeated experiments in one process, you may want to clear propagation caches:

import boundlab.prop as prop
prop._UB_CACHE.clear()
prop._LB_CACHE.clear()

Extending BoundLab#

To add a new nonlinear operator in the zonotope domain:

Implement a linearizer returning zono.ZonoBounds.
Register it with @zono._register_linearizer("op_name").
Add dispatcher entries for both function and module forms if needed.
Add soundness tests (sample-based checks are a good baseline).

Differential Verification (`boundlab.diff.zono3`)#

Standard zonotope verification bounds each network independently and then subtracts the interval, which can be very loose when two networks are similar. boundlab.diff.zono3 tracks the triple (x, y, diff) jointly, exploiting shared input noise and cancellations in affine layers to get tighter bounds on f₁(x) − f₂(x).

Quick example#

import torch
from torch import nn

import boundlab.expr as expr
from boundlab.diff.expr import DiffExpr3
from boundlab.diff.op import DiffLinear
from boundlab.diff.zono3 import interpret as diff_interpret

fc1, fc2 = nn.Linear(4, 8), nn.Linear(4, 8)
model = nn.Sequential(DiffLinear(fc1, fc2), nn.ReLU(), nn.Linear(8, 3))
exported = torch.export.export(model, (torch.randn(4),))
op = diff_interpret(exported)

c = torch.randn(4)
x = expr.ConstVal(c) + 0.1 * expr.LpEpsilon([4])
out = op(x, x)   # returns DiffExpr3 after the first ReLU

d_ub, d_lb = out.diff.ublb()
print("max |diff|:", d_ub.abs().max().item())

Key types#

Type	Purpose
`DiffExpr2(x, y)`	Pair of branches; produced by `diff_pair_handler`.
`DiffExpr3(x, y, diff)`	Triple with explicit difference expression; output of nonlinear handlers.
`DiffLinear(fc1, fc2)`	`nn.Module` that pairs two linear layers via `diff_pair`.

Input options for `diff_interpret`#

Pass a DiffExpr3 directly when you have two separate zonotopes.
Pass DiffLinear in the model for a single shared-input scenario.
Pass a plain Expr to use diff_interpret as a drop-in replacement for zono.interpret (identical results, no overhead).

For a full explanation, see Differential Verification with diff.zono3.

Where To Go Next#

See Examples for complete scripts.
Inspect API Reference for full signatures and class docs.
Read tests/ for additional usage patterns and soundness checks.