The Abstract Interpreter#
The idea#
BoundLab does not re-implement PyTorch semantics. A model is exported to
ONNX once, and the Interpreter replays the graph node by node over a mixed
environment of concrete tensors and abstract expressions. Each ONNX
operator dispatches over a collection of OpHandlers: a handler declares
which operator it implements (op) and when it applies (condition,
usually a check on the operand classes), and exactly one handler must be
ready for every call. Handlers call each other through the interpreter
(interp.relu(x), interp.matmul(x, y)), so a composite handler like
softmax automatically uses whatever relaxations the enclosing domain
registered.
The base interpreter#
boundlab.interp.base.interpret bundles the operators every domain shares.
They are implemented purely against the Expr primitives — einsum, the
shape operators, same-class add — so they are exact for any component
class and never need a domain-specific version: add/sub/neg,
mul/div/matmul with at least one concrete operand, gemm,
reshape/transpose/broadcast_to, reduce_sum/reduce_mean,
unsqueeze/squeeze, identity, and cast.
What the base interpreter deliberately does not contain is anything that loses precision: activations and products of two abstract operands are the domain’s decisions.
How a domain assembles its interpret#
A domain is one Interpreter(...) call: start from the base handlers, add
the domain’s non-linear handlers, and optionally normalize every node result
with after_each. This is the zonotope domain’s actual assembly:
from boundlab.ibp.matmul import MatmulBiased
from boundlab.ibp.softmax import Softmax2ExpReciprocal
from boundlab.interp import Interpreter, base
from boundlab.zono import keep_zono, linearizers, matmul
interpret = Interpreter(
base.interpret, # shared exact operators
MatmulBiased(), # peel Bias centers off matmul operands
matmul.Matmul(), # zonotope x zonotope estimate
linearizers.MaxWithConst2Relu(), # max(x, c) -> relu(x - c) + c
linearizers.Relu(), # affine enclosures with Noise residual
linearizers.Exp(),
linearizers.Tanh(),
linearizers.Reciprocal(),
Softmax2ExpReciprocal(), # softmax -> exp / sum / reciprocal
after_each=keep_zono, # fresh Noise -> new error symbols
)
For one dispatched call, say matmul(x, y), every handler with
op == "matmul" is asked for its condition: MatmulSimple (from the base
set) accepts when one operand is concrete, MatmulBiased when both carry a
Bias component, Matmul when both are pure zonotopes — the conditions
partition the input space, and the single ready handler runs. When two
handlers legitimately overlap, overrides declares the winner.
Using an interpreter#
Calling the interpreter on a model exports it (once) and returns an evaluator over abstract inputs:
import torch
from torch import nn
from boundlab import Error
from boundlab.utils import ShapeDtype
from boundlab.zono import Zono, interpret
center = torch.zeros(4)
err = Error("input", ShapeDtype(center.shape, center.dtype))
x = Zono.error(err) * 0.1 + center
model = nn.Sequential(nn.Linear(4, 4), nn.ReLU()).eval()
lb, ub = interpret(model)(x).lbub()
The same dispatch is available directly by attribute — useful for building values by hand or testing a single relaxation:
y = interpret.relu(x) # dispatch the "relu" operator
z = interpret.matmul(y, torch.randn(4, 2))
A custom operator is a small class (or a handler_fn-decorated function)
added to a new Interpreter — existing handlers are inherited by passing
the old interpreter first, and override_handler declares precedence over
the handler being shadowed:
from boundlab.interp import Interpreter, base, handler_fn
@handler_fn("neg")
def LoudNeg(value, **params):
"""Negation that announces itself."""
print("negating!")
return -value
mine = Interpreter(
interpret,
LoudNeg().override_handler(base.Neg), # wins over the base neg handler
after_each=interpret.after_each,
)
_ = mine.neg(x)
interpret(model, verbose=True) prints, per node, which handlers fired and
a one-line diagnostic of the result — the first tool to reach for when a
bound is wider than expected.
Export#
onnx_export(function, args) wraps torch.export + torch.onnx.export
and returns the in-memory ONNX IR the interpreter consumes; paths and
ONNXPrograms are accepted by the interpreter directly.
InterpreterModule packages a bound graph as an nn.Module returning
(lb, ub), so a whole verification pipeline can itself be exported or
AOT-compiled.