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.