diff --git a/.github/workflows/act-bab.yml b/.github/workflows/act-bab.yml index c880260d0..d19fad8d9 100644 --- a/.github/workflows/act-bab.yml +++ b/.github/workflows/act-bab.yml @@ -72,7 +72,7 @@ jobs: - name: VNNLIB — download benchmarks run: | cd ${{ github.workspace }} - for cat in acasxu_2023 dist_shift_2023 safenlp_2024 relusplitter_2026 isomorphic_acasxu_2026 smart_turn_multimodal_2026; do + for cat in acasxu_2023 dist_shift_2023 safenlp_2024 relusplitter_2026 isomorphic_acasxu_2026; do coverage run -p -m act.front_end --download $cat done @@ -98,17 +98,6 @@ jobs: coverage run -p -m act.pipeline --verify vnnlib --category relusplitter_2026 --max-instances 2 --merge-split-relus \ --solvers dual --device cpu --dtype float64 - # =================================================================== - # ONNX QuantizeLinear/DequantizeLinear ingestion: smart_turn is the only - # benchmark with QDQ nodes, exercising the onnx_converter Q/DQ converters - # + initializer folding. - # =================================================================== - - name: VNNLIB E2E — smart_turn_multimodal_2026 (ONNX Q/DQ ingestion) - run: | - cd ${{ github.workspace }} - timeout 30 coverage run -p -m act.pipeline --verify vnnlib --category smart_turn_multimodal_2026 \ - --max-instances 1 --solvers dual --device cpu --dtype float64 --ignore-errors || true - # =================================================================== # Isomorphic dual-network specs (2026): _isomorphic_multinet_rewrite, # _CombinedDualModel loading, multi-row LINEAR_LE conjunction encode. @@ -143,8 +132,8 @@ jobs: --device cpu --dtype float64 # =================================================================== - # Gain branching + joint multi-neuron splitting: _multi_split_from_decision - # / _multi_split_from_groups and the gain-tier bab loop (2^k child cubes). + # NEURON SPLIT and MULTI-SPLIT invariants are asserted in node.py and bab.py. + # Exit status 0 means every lane received a verdict; stdout is not validation. # =================================================================== - name: BaB gain branching + multi-split run: | @@ -153,71 +142,46 @@ jobs: --solver dual --bab --bab-branching gain --bab-multi-split-levels 3 \ --bab-max-depth 4 --bab-max-subproblems 40 --device cpu --dtype float64 - # Joint multi-neuron splitting is orthogonal to the branching method: the - # k neurons come from the BaBSR heuristic either way, so this must work - # under --bab-branching babsr. The greps make the job non-vacuous — the - # telemetry counters are emitted only when the joint path actually runs, - # and k_requested proves the explicit flag survives the --solver dual - # preset instead of being silently forced to 4. + # Joint-split depth, counters, and child cubes are asserted in bab.py and node.py. + # Exit status 0 means every lane received a verdict; stdout is not validation. - name: BaB joint multi-split — decoupled from gain (babsr + dual_alpha_eta) run: | cd ${{ github.workspace }} - out=$(coverage run -p -m act.back_end --verify \ + coverage run -p -m act.back_end --verify \ --network "$ACT_NETS_DIR/mlp_plain_3x8_64x64_3962224133.json" \ --solver dual --bab --bab-branching babsr --bab-solver-tier dual_alpha_eta \ --bab-multi-split-levels 3 --bab-max-batch-size 256 \ --bab-max-depth 6 --bab-max-subproblems 60 \ - --device cpu --dtype float64 --verbose 2>&1) - echo "$out" - grep -qE "multi_split_k_requested: 3" <<<"$out" - grep -qE "multi_split_wave_count: [1-9]" <<<"$out" - grep -qE "multi_split_k_used: [2-9]" <<<"$out" - - # --bab-multi-split-levels 1 is load-bearing: when joint splitting fires it - # replaces the brancher entirely, so compute_scores is never called and this - # step would silently exercise nothing. wave_count: 0 asserts that. + --device cpu --dtype float64 --verbose + + # Witness decision and neuron-split invariants are asserted in branching.py and node.py. + # Exit status 0 means every lane received a verdict; stdout is not validation. - name: BaB branching — witness residual (dual_alpha_eta) run: | cd ${{ github.workspace }} - out=$(coverage run -p -m act.back_end --verify \ + coverage run -p -m act.back_end --verify \ --network "$ACT_NETS_DIR/mlp_plain_3x8_64x64_3962224133.json" \ --solver dual --bab --bab-branching witness_residual \ --bab-solver-tier dual_alpha_eta --bab-multi-split-levels 1 \ --bab-max-batch-size 64 --bab-max-depth 5 --bab-max-subproblems 30 \ - --device cpu --dtype float64 --verbose 2>&1) - echo "$out" - grep -qE "multi_split_wave_count: 0" <<<"$out" - grep -qE "witness_residual_fallback_count: 0" <<<"$out" - grep -qE "^ nodes: [1-9]" <<<"$out" - - # top-k is a real parameter on the four TopKBounding-family strategies and - # a hard error on the other two. Both halves are asserted. + --device cpu --dtype float64 --verbose + + # Top-k cap and pool-size invariants are asserted in branching/bounding.py. + # Exit status 0 means every lane received a verdict. - name: BaB bounding — top-k parameter run: | cd ${{ github.workspace }} for v in depth_bound_blend greedy annealed diverse_split_signs; do - out=$(coverage run -p -m act.back_end --verify \ - --network "$ACT_NETS_DIR/mlp_plain_3x8_64x64_3962224133.json" \ - --bab --bab-bounding $v --bab-top-k 2 \ - --solver torchlp --device cpu --dtype float64 \ - --timeout 30 --bab-max-depth 4 --bab-max-subproblems 40 --verbose 2>&1) - echo "$out" - grep -qE "^ nodes: [1-9]" <<<"$out" - grep -qE "bounding_top_k_effective: 2" <<<"$out" - done - for v in random mcts; do - if coverage run -p -m act.back_end --verify \ + coverage run -p -m act.back_end --verify \ --network "$ACT_NETS_DIR/mlp_plain_3x8_64x64_3962224133.json" \ --bab --bab-bounding $v --bab-top-k 2 \ --solver torchlp --device cpu --dtype float64 \ - --timeout 20 --bab-max-depth 2 --bab-max-subproblems 5 >/dev/null 2>&1; then - echo "expected --bab-top-k to be rejected for $v"; exit 1 - fi + --timeout 30 --bab-max-depth 4 --bab-max-subproblems 40 --verbose done # =================================================================== # SST/Yelp text-method selection: select_bert_method + BackendConfig - # text-method resolution (planar -> baf/alpha_mode/solver_tier). + # text-method resolution (planar -> solver_tier). # =================================================================== - name: Text-method config resolution (planar) run: | @@ -349,10 +313,8 @@ jobs: --timeout 30 --bab-max-depth 8 --bab-max-subproblems 300 done - # Greedy / annealed / diverse_split_signs / MCTS pool-selection strategies end-to-end through - # the back-end CLI. mlp_plain_3x8 survives LP presolve and enters BaB; the - # greps reject a vacuous run (nodes_minted is emitted only by the BaB wave - # loop). stderr is folded in because the mcts wave summary is a log record. + # Pool transitions and MCTS visit-tree invariants are asserted in branching/bounding.py. + # Exit status 0 means every lane received a verdict; stdout is not validation. - name: BaB bounding — greedy / SA / diverse / MCTS (ACT net) run: | cd ${{ github.workspace }} @@ -361,16 +323,10 @@ jobs: "--bab-bounding annealed --bab-sa-cooling-rate 0.95" \ "--bab-bounding diverse_split_signs" \ "--bab-bounding mcts --bab-solver-tier dual_alpha --bab-mcts-exploration 1.4 --bab-mcts-lambda 0.5 --bab-mcts-virtual-loss 1.0"; do - out=$(coverage run -p -m act.back_end --verify \ + coverage run -p -m act.back_end --verify \ --network "$ACT_NETS_DIR/mlp_plain_3x8_64x64_3962224133.json" --bab $opts \ --solver torchlp --device cpu --dtype float64 \ - --timeout 30 --bab-max-depth 3 --bab-max-subproblems 10 --verbose 2>&1) - echo "$out" - grep -qE "^ nodes: [1-9]" <<<"$out" - grep -q "nodes_minted:" <<<"$out" - case "$opts" in - *mcts*) grep -qE "mcts: nodes=[1-9].*n_tot=[1-9].*histogram=\{" <<<"$out" ;; - esac + --timeout 30 --bab-max-depth 3 --bab-max-subproblems 10 --verbose done - name: BaB split + bounding — N-ary split + SA order (VNNLIB) @@ -391,6 +347,31 @@ jobs: --bab-bounding depth_bound_blend --bab-provenance \ --bab-max-depth 3 --bab-max-nodes 10 --device cpu --dtype float64 + # Correctness invariants are asserted in act/back_end/bab/climb.py; a + # violation raises "CLIMB invariant violated: ..." at the failing line. + # Exit status 0 means every lane received a verdict. The generated net's + # weights change whenever nets are regenerated; 250 subproblems keep core + # learning and propagation exercised across regenerations (17-75 cores). + - name: BaB CLIMB — certificate replay and propagation + run: | + cd ${{ github.workspace }} + coverage run -p -m act.back_end --verify \ + --network "$ACT_NETS_DIR/mlp_plain_3x8_64x64_3962224133.json" \ + --solver dual --bab --bab-preset climb --bab-max-batch-size 16 \ + --bab-max-depth 14 --bab-max-subproblems 250 --timeout 600 --verbose + + # These pool nets exercise symbolic MaxPool forward and cached one-hot dual + # backward. Exit status 0 means every lane received a verdict (CERTIFIED / + # FALSIFIED / UNKNOWN / TIMEOUT); source assertions catch invariant violations. + - name: Dual MaxPool — symbolic forward and one-hot backward + run: | + cd ${{ github.workspace }} + for net in layer_testing_cnn_pool cnn2d_plain_4x1x8x8_32x8x16_4023708073; do + coverage run -p -m act.back_end --verify \ + --network "$ACT_NETS_DIR/$net.json" --solver dual --bab \ + --bab-branching babsr --bab-multi-split-levels 1 --verbose + done + # ── Coverage ─────────────────────────────────────────────────────── - name: Combine coverage data if: always() diff --git a/act/README.md b/act/README.md index 65afc9e24..79bb40f2d 100644 --- a/act/README.md +++ b/act/README.md @@ -67,8 +67,12 @@ act/ │ ├── layer_util.py # Layer validation and creation utilities │ ├── bab/ # Branch-and-bound refinement package │ │ ├── bab.py # BaB engine: verify_bab(), verify_bab_batched() -│ │ ├── node.py # BaB tree node representation +│ │ ├── climb.py # CLIMB certificate replay, core learning, propagation +│ │ ├── violation.py # Concrete counterexample checks for BaB candidates +│ │ ├── node.py # Subproblem batches and all child construction │ │ └── branching/ # Branching and bounding strategies +│ │ ├── branching.py # Split decisions and branching heuristics +│ │ └── multi_split.py # Multi-neuron selection and gain testing │ ├── utils.py # Backend utilities (affine_bounds, validate_constraints) │ ├── analyze.py # Network analysis and bounds propagation │ ├── cons_exportor.py # Constraint export to solvers @@ -77,7 +81,7 @@ act/ │ │ ├── solver_base.py # Base solver interface │ │ ├── solver_gurobi.py # Gurobi MILP solver integration │ │ ├── solver_torchlp.py # PyTorch-based LP solver -│ │ ├── solver_dual.py # Dual certified bounds solver +│ │ ├── solver_dual.py # Dual certified bounds solver with solve_spec_batch() │ │ └── solver_hz.py # HybridZ-based solver │ ├── interval_tf/ # Interval-based transfer functions │ │ ├── interval_tf.py # Interval TF implementation @@ -454,4 +458,4 @@ The back-end CLI provides comprehensive tools for working with ACT networks: 4. **Verify Networks**: `--verify` runs verification with optional `--bab` for branch-and-bound 5. **Test Serialization**: `--test-serialization` validates save/load functionality -All commands support unified `--device {cpu,cuda,gpu}` and `--dtype {float32,float64}` arguments. \ No newline at end of file +All commands support unified `--device {cpu,cuda,gpu}` and `--dtype {float32,float64}` arguments. diff --git a/act/back_end/README.md b/act/back_end/README.md index 219b69da7..fb13286f3 100644 --- a/act/back_end/README.md +++ b/act/back_end/README.md @@ -222,14 +222,18 @@ act/back_end/ verifier.py layer_schema.py bab/ - bab.py - node.py + bab.py # BaB engine wiring and orchestration + climb.py + violation.py + node.py # Subproblem batches and all child construction branching/ + branching.py # Split decisions and branching heuristics + multi_split.py # Multi-neuron selection and gain testing solver/ solver_base.py solver_gurobi.py solver_torchlp.py - solver_dual.py + solver_dual.py # DualSolver.solve_spec_batch() handles K-lane dual solves solver_hz.py interval_tf/ interval_tf.py @@ -286,7 +290,7 @@ Config (`BaBConfig`) / CLI flags: | `llm_probe_model` | `--bab-llm-probe-model` | `""` | | `llm_probe_base_url` | `--bab-llm-probe-base-url` | `""` | | `llm_probe_cadence` | `--bab-llm-probe-cadence` | `1` | -| `llm_probe_decisions` (e.g. `split,frontier,refine,neuron`) / `llm_probe_max_candidates_total` / `llm_probe_api_key_env` / `_temperature` / `_max_candidates` / `_history` / `_max_failures` / `_log` | (config/YAML only) | see `act/config/config.py` | +| `llm_probe_decisions` (e.g. `split,frontier,refine,neuron`) / `llm_probe_max_candidates_total` / `llm_probe_api_key_env` / `_temperature` / `_max_candidates` / `_history` / `_max_failures` / `_log` | `--bab-llm-probe-` (e.g. `--bab-llm-probe-decisions`); also YAML | see `act/config/config.py` | Implementation lives in `act/pipeline/verification/llm_probe.py`. That file has two independent sections: the legacy TinyLlama input/output probe (Section A) and this BaB controller (Section B); diff --git a/act/back_end/__init__.py b/act/back_end/__init__.py index 66063e050..079cb616b 100644 --- a/act/back_end/__init__.py +++ b/act/back_end/__init__.py @@ -124,9 +124,7 @@ "BackendConfig", "GenerationConfig", # BaB - "BabNode", "SubproblemBatch", - "split_subproblems", "BaBConfig", "verify_bab", "BranchingStrategy", diff --git a/act/back_end/bab/__init__.py b/act/back_end/bab/__init__.py index 5ce77660e..1295e991e 100644 --- a/act/back_end/bab/__init__.py +++ b/act/back_end/bab/__init__.py @@ -8,7 +8,7 @@ from act.config.config import BaBConfig from act.back_end.bab.bab import verify_bab, verify_bab_batched -from act.back_end.bab.node import BabNode, SubproblemBatch, split_subproblems +from act.back_end.bab.node import SubproblemBatch from act.back_end.bab.branching.branching import BranchingStrategy from act.back_end.bab.branching.bounding import BoundingStrategy @@ -16,9 +16,7 @@ "BaBConfig", "verify_bab", "verify_bab_batched", - "BabNode", "SubproblemBatch", - "split_subproblems", "BranchingStrategy", "BoundingStrategy", ] diff --git a/act/back_end/bab/bab.py b/act/back_end/bab/bab.py index 88a76261b..c2192a570 100644 --- a/act/back_end/bab/bab.py +++ b/act/back_end/bab/bab.py @@ -7,8 +7,11 @@ # ===---------------------------------------------------------------------====# # # Purpose: -# BaB loop on a single-spec instance. Subproblems explored in K-batched -# waves via solve_batch; CE validation per SAT lane. Solver-agnostic. +# BaB loop on a single-spec instance. Subproblems are explored in K-batched +# waves via solve_batch with CE validation per SAT lane; optional certificate +# reuse is delegated to ``act.back_end.bab.climb.ClimbSession``. Dual-tier +# bound policy lives here, concrete CE checks in ``violation``, branching +# decisions in ``branching/``, and child construction in ``node.py``. # # ===---------------------------------------------------------------------====# @@ -16,17 +19,15 @@ import logging import math -import sys import time -import inspect -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Dict, Final, List, Optional, Tuple, Union, cast +from typing import Any, Callable, Dict, List, Optional, Union, cast import torch from act.config.config import ( BaBConfig, DualConfig, + NEURON_BRANCHING_METHODS, TOP_K_BOUNDINGS, TOP_K_INCOMPATIBLE_BOUNDINGS, VALID_BOUNDINGS, @@ -34,20 +35,25 @@ ) from act.back_end.bab.node import ( SubproblemBatch, - concat_children, - rederive_embedding_block_eps, - split_input, - split_input_nary, - split_neuron_subproblems, + _install_embedding_child_block_eps, + _restore_embedding_child_block_eps, + slice_branching_state, + split_neurons, ) +from act.back_end.bab.climb import ClimbSession from act.back_end.bab.branching.branching import ( BranchingStrategy, SplitDecision, + _assert_witness_residual_decision_unstable, _build_branching_strategy as _build_branching_strategy_impl, - _collect_neuron_candidates, + witness_relu_preactivations, +) +from act.back_end.bab.branching.multi_split import ( _multi_split_from_decision, - _multi_split_from_groups, enumerate_unstable_candidates, + gain_tested_decision, + groups_to_tensors, + presplit_root, ) from act.back_end.bab.branching.bounding import ( BoundingStrategy, @@ -61,8 +67,16 @@ SAOrder, ) -from act.back_end.core import Bounds, Layer, Net, ParamValue -from act.back_end.solver.solver_base import BatchLPSolution, Solver, SolveStatus +from act.back_end.core import Bounds, Layer, Net +from act.back_end.dual_tf.tf_forward import compute_forward_bounds +from act.back_end.layer_schema import LayerKind +from act.back_end.bab.violation import _check_input_specs_batched, check_violations_batched +from act.back_end.solver.solver_base import Solver, SolveStatus +from act.back_end.solver.solver_dual import ( + DualBatchResult, + DualSolver, + expand_bounds_dict, +) from act.back_end.verifier import ( gather_input_spec_layers, get_assert_layer, @@ -71,812 +85,11 @@ setup_and_solve_batch, ) from act.front_end.specs import OutKind, OutputSpec -from act.front_end.specs import InKind, normalize_position_mask -from act.util.model_inference import infer_single_model from act.util.stats import VerifyStatus, VerifyResult -if TYPE_CHECKING: - from act.back_end.interval_tf.tf_attention import LinearBounds - log = logging.getLogger(__name__) -@dataclass -class DualSolveResult: - solution: BatchLPSolution - bounds_dict: Optional[Dict[int, Bounds]] = None - nu_per_layer: Optional[Dict[int, torch.Tensor]] = None - witness_input: Optional[torch.Tensor] = None - row_slack: Optional[torch.Tensor] = None - """Per-spec-row slack ``[K, m]``; ``slack >= 0`` means the row is certified - (ALL-rows kinds). Consumed by the root spec-pruning presolve.""" - - -def _select_spec_rows( - state: Optional[Dict[int, torch.Tensor]], - keep_rows: torch.Tensor, -) -> Optional[Dict[int, torch.Tensor]]: - """Slice the spec axis (dim 1 of ``[N, M, n]``) of per-layer dual state.""" - if state is None: - return None - return { - lid: t.index_select(1, keep_rows.to(t.device)) if t.dim() >= 3 else t - for lid, t in state.items() - } - - -def _presplit_root( - root: SubproblemBatch, - bounds_dict: Dict[int, Bounds], - nu_per_layer: Dict[int, torch.Tensor], - k: int, -) -> Optional[SubproblemBatch]: - """Materialize the 2^k descendants of the root's top-k scored neurons. - - Score = triangle relaxation area x |nu| (BaBSR essence); the layer with - the strongest top score wins. The 2^k sign assignments exactly partition - the root region (each unstable neuron is either >=0 or <=0), so replacing - the root by these children is sound. - """ - best: Optional[tuple[int, torch.Tensor, torch.Tensor]] = None - for lid, nu in nu_per_layer.items(): - b = bounds_dict.get(lid) - if b is None: - continue - lb = b.lb.flatten(start_dim=1)[0] - ub = b.ub.flatten(start_dim=1)[0] - n = min(lb.shape[-1], nu.shape[-1]) - lb, ub = lb[:n], ub[:n] - amb = (lb < 0) & (ub > 0) - if not bool(amb.any().item()): - continue - area = (-lb * ub / (ub - lb).clamp(min=1e-12)).clamp(min=0) - score = area * nu.reshape(-1, nu.shape[-1])[:, :n].abs().sum(dim=0) - score = torch.where(amb, score, torch.zeros_like(score)) - if best is None or float(score.max()) > float(best[1].max()): - best = (lid, score, lb) - if best is None: - return None - lid, score, _ = best - k = min(k, int((score > 0).sum().item())) - if k < 1: - return None - top_idx = torch.topk(score, k=k).indices - n_children = 2 ** k - n_layer = score.shape[-1] - m = root.incremental_alpha[next(iter(root.incremental_alpha))].shape[1] if root.incremental_alpha else 1 - - signs = torch.zeros(n_children, m, n_layer, device=root.lb.device, dtype=root.lb.dtype) - for j in range(n_children): - for bit, neuron in enumerate(top_idx.tolist()): - signs[j, :, neuron] = 1.0 if (j >> bit) & 1 else -1.0 - - def _rep(state: Optional[Dict[int, torch.Tensor]]) -> Optional[Dict[int, torch.Tensor]]: - if state is None: - return None - return {l: t.repeat(n_children, *([1] * (t.dim() - 1))) for l, t in state.items()} - - merged_signs = _rep(root.split_signs) or {} - if lid in merged_signs: - merged_signs[lid] = merged_signs[lid] + signs - else: - merged_signs[lid] = signs - return SubproblemBatch( - lb=root.lb.repeat(n_children, 1), - ub=root.ub.repeat(n_children, 1), - depths=torch.full((n_children,), k, dtype=torch.long, device=root.lb.device), - incremental_alpha=_rep(root.incremental_alpha), - incremental_eta=_rep(root.incremental_eta), - split_signs=merged_signs, - ) - - -def _interval_refresh_bounds( - net: Net, - base: Dict[int, Bounds], - split_signs: Dict[int, torch.Tensor], -) -> Optional[Dict[int, Bounds]]: - """Cheap batched IBP re-propagation of split phases, intersected with base. - - Frozen root bounds lose the downstream effect of hardened splits; a plain - interval pass (no LinearBound A matrices, milliseconds per batch) restores - it. Every entry is intersected with the base bounds, so the result can - only tighten valid over-approximations (sound). Returns None when an - unsupported layer kind is encountered - the caller keeps the base dict. - """ - from act.back_end.dual_tf.tf_forward import _fwd_conv2d_interval - - out = dict(base) - vals: Dict[int, tuple[torch.Tensor, torch.Tensor]] = {} - for layer in net.layers: - k = layer.kind.upper() if isinstance(layer.kind, str) else layer.kind - lid = layer.id - if k == "ASSERT": - continue - if k in ("INPUT", "INPUT_SPEC"): - b = out.get(lid) - if b is None: - return None - vals[lid] = (b.lb.flatten(start_dim=1), b.ub.flatten(start_dim=1)) - continue - preds = net.preds.get(lid, []) - try: - if k == "CONV2D": - plb, pub = vals[preds[0]] - lb, ub = _fwd_conv2d_interval(layer, plb, pub) - lb, ub = lb.flatten(start_dim=1), ub.flatten(start_dim=1) - elif k == "DENSE": - w = layer.params["weight"] - bias = layer.params.get("bias") - if not isinstance(w, torch.Tensor): - return None - plb, pub = vals[preds[0]] - w_pos, w_neg = w.clamp(min=0), w.clamp(max=0) - lb = plb @ w_pos.T + pub @ w_neg.T - ub = pub @ w_pos.T + plb @ w_neg.T - if isinstance(bias, torch.Tensor): - lb, ub = lb + bias, ub + bias - elif k == "ADD": - (alb, aub), (blb, bub) = vals[preds[0]], vals[preds[1]] - lb, ub = alb + blb, aub + bub - elif k in ("FLATTEN", "RESHAPE"): - lb, ub = vals[preds[0]] - elif k == "RELU": - lb, ub = vals[preds[0]] - else: - return None - except (KeyError, IndexError, ValueError): - return None - - b = out.get(lid) - if b is not None: - lb = torch.maximum(lb, b.lb.flatten(start_dim=1)) - ub = torch.minimum(ub, b.ub.flatten(start_dim=1)) - ub = torch.maximum(ub, lb) - if k == "RELU": - s = split_signs.get(lid) - if s is not None: - sl = s[:, 0, :] if s.dim() == 3 else s - n = min(lb.shape[-1], sl.shape[-1]) - sn = sl[..., :n].to(lb.device) - lb, ub = lb.clone(), ub.clone() - lb[..., :n] = torch.where(sn > 0, lb[..., :n].clamp(min=0.0), lb[..., :n]) - ub[..., :n] = torch.where(sn < 0, ub[..., :n].clamp(max=0.0), ub[..., :n]) - ub[..., :n] = torch.maximum(ub[..., :n], lb[..., :n]) - if b is not None: - out[lid] = Bounds(lb.view_as(b.lb).clone(), ub.view_as(b.ub).clone()) - if k == "RELU": - vals[lid] = (lb.clamp(min=0.0), ub.clamp(min=0.0)) - else: - vals[lid] = (lb, ub) - return out - - -def _neuron_branching_supported(config: BaBConfig) -> bool: - return ( - getattr(config, "branching_method", "random") in ("babsr", "fsb", "gain", "witness_residual") - and getattr(config, "solver_tier", "lp") in ("dual_alpha", "dual_alpha_eta") - ) - - -def _witness_residual_branching_active(config: BaBConfig) -> bool: - return getattr(config, "branching_method", "random") == "witness_residual" - - -def _witness_relu_preactivations( - net: Net, - witness_input: torch.Tensor, - input_shape: tuple[int, ...], -) -> Optional[Dict[int, torch.Tensor]]: - """Concrete per-ReLU pre-activations at the dual solver's spurious CE. - - ``witness_residual`` scores each neuron by the relaxation gap realized at - that one point, so it needs ``z`` there; ``bounds_dict`` only carries - ``[l, u]`` over the whole sub-box. A concrete forward is the degenerate case - of the interval one, so ``compute_forward_bounds`` on ``lb == ub == witness`` - supplies it — already keyed by layer id, already storing ReLU boxes - pre-activation. - """ - from act.back_end.dual_tf.tf_forward import compute_forward_bounds - - if witness_input.numel() == 0: - return None - x = witness_input - if input_shape and x.dim() == 2 and x.shape[1] == int(math.prod(input_shape)): - x = x.reshape(x.shape[0], *input_shape) - try: - bounds_dict = compute_forward_bounds(net, x, x) - except (ValueError, RuntimeError, KeyError, IndexError): - # Unregistered layer kind, or a shape/device mismatch in the witness. - # Both are recoverable: the caller falls back to BaBSR scoring. - return None - preacts = { - lid: bounds.lb.flatten(start_dim=1) - for lid, bounds in bounds_dict.items() - if _layer_kind_upper(net.by_id[lid]) == "RELU" - } - return preacts or None - - -def _layer_kind_upper(layer: Layer) -> str: - kind = layer.kind - return kind.upper() if isinstance(kind, str) else str(kind) - - -def _gain_tested_decision( - branch_batch: SubproblemBatch, - net: Net, - assert_layer: Layer, - config: BaBConfig, - dual_config: DualConfig, - keep_rows: Optional[torch.Tensor], - root_bounds_dict: Optional[Dict[int, Bounds]], - bounds_dict: Optional[Dict[int, Bounds]], - nu_per_layer: Optional[Dict[int, torch.Tensor]], - input_shape: tuple[int, ...], - n_candidates: int = 3, -) -> Optional[SplitDecision]: - """Pick each lane's split by measured child bounds, not by score proxy. - - BaBSR scores can rank a regression (-0.07) above the true best split - (+0.07); evaluating the top candidates' actual children with one cheap - non-optimized dual batch restores monotone progress (kfsb-style). - """ - if bounds_dict is None or nu_per_layer is None: - return None - kb = branch_batch.batch_size - device = branch_batch.lb.device - - cand = _collect_neuron_candidates(branch_batch, bounds_dict, nu_per_layer) - if cand is None: - return None - all_scores, all_layers, all_neurons = cand - n_c = min(n_candidates, all_scores.shape[1]) - top = torch.topk(all_scores, k=n_c, dim=1).indices - top_layers = all_layers.gather(1, top) - top_neurons = all_neurons.gather(1, top) - - rep_idx = torch.arange(kb, device=device).repeat_interleave(2 * n_c) - - def _rep_state(state): - if state is None: - return None - return {l: t.index_select(0, rep_idx.to(t.device)) for l, t in state.items()} - - m_specs = 1 - if branch_batch.incremental_alpha: - m_specs = int(next(iter(branch_batch.incremental_alpha.values())).shape[1]) - signs = _rep_state(branch_batch.split_signs) or {} - for lid_val in torch.unique(top_layers).tolist(): - lid_int = int(lid_val) - layer = net.by_id[lid_int] - n_neurons = int(layer.out_vars[-1] - layer.out_vars[0] + 1) - if lid_int not in signs: - signs[lid_int] = torch.zeros( - 2 * n_c * kb, m_specs, n_neurons, device=device, dtype=branch_batch.lb.dtype, - ) - else: - signs[lid_int] = signs[lid_int].clone() - for lane in range(kb): - for c in range(n_c): - if int(top_layers[lane, c]) != lid_int: - continue - row = lane * 2 * n_c + 2 * c - neuron = int(top_neurons[lane, c]) - signs[lid_int][row, :, neuron] = 1.0 - signs[lid_int][row + 1, :, neuron] = -1.0 - - probe = SubproblemBatch( - lb=branch_batch.lb.index_select(0, rep_idx), - ub=branch_batch.ub.index_select(0, rep_idx), - depths=branch_batch.depths.index_select(0, rep_idx), - incremental_alpha=_rep_state(branch_batch.incremental_alpha), - incremental_eta=_rep_state(branch_batch.incremental_eta), - split_signs=signs, - ) - n_probe = probe.batch_size - probe_bounds = Bounds( - probe.lb.reshape(n_probe, *input_shape) if input_shape else probe.lb, - probe.ub.reshape(n_probe, *input_shape) if input_shape else probe.ub, - ) - res = _dispatch_dual_solve( - net=net, - assert_layer=assert_layer, - batched_bounds=probe_bounds, - k_actual=n_probe, - batch=probe, - config=config, - dual_config=dual_config, - optimize=False, - keep_rows=keep_rows, - root_bounds_dict=root_bounds_dict, - ) - child_lbs = (-res.solution.max_viol).view(kb, n_c, 2) - pair_gain = child_lbs.min(dim=2).values - best_c = pair_gain.argmax(dim=1) - lane_idx = torch.arange(kb, device=device) - return SplitDecision( - kind="neuron", - layer_id=top_layers[lane_idx, best_c], - neuron_idx=top_neurons[lane_idx, best_c], - ) - - -def _input_axis_decision_tensor( - decision: SplitDecision, - batch: SubproblemBatch, -) -> torch.Tensor: - if decision.input_axis is None: - raise ValueError("input-axis decision missing input_axis") - input_axis = torch.as_tensor(decision.input_axis, device=batch.lb.device, dtype=torch.long).reshape(-1) - if input_axis.numel() == 1: - input_axis = input_axis.expand(batch.batch_size) - if input_axis.numel() != batch.batch_size: - raise ValueError( - f"input-axis decision has {input_axis.numel()} lanes for batch size {batch.batch_size}" - ) - return input_axis.contiguous() - - -def _split_from_decision( - batch: SubproblemBatch, - decision: SplitDecision, - net: Net, -) -> tuple[SubproblemBatch, torch.Tensor]: - fanout = max(2, int(getattr(decision, "fanout", 2))) - if decision.kind == "input_axis": - dims = ( - _input_axis_decision_tensor(SplitDecision(kind="input_axis", input_axis=decision.cut_dim), batch) - if decision.cut_dim is not None - else _input_axis_decision_tensor(decision, batch) - ) - if fanout == 2: - return split_input(batch, dims) - return split_input_nary(batch, dims, fanout) - - if decision.kind == "neuron": - if decision.layer_id is None or decision.neuron_idx is None: - raise ValueError("neuron decision missing layer_id or neuron_idx") - - layer_id_tensor = decision.layer_id.reshape(-1) - neuron_idx_tensor = decision.neuron_idx.reshape(-1) - if layer_id_tensor.numel() == 0 or neuron_idx_tensor.numel() == 0: - raise ValueError("neuron decision tensors must be non-empty") - - rep_lid = int(layer_id_tensor[0].item()) - rep_idx = int(neuron_idx_tensor[0].item()) - if rep_lid < 0: - fallback_dims = (batch.ub - batch.lb).argmax(dim=-1) - if fanout == 2: - return split_input(batch, fallback_dims) - return split_input_nary(batch, fallback_dims, fanout) - - n_lanes = batch.batch_size - lids = layer_id_tensor.expand(n_lanes) if layer_id_tensor.numel() == 1 else layer_id_tensor - idxs = neuron_idx_tensor.expand(n_lanes) if neuron_idx_tensor.numel() == 1 else neuron_idx_tensor - if lids.numel() != n_lanes or idxs.numel() != n_lanes: - raise ValueError( - f"neuron decision has {lids.numel()}/{idxs.numel()} entries " - f"for batch size {n_lanes}" - ) - - # Per-lane split: lane i hardens ITS OWN (layer, neuron); collapsing - # to lane 0's choice makes the other K-1 lanes split an irrelevant - # neuron and stalls deep convergence. - device = batch.lb.device - parent_index = torch.arange(n_lanes, device=device).repeat(2) - - def _gather(state: Optional[Dict[int, torch.Tensor]]) -> Optional[Dict[int, torch.Tensor]]: - if state is None: - return None - return { - l: t.index_select(0, parent_index.to(t.device)) for l, t in state.items() - } - - m_specs = 1 - if batch.incremental_alpha: - m_specs = int(next(iter(batch.incremental_alpha.values())).shape[1]) - elif batch.split_signs: - m_specs = int(next(iter(batch.split_signs.values())).shape[1]) - - signs = _gather(batch.split_signs) or {} - for lid_val in torch.unique(lids).tolist(): - lid_int = int(lid_val) - layer = net.by_id[lid_int] - n_neurons = int(layer.out_vars[-1] - layer.out_vars[0] + 1) - if lid_int not in signs: - signs[lid_int] = torch.zeros( - 2 * n_lanes, m_specs, n_neurons, device=device, dtype=batch.lb.dtype, - ) - else: - signs[lid_int] = signs[lid_int].clone() - lane_sel = torch.where(lids == lid_val)[0] - neuron_sel = idxs[lane_sel].to(device=device, dtype=torch.long) - signs[lid_int][lane_sel, :, neuron_sel] = 1.0 - signs[lid_int][lane_sel + n_lanes, :, neuron_sel] = -1.0 - - children = SubproblemBatch( - lb=batch.lb.index_select(0, parent_index), - ub=batch.ub.index_select(0, parent_index), - depths=batch.depths.index_select(0, parent_index) + 1, - incremental_alpha=_gather(batch.incremental_alpha), - incremental_eta=_gather(batch.incremental_eta), - split_signs=signs, - parent_margins=( - batch.parent_margins.index_select(0, parent_index) - if batch.parent_margins is not None - else None - ), - lower_bound=( - batch.lower_bound.index_select(0, parent_index) - if batch.lower_bound is not None - else None - ), - ) - return children, parent_index - - raise ValueError(f"Unknown SplitDecision.kind: {decision.kind!r}") - - -def _slice_branching_state( - bounds_dict: Optional[Dict[int, Bounds]], - nu_per_layer: Optional[Dict[int, torch.Tensor]], - lane_idx: torch.Tensor, - k_actual: int, -) -> tuple[Optional[Dict[int, Bounds]], Optional[Dict[int, torch.Tensor]]]: - # ν/bounds are computed over the full k_actual wave; the brancher runs on the - # sub-batch actually being split. Bounds are [k_actual, *]; ν is [k_actual*M, n] - # packed sample-major (row b*M+j), so ν rows expand per selected lane. - bd_out: Optional[Dict[int, Bounds]] = None - if bounds_dict is not None: - bd_out = { - lid: Bounds( - b.lb.index_select(0, lane_idx.to(b.lb.device)), - b.ub.index_select(0, lane_idx.to(b.ub.device)), - ) - for lid, b in bounds_dict.items() - } - nu_out: Optional[Dict[int, torch.Tensor]] = None - if nu_per_layer is not None: - nu_out = {} - for lid, t in nu_per_layer.items(): - total = int(t.shape[0]) - if k_actual > 0 and total != k_actual and total % k_actual == 0: - m = total // k_actual - rows = ( - lane_idx.to(t.device).unsqueeze(1) * m - + torch.arange(m, device=t.device) - ).reshape(-1) - else: - rows = lane_idx.to(t.device) - nu_out[lid] = t.index_select(0, rows) - return bd_out, nu_out - - - -def _as_batched_vector( - value: object, - n_batch: int, - width: int, - device: torch.device, - dtype: torch.dtype, - name: str, -) -> torch.Tensor: - t = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) - t = t.to(device=device, dtype=dtype) - if t.dim() == 0: - return t.expand(n_batch, width).contiguous() - if t.dim() == 1: - if t.numel() == width: - return t.unsqueeze(0).expand(n_batch, -1).contiguous() - if width == 1 and t.numel() == n_batch: - return t.reshape(n_batch, 1).contiguous() - if t.dim() == 2: - if t.shape == (1, width): - return t.expand(n_batch, -1).contiguous() - if t.shape == (n_batch, width): - return t.contiguous() - raise ValueError( - f"{name}: expected scalar, ({width},), (1,{width}), or " - f"({n_batch},{width}); got {tuple(t.shape)}" - ) - - -def _as_batched_index( - value: object, - n_batch: int, - n_out: int, - device: torch.device, - name: str, -) -> torch.Tensor: - t = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) - t = t.to(device=device, dtype=torch.long).reshape(-1) - if t.numel() == 1: - t = t.expand(n_batch) - if t.numel() != n_batch: - raise ValueError( - f"{name}: expected 1 or {n_batch} indices, got {t.numel()}" - ) - if bool(((t < 0) | (t >= n_out)).any().item()): - raise ValueError(f"{name}: index out of range for n_out={n_out}: {t.tolist()}") - return t.contiguous() - - -# Per-net cache of reconstructed PyTorch nn.Module used for CE validation. -# Without this, every check_violations_batched call rebuilds the module from -# scratch via ACTToTorch.run() — costly under K-batched BaB which can invoke -# CE validation dozens of times for a single net. Cleared per top-level -# verify-all dispatch via clear_violation_check_module_cache(). -_VIOLATION_CHECK_MODULE_CACHE: dict[int, torch.nn.Module] = {} - - -def clear_violation_check_module_cache() -> None: - _VIOLATION_CHECK_MODULE_CACHE.clear() - - -def _module_float_dtype(module: torch.nn.Module, default: torch.dtype) -> torch.dtype: - for tensor in module.parameters(): - if tensor.is_floating_point(): - return tensor.dtype - for tensor in module.buffers(): - if tensor.is_floating_point(): - return tensor.dtype - return default - - -def _forward_for_violation_check(net: object, x_batch: torch.Tensor) -> torch.Tensor: - if isinstance(net, torch.nn.Module): - module = net - else: - key = id(net) - cached = _VIOLATION_CHECK_MODULE_CACHE.get(key) - if cached is None: - from act.pipeline.verification.act2torch import ACTToTorch - - cached = ACTToTorch(cast(Net, net)).run() - # ACTToTorch emits mixed float32 weights + float64 buffers; unify to - # the analysis dtype or the internal forward clashes float32/float64. - if x_batch.is_floating_point(): - cached = cached.to(dtype=x_batch.dtype) - _VIOLATION_CHECK_MODULE_CACHE[key] = cached - module = cached - _ = module.eval() - target_dtype = _module_float_dtype(module, x_batch.dtype) - if x_batch.dtype != target_dtype: - x_batch = x_batch.to(dtype=target_dtype) - success, output, error = infer_single_model("ce_validate_batched", module, x_batch) - if not success or output is None: - raise RuntimeError(f"check_violations_batched: model forward failed: {error}") - if output.dim() < 2: - raise ValueError( - f"check_violations_batched: model output must be batched, got " - f"shape={tuple(output.shape)}" - ) - return output.reshape(output.shape[0], -1) - - -@torch.no_grad() -def check_violations_batched(net: object, x_batch: torch.Tensor, assert_layer: Layer) -> torch.Tensor: - """[BATCHED-API] Return a ``[N]`` bool tensor for concrete ASSERT violations. - - ``x_batch`` is always treated as a tensor-view batch ``[N, *input_shape]``; - N=1 is represented by a length-one leading dimension. ASSERT parameters are - read directly from ``assert_layer.params`` in their batch-native form. - """ - if x_batch.dim() < 2: - raise ValueError( - f"check_violations_batched: x_batch must be [N, *input_shape], " - f"got shape={tuple(x_batch.shape)}" - ) - y_batch = _forward_for_violation_check(net, x_batch) - n_batch = int(x_batch.shape[0]) - if int(y_batch.shape[0]) != n_batch: - raise ValueError( - f"check_violations_batched: output batch {int(y_batch.shape[0])} " - f"!= input batch {n_batch}" - ) - n_out = int(y_batch.shape[1]) - device = y_batch.device - dtype = y_batch.dtype - params = assert_layer.params - kind = params.get("kind") - eps = 1e-8 - - if kind == OutKind.TOP1_ROBUST: - y_true = _as_batched_index(params["y_true"], n_batch, n_out, device, "y_true") - y_true_scores = y_batch.gather(1, y_true.unsqueeze(1)).squeeze(1) - mask = torch.ones_like(y_batch, dtype=torch.bool) - _ = mask.scatter_(1, y_true.unsqueeze(1), False) - other_scores = y_batch.masked_fill(~mask, -float("inf")) - return (other_scores.max(dim=1).values - y_true_scores) >= 0 - - if kind == OutKind.MARGIN_ROBUST: - y_true = _as_batched_index(params["y_true"], n_batch, n_out, device, "y_true") - margin = _as_batched_vector( - params["margin"], n_batch, 1, device, dtype, "margin" - ).reshape(n_batch) - y_true_scores = y_batch.gather(1, y_true.unsqueeze(1)).squeeze(1) - mask = torch.ones_like(y_batch, dtype=torch.bool) - _ = mask.scatter_(1, y_true.unsqueeze(1), False) - other_scores = y_batch.masked_fill(~mask, -float("inf")) - # ``-margin``, not ``+margin``: ``encode_linear`` emits rows ``e_j - e_t`` - # with ``thresholds = -margin``, so a lane is certified iff - # ``max_j(z_j - z_t) < -margin``. The negative sign is deliberate. - return (other_scores.max(dim=1).values - y_true_scores) >= -margin - - if kind == OutKind.LINEAR_LE: - c_raw = params["c"] - c_t = (c_raw if isinstance(c_raw, torch.Tensor) else torch.as_tensor(c_raw)).to(device=device, dtype=dtype) - if c_t.dim() <= 1: - rows = 1 - elif c_t.dim() == 2: - if c_t.shape[1] != n_out: - raise ValueError(f"LINEAR_LE: c cols {c_t.shape[1]} != n_out {n_out}") - rows = int(c_t.shape[0]) - else: - rows = int(c_t.shape[1]) - if rows <= 1: - coeff = _as_batched_vector(params["c"], n_batch, n_out, device, dtype, "c") - bound = _as_batched_vector(params["d"], n_batch, 1, device, dtype, "d").reshape(n_batch) - return (coeff * y_batch).sum(dim=1) >= bound + eps - if c_t.dim() == 2: - c_view = c_t.unsqueeze(0).expand(n_batch, -1, -1) - else: - c_view = c_t if c_t.shape[0] == n_batch else c_t.expand(n_batch, -1, -1) - d_raw = params["d"] - d_t = (d_raw if isinstance(d_raw, torch.Tensor) else torch.as_tensor(d_raw)).to(device=device, dtype=dtype).flatten() - if d_t.numel() == rows: - d_view = d_t.unsqueeze(0).expand(n_batch, -1) - elif d_t.numel() == n_batch * rows: - d_view = d_t.reshape(n_batch, rows) - else: - raise ValueError(f"LINEAR_LE: d numel {d_t.numel()} incompatible with rows {rows}") - lhs = torch.einsum("bmo,bo->bm", c_view.contiguous(), y_batch) - return (lhs > d_view + eps).any(dim=1) - - if kind == OutKind.RANGE: - result = torch.zeros(n_batch, dtype=torch.bool, device=device) - lb_raw = params.get("lb") - ub_raw = params.get("ub") - if lb_raw is not None: - lb = _as_batched_vector(lb_raw, n_batch, n_out, device, dtype, "lb") - result = result | (y_batch < lb - eps).any(dim=1) - if ub_raw is not None: - ub = _as_batched_vector(ub_raw, n_batch, n_out, device, dtype, "ub") - result = result | (y_batch > ub + eps).any(dim=1) - return result - - if kind == OutKind.UNSAFE_LINEAR: - m_raw = params.get("M", 1) - if isinstance(m_raw, torch.Tensor): - m_rows = int(m_raw.item()) - elif isinstance(m_raw, int): - m_rows = m_raw - else: - raise ValueError(f"UNSAFE_LINEAR: M must be int, got {m_raw!r}") - c_raw = params.get("C", params.get("c")) - if c_raw is None: - raise ValueError("UNSAFE_LINEAR requires C or c params") - c_tensor = c_raw if isinstance(c_raw, torch.Tensor) else torch.as_tensor(c_raw) - c_tensor = c_tensor.to(device=device, dtype=dtype) - if c_tensor.dim() == 2: - if c_tensor.shape == (m_rows, n_out): - c_view = c_tensor.unsqueeze(0).expand(n_batch, -1, -1).contiguous() - elif c_tensor.shape == (n_batch * m_rows, n_out): - c_view = c_tensor.reshape(n_batch, m_rows, n_out).contiguous() - else: - raise ValueError( - f"UNSAFE_LINEAR: C shape {tuple(c_tensor.shape)} incompatible " - f"with N={n_batch}, M={m_rows}, n_out={n_out}" - ) - elif c_tensor.dim() == 3: - if c_tensor.shape == (1, m_rows, n_out): - c_view = c_tensor.expand(n_batch, -1, -1).contiguous() - elif c_tensor.shape == (n_batch, m_rows, n_out): - c_view = c_tensor.contiguous() - else: - raise ValueError( - f"UNSAFE_LINEAR: c shape {tuple(c_tensor.shape)} incompatible " - f"with N={n_batch}, M={m_rows}, n_out={n_out}" - ) - else: - raise ValueError(f"UNSAFE_LINEAR: unsupported C dim {c_tensor.dim()}") - d_raw = params.get("thresholds", params.get("d")) - if d_raw is None: - raise ValueError("UNSAFE_LINEAR requires thresholds or d params") - d_tensor = d_raw if isinstance(d_raw, torch.Tensor) else torch.as_tensor(d_raw) - d_tensor = d_tensor.to(device=device, dtype=dtype) - if d_tensor.dim() == 1 and d_tensor.numel() == m_rows: - d_view = d_tensor.unsqueeze(0).expand(n_batch, -1).contiguous() - elif d_tensor.shape == (1, m_rows): - d_view = d_tensor.expand(n_batch, -1).contiguous() - elif d_tensor.shape == (n_batch, m_rows): - d_view = d_tensor.contiguous() - else: - raise ValueError( - f"UNSAFE_LINEAR: d shape {tuple(d_tensor.shape)} incompatible " - f"with N={n_batch}, M={m_rows}" - ) - lhs = torch.einsum("bmo,bo->bm", c_view, y_batch) - return (lhs <= d_view + eps).all(dim=1) - - raise NotImplementedError(f"ASSERT kind not supported: {kind}") - - -def _check_input_specs_batched(x_batch: torch.Tensor, spec_layers: List[Layer]) -> torch.Tensor: - result = torch.ones(x_batch.shape[0], device=x_batch.device, dtype=torch.bool) - tol = 1e-7 - for layer in spec_layers: - kind = layer.params.get("kind") - if kind not in (InKind.BOX, InKind.LINF_BALL, InKind.LP_EMBEDDING): - continue - lb = layer.params.get("lb") - ub = layer.params.get("ub") - if isinstance(lb, torch.Tensor) and isinstance(ub, torch.Tensor): - lb_t = lb.to(device=x_batch.device, dtype=x_batch.dtype) - ub_t = ub.to(device=x_batch.device, dtype=x_batch.dtype) - result &= ((x_batch >= lb_t - tol) & (x_batch <= ub_t + tol)).flatten(start_dim=1).all(dim=1) - if kind != InKind.LP_EMBEDDING: - continue - center = layer.params.get("center") - eps = layer.params.get("eps") - p_norm = layer.params.get("p_norm") - if not isinstance(center, torch.Tensor) or eps is None or p_norm is None: - result &= torch.zeros_like(result) - continue - center_t = center.to(device=x_batch.device, dtype=x_batch.dtype) - if isinstance(eps, torch.Tensor): - eps_t = eps.to(device=x_batch.device, dtype=x_batch.dtype) - elif isinstance(eps, (int, float, bool)): - eps_t = center_t.new_tensor(float(eps)) - else: - result &= torch.zeros_like(result) - continue - if isinstance(p_norm, torch.Tensor): - p_value = float(p_norm.reshape(-1)[0].item()) - elif isinstance(p_norm, (int, float, bool)): - p_value = float(p_norm) - else: - result &= torch.zeros_like(result) - continue - positions_raw = layer.params.get("perturbed_positions") - positions = ( - positions_raw - if positions_raw is None or isinstance(positions_raw, (torch.Tensor, list, tuple)) - else None - ) - mask = normalize_position_mask( - positions, - int(center_t.shape[-2]), - batch_shape=tuple(center_t.shape[:-2]), - device=x_batch.device, - ) - if mask.shape[0] == 1 and x_batch.shape[0] != 1: - mask_b = mask.expand(x_batch.shape[0], *mask.shape[1:]) - center_b = center_t.expand_as(x_batch) - else: - mask_b = mask - center_b = center_t.expand_as(x_batch) - delta = x_batch - center_b - clean = (~mask_b).unsqueeze(-1).expand_as(delta) - if bool(clean.any().item()): - clean_ok = torch.where(clean, delta.abs(), torch.zeros_like(delta)).flatten(start_dim=1).amax(dim=1) <= tol - result &= clean_ok - perturbed_delta = delta[mask_b.unsqueeze(-1).expand_as(delta)].reshape(x_batch.shape[0], -1, center_t.shape[-1]) - if perturbed_delta.numel() == 0: - continue - if p_value == float("inf"): - norms = perturbed_delta.abs().amax(dim=-1) - elif p_value == 1.0: - norms = perturbed_delta.abs().sum(dim=-1) - elif p_value == 2.0: - norms = torch.linalg.vector_norm(perturbed_delta, ord=2, dim=-1) - else: - norms = torch.linalg.vector_norm(perturbed_delta, ord=p_value, dim=-1) - result &= (norms <= eps_t.reshape(-1)[0] + tol).all(dim=1) - return result - - # --------------------------------------------------------------------------- # Strategy factories # --------------------------------------------------------------------------- @@ -935,36 +148,63 @@ def _build_bounding( return TopKBounding(order, k=top_k) -def _groups_to_tensors(groups: Dict[int, Any], batch: SubproblemBatch): - # Duplicate (layer, neuron) pairs within a lane are SOUND but wasteful: the - # sign write in _multi_split_from_groups is last-write-wins over the bit - # index, applied uniformly across all 2^k children, so a repeated pair - # yields redundant-but-covering children (the union still constrains that - # neuron both >=0 and <=0). It simply burns 2^k child rows on fewer than - # 2^k distinct regions, so we drop the repeats and shrink k_eff instead. - bb = batch.batch_size - if len(groups) != bb: - return None, None, 0 - deduped: Dict[int, List[Tuple[int, int]]] = {} - for lane in range(bb): - seen: Dict[Tuple[int, int], None] = {} - for lid, nidx in groups.get(lane, []): - seen.setdefault((int(lid), int(nidx)), None) - deduped[lane] = list(seen) - k_eff = min((len(entries) for entries in deduped.values()), default=0) - if k_eff < 1: - return None, None, 0 - device = batch.lb.device - top_layers = torch.zeros(bb, k_eff, dtype=torch.long, device=device) - top_neurons = torch.zeros(bb, k_eff, dtype=torch.long, device=device) - for lane in range(bb): - for j, (lid, nidx) in enumerate(deduped[lane][:k_eff]): - top_layers[lane, j] = lid - top_neurons[lane, j] = nidx - return top_layers, top_neurons, k_eff - - -def _dispatch_dual_solve( +# --------------------------------------------------------------------------- +# BaB engine +# --------------------------------------------------------------------------- + + +def _net_bound_elements(net: Net) -> int: + """Total bound-carrying variables; a proxy for per-lane memory cost.""" + return sum(len(l.out_vars) for l in net.layers if l.out_vars) + + +def _select_spec_rows( + state: Optional[Dict[int, torch.Tensor]], + keep_rows: torch.Tensor, +) -> Optional[Dict[int, torch.Tensor]]: + """Slice the spec axis of per-layer incremental dual state.""" + if state is None: + return None + return { + lid: tensor.index_select(1, keep_rows.to(tensor.device)) + if tensor.dim() >= 3 + else tensor + for lid, tensor in state.items() + } + + +def _neuron_branching_supported(config: BaBConfig) -> bool: + return ( + config.branching_method in NEURON_BRANCHING_METHODS + and config.solver_tier in ("dual_alpha", "dual_alpha_eta") + ) + + +def _witness_residual_branching_active(config: BaBConfig) -> bool: + return config.branching_method == "witness_residual" + + +def _assert_multi_split_invariants( + stats: Dict[str, int], active_k: Optional[int] = None +) -> None: + requested = stats["multi_split_k_requested"] + used = stats["multi_split_k_used"] + waves = stats["multi_split_wave_count"] + clamped = stats["multi_split_clamped_wave_count"] + starved = stats["multi_split_lane_starved_count"] + counters_ok = waves >= 0 and clamped >= 0 and 0 <= starved <= waves + metadata_ok = (waves == 0 and used == 1) or ( + waves > 0 and 2 <= used <= requested + ) + active_ok = active_k is None or 1 <= active_k <= requested + assert requested >= 1 and counters_ok and metadata_ok and active_ok, ( + "MULTI-SPLIT invariant violated: split depth or wave counters are inconsistent " + f"(k_requested={requested}, k_used={used}, active_k={active_k}, " + f"waves={waves}, clamped_waves={clamped}, starved_waves={starved})" + ) + + +def _solve_dual_batch( *, net: Net, assert_layer: Layer, @@ -977,328 +217,131 @@ def _dispatch_dual_solve( keep_rows: Optional[torch.Tensor] = None, root_bounds_dict: Optional[Dict[int, Bounds]] = None, round_policy: Optional[Any] = None, -) -> DualSolveResult: - """Run one dual-family BaB bound pass and decode lane statuses. - - ``keep_rows`` restricts the encoded spec to the given row indices - (ALL-rows kinds only). ``root_bounds_dict`` replaces the per-node forward - pass with the root box's bounds (input-layer entries overridden by each - lane's sub-box). Both are sound by bound monotonicity: certified rows and - per-layer bounds of an ancestor box remain valid on every descendant. - """ - from act.back_end.dual_tf.tf_forward import compute_forward_bounds - from act.back_end.solver.solver_dual import DualSolver, expand_bounds_dict - - solver_tier = getattr(config, "solver_tier", "lp") - block_eps_updates = _install_embedding_child_block_eps(net, batched_bounds, batch) - if root_bounds_dict is not None: - bounds_dict_dual = expand_bounds_dict(root_bounds_dict, k_actual) - lane_box = Bounds(batched_bounds.lb, batched_bounds.ub) - for layer in net.layers: - kind_up = layer.kind.upper() if isinstance(layer.kind, str) else layer.kind - if kind_up in ("INPUT", "INPUT_SPEC") and layer.id in bounds_dict_dual: - bounds_dict_dual[layer.id] = lane_box - if batch.split_signs: - refreshed = _interval_refresh_bounds(net, bounds_dict_dual, batch.split_signs) - if refreshed is not None: - bounds_dict_dual = refreshed - psr_mode = getattr(config, "per_subproblem_refine", "none") - psr_rows_cap = getattr(config, "per_subproblem_refine_rows_cap", 64) - psr_iters = getattr(config, "per_subproblem_refine_iters", 0) - if round_policy is not None: - if round_policy.refine_mode is not None: - psr_mode = round_policy.refine_mode - if round_policy.refine_rows_cap is not None: - psr_rows_cap = round_policy.refine_rows_cap - if round_policy.refine_iters is not None: - psr_iters = round_policy.refine_iters - if psr_mode != "none": - bounds_dict_dual = DualSolver().refine_intermediate_bounds_batched( - net, - bounds_dict_dual, - split_signs=batch.split_signs, - mode=psr_mode, - rows_cap=psr_rows_cap, - optimize_iters=psr_iters, - ) - else: - bounds_dict_dual = compute_forward_bounds(net, batched_bounds.lb, batched_bounds.ub) - out_kind_raw = assert_layer.params["kind"] - if not isinstance(out_kind_raw, str): - raise TypeError(f"ASSERT kind must be str, got {type(out_kind_raw).__name__}") - - out_spec_fields = OutputSpec(kind=out_kind_raw)._gather_rows( - rows=None, - batch_size=1, - device=batched_bounds.lb.device, - dtype=batched_bounds.lb.dtype, - shared_ndim={}, - source=assert_layer.params, - source_batch_size=1, - drop_singleton_batch=True, +) -> DualBatchResult: + """Prepare BaB-specific bounds/state policy, then call ``DualSolver``.""" + solver = DualSolver() + block_eps_updates = _install_embedding_child_block_eps( + net, batched_bounds, batch ) - - out_spec = OutputSpec( - kind=out_kind_raw, - y_true=out_spec_fields.get("y_true"), - margin=out_spec_fields.get("margin"), - c=out_spec_fields.get("c"), - d=out_spec_fields.get("d"), - lb=out_spec_fields.get("lb"), - ub=out_spec_fields.get("ub"), - ) - sample_bounds = next(iter(bounds_dict_dual.values())) - device = sample_bounds.lb.device - dtype = sample_bounds.lb.dtype - assert_preds = net.preds.get(assert_layer.id, []) - if len(assert_preds) != 1: - raise ValueError( - f"ASSERT layer {assert_layer.id} must have exactly 1 predecessor, " - f"got {len(assert_preds)}" + try: + input_split_child = ( + batch.depths.numel() > 0 + and bool((batch.depths.max() > 0).item()) + and not batch.split_signs ) - output_bounds = bounds_dict_dual[assert_preds[0]] - n_out = int(output_bounds.lb.flatten(start_dim=1).shape[-1]) - encoded_spec = out_spec.encode_linear(B=k_actual, n_out=n_out, device=device, dtype=dtype) - m_specs = int(encoded_spec["M"]) - dual = DualSolver() - - if out_spec.kind == OutKind.UNSAFE_LINEAR: - c_rows = cast(torch.Tensor, encoded_spec["C"]).contiguous() - thresholds = cast(torch.Tensor, encoded_spec["thresholds"]).contiguous() - else: - c_rows = -cast(torch.Tensor, encoded_spec["C"]).contiguous() - thresholds = -cast(torch.Tensor, encoded_spec["thresholds"]).contiguous() - if keep_rows is not None: - idx = keep_rows.to(device=device, dtype=torch.long) - c_rows = ( - c_rows.view(k_actual, m_specs, n_out) - .index_select(1, idx) - .reshape(k_actual * int(idx.numel()), n_out) - .contiguous() + use_root_dict = root_bounds_dict is not None and not input_split_child + if use_root_dict: + assert root_bounds_dict is not None + bounds_dict = expand_bounds_dict(root_bounds_dict, k_actual) + lane_box = Bounds(batched_bounds.lb, batched_bounds.ub) + for layer in net.layers: + kind = ( + layer.kind.upper() + if isinstance(layer.kind, str) + else layer.kind + ) + if ( + kind in (LayerKind.INPUT.value, LayerKind.INPUT_SPEC.value) + and layer.id in bounds_dict + ): + bounds_dict[layer.id] = lane_box + if config.root_bounds_reuse == "split_refresh" and batch.split_signs: + refreshed = solver._interval_refresh_bounds( + net, bounds_dict, batch.split_signs + ) + if refreshed is not None: + bounds_dict = refreshed + refine_mode = config.per_subproblem_refine + refine_rows_cap = config.per_subproblem_refine_rows_cap + refine_iters = config.per_subproblem_refine_iters + if round_policy is not None: + if round_policy.refine_mode is not None: + refine_mode = round_policy.refine_mode + if round_policy.refine_rows_cap is not None: + refine_rows_cap = round_policy.refine_rows_cap + if round_policy.refine_iters is not None: + refine_iters = round_policy.refine_iters + if refine_mode != "none": + bounds_dict = solver.refine_intermediate_bounds_batched( + net, + bounds_dict, + split_signs=batch.split_signs, + mode=refine_mode, + rows_cap=refine_rows_cap, + optimize_iters=refine_iters, + ) + else: + bounds_dict = compute_forward_bounds( + net, + batched_bounds.lb, + batched_bounds.ub, + forward_lin_max_perturbed=dual_config.forward_lin_max_perturbed, ) - thresholds = thresholds.index_select(1, idx).contiguous() - m_specs = int(idx.numel()) - active_mask = torch.ones(k_actual, m_specs, dtype=torch.bool, device=device) - - return_nu = _neuron_branching_supported(config) - supports_return_nu = "return_nu_per_layer" in inspect.signature( - dual.compute_certified_bound - ).parameters - compute_certified_bound = cast(Any, dual.compute_certified_bound) - - is_child_batch = bool(batch.depths.min().item() > 0) if batch.depths.numel() else False - if optimize: - dual_result = compute_certified_bound( - net, - bounds_dict_dual, - c_rows, - M=m_specs, - optimize=True, - optimize_alpha=not ( - getattr(config, "eta_only_children", False) and is_child_batch - ), - refresh_forward=root_bounds_dict is None, - n_iters=dual_config.n_iters, - lr_alpha=dual_config.lr_alpha, - lr_beta=dual_config.lr_beta, - lr_decay=dual_config.lr_decay, - eta=batch.incremental_eta if solver_tier == "dual_alpha_eta" else None, - incremental_alphas=batch.incremental_alpha if dual_config.incremental_start_enabled else None, - incremental_etas=( - batch.incremental_eta - if solver_tier == "dual_alpha_eta" - and dual_config.incremental_start_enabled - else None - ), - split_signs=batch.split_signs if solver_tier == "dual_alpha_eta" else None, - return_optimized=True, - return_sce=True, - per_class_alpha=dual_config.per_class_alpha, - **({"return_nu_per_layer": True} if return_nu and supports_return_nu else {}), - ) - margins_flat = dual_result.margins - sce = cast(Optional[torch.Tensor], dual_result.sce) - batch.incremental_alpha = dual_result.alpha_state - if solver_tier == "dual_alpha_eta": - batch.incremental_eta = dual_result.eta_state - else: - dual_result = compute_certified_bound( - net, - bounds_dict_dual, - c_rows, - M=m_specs, - return_sce=True, - **({"return_nu_per_layer": True} if return_nu and supports_return_nu else {}), + out_kind = assert_layer.params["kind"] + if not isinstance(out_kind, str): + raise TypeError( + f"ASSERT kind must be str, got {type(out_kind).__name__}" + ) + fields = OutputSpec(kind=out_kind)._gather_rows( + rows=None, + batch_size=1, + device=batched_bounds.lb.device, + dtype=batched_bounds.lb.dtype, + shared_ndim={}, + source=assert_layer.params, + source_batch_size=1, + drop_singleton_batch=True, ) - margins_flat = dual_result.margins - sce = cast(Optional[torch.Tensor], dual_result.sce) - - margins = margins_flat.view(k_actual, m_specs) - slack = margins - thresholds - if out_spec.kind == OutKind.UNSAFE_LINEAR: - certified = ((slack > 0) & active_mask).any(dim=-1) - candidate_rows = torch.zeros(k_actual, dtype=torch.long, device=device) - else: - violations = (slack < 0) & active_mask - certified = ~violations.any(dim=-1) - candidate_rows = torch.where( - violations.any(dim=1), - violations.to(torch.int64).argmax(dim=1), - torch.zeros(k_actual, dtype=torch.long, device=device), + out_spec = OutputSpec( + kind=out_kind, + y_true=fields.get("y_true"), + margin=fields.get("margin"), + c=fields.get("c"), + d=fields.get("d"), + lb=fields.get("lb"), + ub=fields.get("ub"), ) - - statuses = tuple( - SolveStatus.UNSAT if bool(is_certified.item()) else SolveStatus.SAT - for is_certified in certified - ) - nvars = max((max(layer.out_vars) for layer in net.layers if layer.out_vars), default=-1) + 1 - x_candidate = torch.zeros(k_actual, nvars, device=device, dtype=dtype) - input_ids_list = get_input_ids(net) - if sce is not None: - sce_flat = sce.flatten(start_dim=1).to(device=device) - row_offsets = torch.arange(k_actual, device=device) * m_specs + candidate_rows.to(device=device) - chosen_sce = sce_flat.index_select(0, row_offsets) - input_ids = torch.tensor(input_ids_list, device=device, dtype=torch.long) - x_candidate[:, input_ids] = chosen_sce.to(device=device, dtype=dtype) - else: - # TODO: extend CE-candidate generation for dual paths that do not return SCE. - statuses = tuple( - SolveStatus.UNSAT if status == SolveStatus.UNSAT else SolveStatus.UNKNOWN - for status in statuses + eta_enabled = config.solver_tier == "dual_alpha_eta" + is_child_batch = ( + bool(batch.depths.min().item() > 0) if batch.depths.numel() else False ) - solution = BatchLPSolution( - statuses=statuses, - x=x_candidate, - max_viol=-slack.min(dim=1).values.detach(), - ) - branch_bounds: Optional[Dict[int, Bounds]] = None - branch_nu: Optional[Dict[int, torch.Tensor]] = None - if return_nu and root_bounds_dict is not None: - # Heuristic-only consumer (BaBSR/FSB scores). The optimize path does - # not emit nu, and nu=None silently degrades neuron branching to - # input-axis splits - so run one grad-free backward at the converged - # alpha/eta to extract per-layer nu on the same (reused) bounds. - branch_bounds = bounds_dict_dual - branch_nu = getattr(dual_result, "nu_per_layer", None) - if branch_nu is None: - nu_pass = dual.compute_certified_bound( - net, - bounds_dict_dual, - c_rows, - M=m_specs, - alpha=getattr(dual_result, "alpha_state", None), - eta=( - getattr(dual_result, "eta_state", None) - if solver_tier == "dual_alpha_eta" - else None - ), - split_signs=( - batch.split_signs if solver_tier == "dual_alpha_eta" else None - ), - return_nu_per_layer=True, - ) - branch_nu = nu_pass.nu_per_layer - elif return_nu: - branch_bounds, branch_nu = dual.recompute_bounds_and_nu( + result = solver.solve_spec_batch( net, - bounds_dict_dual, - c_rows, - m_specs, - alpha_state=getattr(dual_result, "alpha_state", None), - eta_state=( - getattr(dual_result, "eta_state", None) - if solver_tier == "dual_alpha_eta" + bounds_dict, + out_spec, + optimize=optimize, + dual_config=dual_config, + input_shape=tuple(batched_bounds.lb.shape[1:]), + keep_rows=keep_rows, + split_signs=batch.split_signs if eta_enabled else None, + eta=batch.incremental_eta if eta_enabled else None, + incremental_alphas=( + batch.incremental_alpha + if dual_config.incremental_start_enabled else None ), - split_signs=( - batch.split_signs if solver_tier == "dual_alpha_eta" else None + incremental_etas=( + batch.incremental_eta + if eta_enabled and dual_config.incremental_start_enabled + else None ), - per_class_alpha=dual_config.per_class_alpha, - ) - try: - return DualSolveResult( - solution=solution, - bounds_dict=branch_bounds, - nu_per_layer=branch_nu, - witness_input=x_candidate[:, torch.tensor(input_ids_list, device=device, dtype=torch.long)] - .reshape(k_actual, *batched_bounds.lb.shape[1:]) - .detach() - if sce is not None - else None, - row_slack=slack.detach(), + optimize_alpha=not ( + config.eta_only_children and is_child_batch + ), + refresh_forward=not use_root_dict, + return_nu=_neuron_branching_supported(config), + reuse_bounds_for_branching=use_root_dict, ) + if optimize: + batch.incremental_alpha = result.alpha_state + if eta_enabled: + batch.incremental_eta = result.eta_state + return result finally: _restore_embedding_child_block_eps(block_eps_updates) -def _finite_embedding_spec(net: Net) -> Optional[Layer]: - for layer in net.layers: - if layer.kind == "INPUT_SPEC" and layer.params.get("kind") == InKind.LP_EMBEDDING: - p_norm = layer.params.get("p_norm", float("inf")) - if isinstance(p_norm, torch.Tensor): - p_value = float(p_norm.reshape(-1)[0].item()) - elif isinstance(p_norm, (int, float, bool)): - p_value = float(p_norm) - else: - continue - if p_value != float("inf"): - return layer - return None - - -def _install_embedding_child_block_eps( - net: Net, - batched_bounds: Bounds, - batch: SubproblemBatch, -) -> list[tuple[Layer, ParamValue]]: - spec = _finite_embedding_spec(net) - if spec is None or batch.depths.numel() == 0 or int(batch.depths.max().item()) == 0: - return [] - p_raw = spec.params.get("p_norm", float("inf")) - if isinstance(p_raw, torch.Tensor): - p_norm = float(p_raw.reshape(-1)[0].item()) - elif isinstance(p_raw, (int, float, bool)): - p_norm = float(p_raw) - else: - return [] - input_shape = tuple(batched_bounds.lb.shape[1:]) - positions_raw = spec.params.get("perturbed_positions") - positions = positions_raw if isinstance(positions_raw, torch.Tensor) else None - block_eps = rederive_embedding_block_eps( - batched_bounds.lb.flatten(start_dim=1), - batched_bounds.ub.flatten(start_dim=1), - input_shape, - positions, - p_norm, - ) - old_values: list[tuple[Layer, ParamValue]] = [] - for layer in net.layers: - kind_up = layer.kind.upper() if isinstance(layer.kind, str) else layer.kind - if kind_up in ("INPUT", "INPUT_SPEC"): - old_values.append((layer, layer.params.get("bab_block_eps", None))) - layer.params["bab_block_eps"] = block_eps - return old_values - - -def _restore_embedding_child_block_eps(updates: list[tuple[Layer, ParamValue]]) -> None: - for layer, old in updates: - if old is None: - layer.params.pop("bab_block_eps", None) - else: - layer.params["bab_block_eps"] = old - - -# --------------------------------------------------------------------------- -# BaB engine -# --------------------------------------------------------------------------- - - -def _net_bound_elements(net: Net) -> int: - """Total bound-carrying variables; a proxy for per-lane memory cost.""" - return sum(len(l.out_vars) for l in net.layers if getattr(l, "out_vars", None)) - - def _auto_batch_budget_bytes(safety: float) -> float: """Memory the auto sizer may use: min(safety*total, 90% of what this process can reclaim), so it shares the GPU with other processes.""" @@ -1310,9 +353,9 @@ def _auto_batch_budget_bytes(safety: float) -> float: def _auto_initial_batch(net: Net, config: BaBConfig) -> int: """Conservative first batch from net size; the loop recalibrates it from the measured per-lane peak after the first real round.""" - safety = float(getattr(config, "auto_batch_safety", 0.55)) - cap = int(getattr(config, "auto_batch_cap", 2048)) - floor = int(getattr(config, "auto_batch_floor", 8)) + safety = float(config.auto_batch_safety) + cap = int(config.auto_batch_cap) + floor = int(config.auto_batch_floor) per_lane = 4.0 * max(1, _net_bound_elements(net)) * 256.0 k = int(_auto_batch_budget_bytes(safety) / per_lane) return max(floor, min(cap, k)) @@ -1325,9 +368,9 @@ def _auto_recalibrate_batch(peak_bytes: float, max_k_seen: int, config: BaBConfi folds in the one-time root/presolve peak), so the sizer errs toward fewer lanes - safe against OOM while still ramping up on small nets with spare memory.""" - safety = float(getattr(config, "auto_batch_safety", 0.55)) - cap = int(getattr(config, "auto_batch_cap", 2048)) - floor = int(getattr(config, "auto_batch_floor", 8)) + safety = float(config.auto_batch_safety) + cap = int(config.auto_batch_cap) + floor = int(config.auto_batch_floor) bpl = max(peak_bytes / max(1, max_k_seen), 1.0) k = int(_auto_batch_budget_bytes(safety) / bpl) return max(floor, min(cap, k)) @@ -1339,7 +382,7 @@ def verify_bab_batched( solver_factory: Callable[[], Solver], config: Optional[BaBConfig] = None, *, - max_batch_size: Optional[Union[int, str]] = None, + max_batch_size: Union[int, str], time_budget_s: Optional[float] = None, dual_config: Optional[DualConfig] = None, verbose: bool = False, @@ -1368,10 +411,12 @@ def verify_bab_batched( net: ACT network with a single-instance INPUT_SPEC (B=1 seed). solver_factory: callable returning a fresh ``Solver`` per iteration (no state leakage across iterations). - config: ``BaBConfig``; ``bab_max_batch_size`` (if present, otherwise 8) - caps K. ``max_depth`` and ``max_nodes`` cap the search tree. - max_batch_size: explicit override for K cap; takes precedence over - ``config.bab_max_batch_size``. + config: ``BaBConfig``; ``max_depth`` and ``max_nodes`` cap the search + tree. + max_batch_size: required cap on K, an int >= 1 or ``"auto"`` (sizes K + from GPU memory; ``config.auto_batch_cap`` on CPU). There is no + implicit default: callers pass the central + ``BackendConfig.bab_max_batch_size`` or an explicit value. time_budget_s: wall-clock budget (default 300 s). verbose: reserved. _k_log: diagnostic only — if supplied, the actual K used per iteration @@ -1379,6 +424,10 @@ def verify_bab_batched( """ if config is None: config = BaBConfig() + assert_layer = get_assert_layer(net) + climb = ClimbSession.from_config( + config, net, cast(str, assert_layer.params["kind"]) + ) if dual_config is None: dual_config = DualConfig() auto_batch = isinstance(max_batch_size, str) and max_batch_size == "auto" @@ -1386,12 +435,10 @@ def verify_bab_batched( effective_batch = ( _auto_initial_batch(net, config) if torch.cuda.is_available() - else int(getattr(config, "auto_batch_cap", 512)) + else int(config.auto_batch_cap) ) - elif max_batch_size is None: - effective_batch = int(getattr(config, "bab_max_batch_size", 8)) else: - effective_batch = int(cast(int, max_batch_size)) + effective_batch = int(max_batch_size) if effective_batch < 1: raise ValueError(f"max_batch_size must be >= 1, got {effective_batch}") max_k_seen = 0 @@ -1400,8 +447,6 @@ def verify_bab_batched( fsb_dual_solver = None if config.branching_method == "fsb": - from act.back_end.solver.solver_dual import DualSolver - fsb_dual_solver = DualSolver() # "gain" measures child bounds directly; its fallback (when no measured # decision is available) is BaBSR — it reuses the dual ν scores and @@ -1413,21 +458,24 @@ def verify_bab_batched( brancher = _build_branching_strategy(brancher_method, dual_solver=fsb_dual_solver) multi_split_stats: Dict[str, int] = { - "multi_split_k_requested": int(getattr(config, "multi_split_levels", 1)), + "multi_split_k_requested": int(config.multi_split_levels), "multi_split_k_used": 1, "multi_split_wave_count": 0, "multi_split_clamped_wave_count": 0, "multi_split_lane_starved_count": 0, } - def _branching_metadata() -> Dict[str, int]: - meta: Dict[str, int] = dict(multi_split_stats) + def _branching_metadata() -> Dict[str, Any]: + _assert_multi_split_invariants(multi_split_stats) + meta: Dict[str, Any] = dict(multi_split_stats) meta["bounding_top_k_effective"] = int(getattr(pool, "k", 0)) if _witness_residual_branching_active(config): meta["witness_residual_fallback_count"] = int(getattr(brancher, "fallback_count", 0)) meta["witness_residual_diff_from_babsr_count"] = int( getattr(brancher, "different_from_babsr_count", 0) ) + if climb is not None: + meta.update(climb.metadata()) return meta pool = _build_bounding( @@ -1443,11 +491,11 @@ def _branching_metadata() -> Dict[str, int]: llm_probe: Any = None _llm: Any = None _wave_index = 0 - if getattr(config, "llm_probe_enabled", False): + if config.llm_probe_enabled: from act.pipeline.verification import llm_probe as _llm llm_probe = _llm.build_llm_probe(config) - provenance = bool(getattr(config, "provenance_enabled", False)) or isinstance( + provenance = bool(config.provenance_enabled) or isinstance( pool, MCTSBounding ) if provenance and not isinstance(pool, (TopKBounding, MCTSBounding)): @@ -1461,11 +509,10 @@ def _branching_metadata() -> Dict[str, int]: "pre-split root batch carries no node_id/parent_id provenance" ) node_counter = 0 - fanout = max(2, int(getattr(config, "input_split_fanout", 2))) - frontier_cap = int(getattr(config, "frontier_cap", 0)) + fanout = max(2, int(config.input_split_fanout)) + frontier_cap = int(config.frontier_cap) spec_layers = gather_input_spec_layers(net) - assert_layer = get_assert_layer(net) root_bounds = seed_from_input_specs(spec_layers) input_shape: tuple[int, ...] = tuple(root_bounds.lb.shape[1:]) @@ -1497,36 +544,39 @@ def _branching_metadata() -> Dict[str, int]: # on the root box stay certified on every sub-box, so descendants only # carry the unproven rows. spec_keep_rows: Optional[torch.Tensor] = None - presolve_tier = getattr(config, "solver_tier", "lp") + presolve_tier = config.solver_tier root_fwd: Optional[Dict[int, Bounds]] = None - refine_mode = getattr(config, "intermediate_refine", "none") + refine_mode = config.intermediate_refine + reuse_mode = config.root_bounds_reuse if presolve_tier in ("dual", "dual_alpha", "dual_alpha_eta") and ( - getattr(config, "reuse_root_bounds", False) or refine_mode != "none" + reuse_mode != "none" or refine_mode != "none" ): - from act.back_end.dual_tf.tf_forward import compute_forward_bounds - from act.back_end.solver.solver_dual import DualSolver - - root_fwd = compute_forward_bounds(net, root_bounds.lb, root_bounds.ub) + root_fwd = compute_forward_bounds( + net, + root_bounds.lb, + root_bounds.ub, + forward_lin_max_perturbed=dual_config.forward_lin_max_perturbed, + ) if refine_mode != "none": root_fwd = DualSolver().refine_intermediate_bounds( net, root_fwd, mode=refine_mode, - blowup_ratio=getattr(config, "intermediate_refine_ratio", 10.0), + blowup_ratio=config.intermediate_refine_ratio, ) - # Per-node bound reuse is governed solely by reuse_root_bounds: root_fwd may + # Per-node bound reuse is governed solely by root_bounds_reuse: root_fwd may # exist just for the root presolve/refine above, and passing it to descendant # solves would freeze every child's intermediate bounds at root tightness # (fatal for input-split BaB, where the whole gain comes from recomputing # intermediates on the smaller box). node_root_fwd: Optional[Dict[int, Bounds]] = ( - root_fwd if getattr(config, "reuse_root_bounds", False) else None + root_fwd if reuse_mode != "none" else None ) if ( presolve_tier in ("dual", "dual_alpha", "dual_alpha_eta") and assert_layer.params.get("kind") != OutKind.UNSAFE_LINEAR ): - presolve = _dispatch_dual_solve( + presolve = _solve_dual_batch( net=net, assert_layer=assert_layer, batched_bounds=Bounds(root_bounds.lb, root_bounds.ub), @@ -1549,6 +599,7 @@ def _branching_metadata() -> Dict[str, int]: "spec_rows_total": total_rows, "spec_rows_kept": 0, "resolved_by": "root_presolve", + **(climb.metadata() if climb is not None else {}), }, ) keep = torch.where(unproven)[0] @@ -1563,15 +614,19 @@ def _branching_metadata() -> Dict[str, int]: root_batch.split_signs = _select_spec_rows( root_batch.split_signs, keep, ) - presplit_k = int(getattr(config, "presplit_levels", 0)) + presplit_k = int(config.presplit_levels) if ( presplit_k > 0 and root_batch.batch_size == 1 and presolve.bounds_dict is not None and presolve.nu_per_layer is not None ): - presplit = _presplit_root( - root_batch, presolve.bounds_dict, presolve.nu_per_layer, presplit_k, + presplit = presplit_root( + root_batch, + net, + presolve.bounds_dict, + presolve.nu_per_layer, + presplit_k, ) if presplit is not None: root_batch = presplit @@ -1593,6 +648,11 @@ def _branching_metadata() -> Dict[str, int]: if elapsed >= budget_s or processed >= config.max_nodes: break + if climb is not None: + climb.before_pop(cast(TopKBounding, pool)) + if pool.empty: + break + remaining_nodes = config.max_nodes - processed k_requested = min(len(pool), effective_batch, remaining_nodes) if k_requested <= 0: @@ -1619,6 +679,8 @@ def _branching_metadata() -> Dict[str, int]: k_actual = batch.batch_size if _k_log is not None: _k_log.append(k_actual) + if climb is not None: + climb.after_pop(len(pool) + k_actual) if input_shape: k_lb = batch.lb.reshape(k_actual, *input_shape) @@ -1628,18 +690,19 @@ def _branching_metadata() -> Dict[str, int]: k_ub = batch.ub batched_bounds = Bounds(k_lb, k_ub) - solver_tier = getattr(config, "solver_tier", "lp") + solver_tier = config.solver_tier neuron_branching_supported = _neuron_branching_supported(config) bounds_dict_for_branching: Optional[Dict[int, Bounds]] = None nu_per_layer_for_branching: Optional[Dict[int, torch.Tensor]] = None witness_input_for_branching: Optional[torch.Tensor] = None + dual_solve_result: Optional[DualBatchResult] = None if solver_tier == "lp": solver = solver_factory() solution = setup_and_solve_batch( net, batched_bounds, solver, timelimit=None, ) elif solver_tier == "dual": - dual_solve_result = _dispatch_dual_solve( + dual_solve_result = _solve_dual_batch( net=net, assert_layer=assert_layer, batched_bounds=batched_bounds, @@ -1654,7 +717,7 @@ def _branching_metadata() -> Dict[str, int]: ) solution = dual_solve_result.solution elif solver_tier in ("dual_alpha", "dual_alpha_eta"): - dual_solve_result = _dispatch_dual_solve( + dual_solve_result = _solve_dual_batch( net=net, assert_layer=assert_layer, batched_bounds=batched_bounds, @@ -1676,6 +739,9 @@ def _branching_metadata() -> Dict[str, int]: f"Unknown solver_tier={solver_tier!r}. Valid: {VALID_SOLVER_TIERS}." ) + if climb is not None: + climb.after_main_bound(k_actual) + node_lower_bound = (-solution.max_viol).detach() if batch.lower_bound is not None: # Bound inheritance: a child region is a subset of its parent, so @@ -1747,82 +813,28 @@ def _branching_metadata() -> Dict[str, int]: device=batch.lb.device, dtype=torch.long, ) + if climb is not None: + assert dual_solve_result is not None + climb.learn(batch, solution.statuses, dual_solve_result, k_actual) + if int(unresolved_idx.numel()) > 0: - def _select_incremental_state( - state: Optional[dict[int, torch.Tensor]], - indices: torch.Tensor, - ) -> Optional[dict[int, torch.Tensor]]: - if state is None: - return None - return { - layer_id: tensor.index_select(0, indices.to(tensor.device)) - for layer_id, tensor in state.items() - } - - unresolved = SubproblemBatch( - lb=batch.lb.index_select(0, unresolved_idx.to(batch.lb.device)), - ub=batch.ub.index_select(0, unresolved_idx.to(batch.ub.device)), - depths=batch.depths.index_select(0, unresolved_idx.to(batch.depths.device)), - incremental_alpha=_select_incremental_state(batch.incremental_alpha, unresolved_idx), - incremental_eta=_select_incremental_state(batch.incremental_eta, unresolved_idx), - split_signs=_select_incremental_state(batch.split_signs, unresolved_idx), - parent_margins=( - batch.parent_margins.index_select(0, unresolved_idx.to(batch.parent_margins.device)) - if batch.parent_margins is not None - else None - ), - lower_bound=node_lower_bound.index_select( - 0, unresolved_idx.to(node_lower_bound.device) - ), - node_id=( - batch.node_id.index_select(0, unresolved_idx.to(batch.node_id.device)) - if batch.node_id is not None - else None - ), - parent_id=( - batch.parent_id.index_select(0, unresolved_idx.to(batch.parent_id.device)) - if batch.parent_id is not None - else None - ), + unresolved = batch.select(unresolved_idx) + unresolved.lower_bound = node_lower_bound.index_select( + 0, unresolved_idx.to(node_lower_bound.device) ) + if climb is not None: + unresolved = climb.before_split(cast(TopKBounding, pool), unresolved) branch_mask = unresolved.depths < int(config.max_depth) if bool((~branch_mask).any().item()): any_dropped_max_depth = True branch_idx = torch.where(branch_mask)[0] if int(branch_idx.numel()) > 0: - branch_batch = SubproblemBatch( - lb=unresolved.lb.index_select(0, branch_idx.to(unresolved.lb.device)), - ub=unresolved.ub.index_select(0, branch_idx.to(unresolved.ub.device)), - depths=unresolved.depths.index_select(0, branch_idx), - incremental_alpha=_select_incremental_state(unresolved.incremental_alpha, branch_idx), - incremental_eta=_select_incremental_state(unresolved.incremental_eta, branch_idx), - split_signs=_select_incremental_state(unresolved.split_signs, branch_idx), - parent_margins=( - unresolved.parent_margins.index_select(0, branch_idx) - if unresolved.parent_margins is not None - else None - ), - lower_bound=( - unresolved.lower_bound.index_select(0, branch_idx) - if unresolved.lower_bound is not None - else None - ), - node_id=( - unresolved.node_id.index_select(0, branch_idx.to(unresolved.node_id.device)) - if unresolved.node_id is not None - else None - ), - parent_id=( - unresolved.parent_id.index_select(0, branch_idx.to(unresolved.parent_id.device)) - if unresolved.parent_id is not None - else None - ), - ) + branch_batch = unresolved.select(branch_idx) if neuron_branching_supported: full_branch_idx = unresolved_idx.index_select( 0, branch_idx.to(unresolved_idx.device) ) - bd_branch, nu_branch = _slice_branching_state( + bd_branch, nu_branch = slice_branching_state( bounds_dict_for_branching, nu_per_layer_for_branching, full_branch_idx, @@ -1834,24 +846,25 @@ def _select_incremental_state( 0, full_branch_idx.to(witness_input_for_branching.device), ) - witness_preact_branch = _witness_relu_preactivations( + witness_preact_branch = witness_relu_preactivations( net, witness_branch, input_shape, + dual_config, ) multi = None - multi_k = int(getattr(config, "multi_split_levels", 1)) + multi_k = int(config.multi_split_levels) if llm_probe is not None and _llm is not None and llm_probe.wants_neuron: # neuron_topk>0 => never bail on candidate count: enumerate the # full set (limit=None) and let advise_neuron_groups truncate to # the top-K by score, so the LLM always decides (with a bounded # view) instead of falling back to FSB. neuron_topk==0 keeps the # legacy "bail to FSB when > max_candidates_total" behavior. - _neuron_topk = int(getattr(config, "llm_probe_neuron_topk", 0)) + _neuron_topk = int(config.llm_probe_neuron_topk) _cand_dicts = enumerate_unstable_candidates( branch_batch, bd_branch, nu_branch, limit=None if _neuron_topk > 0 - else getattr(config, "llm_probe_max_candidates_total", 1024), + else config.llm_probe_max_candidates_total, ) if _cand_dicts: _ngroups = llm_probe.advise_neuron_groups(_llm.build_frontier_stats( @@ -1864,9 +877,9 @@ def _select_incremental_state( candidates=[_llm.CandidateSummary(**_d) for _d in _cand_dicts], )) if _ngroups is not None: - _tl, _tn, _keff = _groups_to_tensors(_ngroups, branch_batch) + _tl, _tn, _keff = groups_to_tensors(_ngroups, branch_batch) if _tl is not None and _tn is not None: - multi = _multi_split_from_groups(branch_batch, net, _tl, _tn, _keff) + multi = split_neurons(branch_batch, net, _tl, _tn, _keff) _wave_split_used = _keff # Joint splitting is orthogonal to how a split is SCORED, so # it is no longer keyed on branching_method == "gain": the k @@ -1922,11 +935,14 @@ def _select_incremental_state( if multi[0].batch_size != branch_batch.batch_size * (2 ** k_adaptive): multi_split_stats["multi_split_lane_starved_count"] += 1 if multi is not None: + _assert_multi_split_invariants( + multi_split_stats, active_k=_wave_split_used + ) children, parent_index = multi else: decision = None if config.branching_method == "gain": - decision = _gain_tested_decision( + decision = gain_tested_decision( branch_batch, net, assert_layer, @@ -1937,6 +953,7 @@ def _select_incremental_state( bd_branch, nu_branch, input_shape, + solve_dual=_solve_dual_batch, ) if decision is None: extra_branch_kwargs = ( @@ -1952,50 +969,47 @@ def _select_incremental_state( **extra_branch_kwargs, ) decision = cast(SplitDecision, cast(Any, brancher).select(scores)) + if _witness_residual_branching_active(config): + _assert_witness_residual_decision_unstable( + decision, bd_branch + ) if decision.kind == "input_axis": + if climb is not None: + climb.reject_input_axis_split() decision.fanout = fanout - children, parent_index = _split_from_decision(branch_batch, decision, net) + children, parent_index = decision.apply(branch_batch, net) else: scores = brancher.compute_scores(branch_batch, net) legacy_decision = cast(Any, brancher).select(scores) - split_fanout = fanout if isinstance(legacy_decision, SplitDecision): - if legacy_decision.cut_dim is not None: - split_dims = _input_axis_decision_tensor( - SplitDecision(kind="input_axis", input_axis=legacy_decision.cut_dim), - branch_batch, - ) - else: - if legacy_decision.input_axis is None: - raise ValueError("input-axis decision missing input_axis") - split_dims = _input_axis_decision_tensor( - legacy_decision, - branch_batch, - ) - split_fanout = max(2, int(getattr(legacy_decision, "fanout", fanout))) + decision = legacy_decision + split_dims = decision.input_axes(branch_batch) else: split_dims = torch.as_tensor( legacy_decision, device=branch_batch.lb.device, dtype=torch.long, ).reshape(-1) + decision = SplitDecision( + kind="input_axis", + input_axis=split_dims, + fanout=fanout, + ) widths = branch_batch.widths() _last_input_widths = ( widths.mean(dim=0).tolist() if widths.shape[1] <= 32 else None ) - if _wave_policy is not None and getattr(_wave_policy, "input_split_dim", None) is not None: + if _wave_policy is not None and _wave_policy.input_split_dim is not None: # LLM-advised input dimension (already range-clipped). Lanes where # the advised dim has zero width keep the brancher's choice: # splitting a zero-width dim yields identical children (livelock). advised = torch.full_like(split_dims, int(_wave_policy.input_split_dim)) has_width = widths.gather(1, advised.unsqueeze(1)).squeeze(1) > 0 split_dims = torch.where(has_width, advised, split_dims) - if _wave_policy is not None and getattr(_wave_policy, "input_split_fanout", None) is not None: - split_fanout = int(_wave_policy.input_split_fanout) - if split_fanout == 2: - children, parent_index = split_input(branch_batch, split_dims) - else: - children, parent_index = split_input_nary(branch_batch, split_dims, split_fanout) + decision.input_axis = split_dims + if _wave_policy is not None and _wave_policy.input_split_fanout is not None: + decision.fanout = int(_wave_policy.input_split_fanout) + children, parent_index = decision.apply(branch_batch, net) if provenance: pid = branch_batch.node_id @@ -2010,6 +1024,8 @@ def _select_incremental_state( ) node_counter += nc pool.push(children) + if climb is not None: + climb.after_children(children.batch_size, len(pool)) if frontier_cap > 0 and len(pool) > frontier_cap: if pool.evict_to(frontier_cap) > 0: any_dropped_frontier_cap = True @@ -2118,7 +1134,7 @@ def verify_bab( time_budget_s if time_budget_s is not None else (timelimit if timelimit is not None else 300.0) ) - solver_tier = getattr(config, "solver_tier", "lp") + solver_tier = config.solver_tier if solver_tier not in VALID_SOLVER_TIERS: raise ValueError( f"Unknown solver_tier={solver_tier!r}. Valid: {VALID_SOLVER_TIERS}." @@ -2133,59 +1149,3 @@ def verify_bab( verbose=verbose, dual_config=dual_config, ) - - -# --------------------------------------------------------------------------- -# Module tests -# --------------------------------------------------------------------------- - - - - - - - - -class _IdentityOutput(torch.nn.Module): # pragma: no cover - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x.reshape(x.shape[0], -1) - - -def _make_assert_layer(kind: str, params: dict[str, ParamValue], n_out: int) -> Layer: # pragma: no cover - from act.back_end.layer_schema import LayerKind - - merged: dict[str, ParamValue] = {"kind": kind} - merged.update(params) - if "C" not in merged or "thresholds" not in merged or "M" not in merged: - batch_size = 1 - for key in ("y_true", "margin", "c", "d", "lb", "ub"): - value = merged.get(key) - if isinstance(value, torch.Tensor) and value.dim() > 0: - batch_size = max(batch_size, int(value.shape[0])) - if kind == OutKind.UNSAFE_LINEAR: - c_value = merged.get("c") - d_value = merged.get("d") - if not isinstance(c_value, torch.Tensor) or not isinstance(d_value, torch.Tensor): - raise ValueError("UNSAFE_LINEAR test layer requires tensor c and d") - if c_value.dim() == 3: - batch_size = int(c_value.shape[0]) - m_rows = int(c_value.shape[1]) - merged["C"] = c_value.reshape(batch_size * m_rows, n_out) - elif c_value.dim() == 2: - m_rows = int(c_value.shape[0]) - merged["C"] = c_value - else: - raise ValueError(f"UNSAFE_LINEAR test c dim {c_value.dim()} unsupported") - merged["thresholds"] = d_value.reshape(batch_size, m_rows) - merged["M"] = m_rows - else: - merged["C"] = torch.zeros(batch_size, n_out) - merged["thresholds"] = torch.zeros(batch_size, 1) - merged["M"] = 1 - return Layer( - id=99, - kind=LayerKind.ASSERT.value, - params=merged, - in_vars=list(range(n_out)), - out_vars=list(range(n_out)), - ) diff --git a/act/back_end/bab/branching/bounding.py b/act/back_end/bab/branching/bounding.py index a783e68ad..110d92f56 100644 --- a/act/back_end/bab/branching/bounding.py +++ b/act/back_end/bab/branching/bounding.py @@ -69,6 +69,7 @@ from act.back_end.bab.node import SubproblemBatch from act.back_end.solver.solver_base import SolveStatus +from act.util.device_manager import get_default_dtype # --------------------------------------------------------------------------- @@ -384,6 +385,7 @@ def __init__( self._split_signs: Optional[Dict[int, torch.Tensor]] = None def push(self, batch: SubproblemBatch) -> None: + size_before = len(self) n_new = batch.batch_size device, dtype = batch.lb.device, batch.lb.dtype lower = ( @@ -409,6 +411,7 @@ def push(self, batch: SubproblemBatch) -> None: self._incremental_alpha = _clone_optional_dict(batch.incremental_alpha) self._incremental_eta = _clone_optional_dict(batch.incremental_eta) self._split_signs = _clone_optional_dict(batch.split_signs) + self._assert_push_transition(size_before, n_new) return assert prev_ub is not None and prev_depths is not None @@ -430,6 +433,7 @@ def push(self, batch: SubproblemBatch) -> None: if self._parent_id is not None: assert batch.parent_id is not None self._parent_id = torch.cat([self._parent_id, batch.parent_id.to(self._parent_id.device)], dim=0) + self._assert_push_transition(size_before, n_new) def pop(self, batch_size: int = 1) -> SubproblemBatch: lb = self._lb @@ -455,8 +459,40 @@ def pop(self, batch_size: int = 1) -> SubproblemBatch: self._clear() else: self._restrict(remaining) + self._assert_pop_transition(total, result.batch_size) return result + def _assert_push_transition(self, size_before: int, added: int) -> None: + size_after = len(self) + assert size_after == size_before + added, ( + "TOP-K BOUNDING invariant violated: push changed pool by wrong size " + f"(before={size_before}, batch_size={added}, after={size_after}, " + f"expected_after={size_before + added})" + ) + + def _assert_pop_transition(self, size_before: int, popped: int) -> None: + size_after = len(self) + cap_ok = self.k <= 0 or popped <= self.k + bookkeeping_ok = size_before - popped == size_after + assert cap_ok and bookkeeping_ok, ( + "TOP-K BOUNDING invariant violated: pop cap or pool bookkeeping failed " + f"(before={size_before}, popped={popped}, after={size_after}, " + f"k={self.k}, expected_after={size_before - popped})" + ) + + def view_all(self) -> SubproblemBatch: + """Return a lossless, non-destructive view of the full frontier.""" + if self._lb is None: + raise IndexError("view of empty pool") + indices = torch.arange(self._lb.shape[0], device=self._lb.device) + return self._build(indices) + + def replace_all(self, batch: SubproblemBatch) -> None: + """Replace the full frontier without scoring or advancing schedules.""" + self._clear() + if batch.batch_size > 0: + self.push(batch) + def _priority_scores(self) -> torch.Tensor: depths_t, lb = self._depths, self._lower_bound assert depths_t is not None and lb is not None @@ -579,6 +615,7 @@ def pop(self, batch_size: int = 1) -> SubproblemBatch: self._clear() else: self._restrict(remaining) + self._assert_pop_transition(total, result.batch_size) return result def _dedup_select(self, order: torch.Tensor, n: int) -> torch.Tensor: @@ -688,7 +725,9 @@ def _split_sign_features(self) -> Optional[torch.Tensor]: value = split_signs[layer_id] if value.shape[0] == 0: continue - pieces.append(value.detach().reshape(value.shape[0], -1).to(dtype=torch.float32)) + pieces.append( + value.detach().reshape(value.shape[0], -1).to(dtype=get_default_dtype()) + ) if not pieces: return None features = torch.cat(pieces, dim=1) @@ -699,7 +738,11 @@ def _split_sign_features(self) -> Optional[torch.Tensor]: def _box_center_features(self) -> Optional[torch.Tensor]: if self._lb is None or self._ub is None: return None - centers = ((self._lb + self._ub) / 2.0).detach().to(dtype=torch.float32) + centers = ( + ((self._lb + self._ub) / 2.0) + .detach() + .to(dtype=get_default_dtype()) + ) if centers.ndim > 2: centers = centers.reshape(centers.shape[0], -1) if centers.shape[1] == 0: @@ -710,6 +753,47 @@ def _box_center_features(self) -> Optional[torch.Tensor]: ROOT_PARENT = -1 +def _mcts_visit_accounting_valid( + parent: Dict[int, int], visits: Dict[int, int], n_tot: int +) -> bool: + if n_tot < 0 or any(count < 0 for count in visits.values()): + return False + child_visit_sums: Dict[int, int] = {} + for node, count in visits.items(): + parent_id = parent.get(node, ROOT_PARENT) + if parent_id != ROOT_PARENT: + child_visit_sums[parent_id] = child_visit_sums.get(parent_id, 0) + count + if any(visits.get(node, 0) < total for node, total in child_visit_sums.items()): + return False + root_visits = sum( + count + for node, count in visits.items() + if parent.get(node, ROOT_PARENT) == ROOT_PARENT + ) + return root_visits == n_tot + + +def _mcts_visit_diagnostics( + parent: Dict[int, int], visits: Dict[int, int] +) -> tuple[int, int, int]: + child_visit_sums: Dict[int, int] = {} + for node, count in visits.items(): + parent_id = parent.get(node, ROOT_PARENT) + if parent_id != ROOT_PARENT: + child_visit_sums[parent_id] = child_visit_sums.get(parent_id, 0) + count + root_visits = sum( + count + for node, count in visits.items() + if parent.get(node, ROOT_PARENT) == ROOT_PARENT + ) + min_visit = min(visits.values(), default=0) + max_child_excess = max( + (total - visits.get(node, 0) for node, total in child_visit_sums.items()), + default=0, + ) + return root_visits, min_visit, max_child_excess + + class MCTSBounding(BoundingStrategy): """MCTS side tables (``N``/``Q``) over the BaB tree, maintained as a pure observer. @@ -896,8 +980,8 @@ def observe( ``SAT`` status here always denotes a spurious, still-unresolved lane. """ ids = node_ids.detach().cpu() - lb = lower_bounds.detach().to(device="cpu", dtype=torch.float64) - depth = depths.detach().to(device="cpu", dtype=torch.float64) + lb = lower_bounds.detach().cpu().to(dtype=get_default_dtype()) + depth = depths.detach().cpu().to(dtype=get_default_dtype()) blended = self.lambda_ * depth / max(n_unstable, 1) + ( 1.0 - self.lambda_ ) * self._rank01(ids, lb) @@ -918,6 +1002,13 @@ def observe( while valued != ROOT_PARENT and self.Q.get(valued, -math.inf) < reward: self.Q[valued] = reward valued = self.parent[valued] + assert _mcts_visit_accounting_valid(self.parent, self.N, self.n_tot), ( + "MCTS invariant violated: visit-tree accounting is inconsistent " + f"(n_tot={self.n_tot}, root_visits=" + f"{_mcts_visit_diagnostics(self.parent, self.N)[0]}, min_visit=" + f"{_mcts_visit_diagnostics(self.parent, self.N)[1]}, " + f"max_child_excess={_mcts_visit_diagnostics(self.parent, self.N)[2]})" + ) def frontier_parent_visit_histogram(self) -> Dict[int, int]: parent_ids = self._parent_id diff --git a/act/back_end/bab/branching/branching.py b/act/back_end/bab/branching/branching.py index d166bb20e..2e9e6e78e 100644 --- a/act/back_end/bab/branching/branching.py +++ b/act/back_end/bab/branching/branching.py @@ -14,11 +14,8 @@ # 1. Input splitting — bisect (or N-ary section) the input domain along a dim. # 2. Neuron splitting — fix an unstable ReLU's activation phase (on/off). # -# Joint multi-neuron splitting (``_multi_split_from_decision``): split each lane's -# top-k BaBSR-scored neurons together into all 2^k sign combinations (one "skip" -# instead of k greedy single splits); gains are super-additive. -# * "Mining Verdict Boundaries for Neural Network Verification", -# Jiawei Ren, Guanqin Zhang, Zhenya Zhang, Yulei Sui, FM 2026. +# Joint multi-neuron splitting (2^k sign combinations of each lane's top-k +# BaBSR neurons, FM 2026) lives in ``act.back_end.bab.branching.multi_split``. # # Strategies (subclasses of ``BranchingStrategy``): # * ``RandomBranching`` — uniform-random over eligible dims (width-weighted for @@ -35,7 +32,7 @@ # Decomposition", De Palma et al., 2021. # # Result types: ``BranchingScores`` (per-dim / per-neuron scores) and -# ``SplitDecision`` (``input_axis`` / ``cut_dim`` + ``fanout`` for input splits; +# ``SplitDecision`` (``input_axis`` + ``fanout`` for input splits; # ``layer_id`` / ``neuron_idx`` for neuron splits). Factory: # ``_build_branching_strategy``. # @@ -47,13 +44,24 @@ from abc import ABC, abstractmethod from dataclasses import dataclass +import math from typing import Any, Dict, List, Optional, Tuple import torch -from act.back_end.bab.node import SubproblemBatch +from act.config.config import DualConfig +from act.back_end.bab.node import ( + SubproblemBatch, + _layer_neuron_count, + split_input, + split_input_nary, + split_neurons, +) from act.back_end.core import Bounds, Layer, Net +from act.back_end.dual_tf.tf_forward import compute_forward_bounds +from act.back_end.layer_schema import LayerKind from act.front_end.specs import InKind +from act.util.device_manager import get_default_device, get_default_dtype # --------------------------------------------------------------------------- @@ -66,17 +74,92 @@ class BranchingScores: flat: Optional[torch.Tensor] = None per_layer: Optional[Dict[int, torch.Tensor]] = None intercept_per_layer: Optional[Dict[int, torch.Tensor]] = None + input_fallback: Optional[torch.Tensor] = None + """Width-masked input-axis scores kept alongside neuron scores, so a lane + with no splittable neuron falls back to a *wide* input dimension instead of + a hard-coded dim 0 (zero width on any sparse-perturbation benchmark).""" @dataclass class SplitDecision: kind: str input_axis: Optional[torch.Tensor | int] = None - cut_dim: Optional[torch.Tensor] = None fanout: int = 2 layer_id: Optional[torch.Tensor] = None neuron_idx: Optional[torch.Tensor] = None + def input_axes(self, batch: SubproblemBatch) -> torch.Tensor: + if self.input_axis is None: + raise ValueError("input-axis decision missing input_axis") + input_axis = torch.as_tensor( + self.input_axis, device=batch.lb.device, dtype=torch.long + ).reshape(-1) + if input_axis.numel() == 1: + input_axis = input_axis.expand(batch.batch_size) + if input_axis.numel() != batch.batch_size: + raise ValueError( + f"input-axis decision has {input_axis.numel()} lanes for batch size " + f"{batch.batch_size}" + ) + return input_axis.contiguous() + + def apply( + self, + batch: SubproblemBatch, + net: Net, + ) -> tuple[SubproblemBatch, torch.Tensor]: + fanout = max(2, int(self.fanout)) + if self.kind == "input_axis": + dims = self.input_axes(batch) + if fanout == 2: + return split_input(batch, dims) + return split_input_nary(batch, dims, fanout) + + if self.kind == "neuron": + if self.layer_id is None or self.neuron_idx is None: + raise ValueError("neuron decision missing layer_id or neuron_idx") + + layer_id_tensor = self.layer_id.reshape(-1) + neuron_idx_tensor = self.neuron_idx.reshape(-1) + if layer_id_tensor.numel() == 0 or neuron_idx_tensor.numel() == 0: + raise ValueError("neuron decision tensors must be non-empty") + + representative_layer_id = int(layer_id_tensor[0].item()) + if representative_layer_id < 0: + raise ValueError( + f"neuron decision has negative layer_id {representative_layer_id}; " + "input-axis splits must be emitted as " + "SplitDecision(kind='input_axis')" + ) + + n_lanes = batch.batch_size + layer_ids = ( + layer_id_tensor.expand(n_lanes) + if layer_id_tensor.numel() == 1 + else layer_id_tensor + ) + neuron_indices = ( + neuron_idx_tensor.expand(n_lanes) + if neuron_idx_tensor.numel() == 1 + else neuron_idx_tensor + ) + if layer_ids.numel() != n_lanes or neuron_indices.numel() != n_lanes: + raise ValueError( + f"neuron decision has {layer_ids.numel()}/{neuron_indices.numel()} " + f"entries for batch size {n_lanes}" + ) + + return split_neurons( + batch, + net, + layer_ids.unsqueeze(1), + neuron_indices.unsqueeze(1), + 1, + first_sign=+1.0, + ) + + raise ValueError(f"Unknown SplitDecision.kind: {self.kind!r}") + # --------------------------------------------------------------------------- # Abstract base @@ -88,9 +171,9 @@ class BranchingStrategy(ABC): Lifecycle (called by the BaB engine per iteration):: - scores = strategy.compute_scores(batch, net, unstable_mask) - split_dims = strategy.select(scores) - left, right = split_subproblems(batch, split_dims) + scores = strategy.compute_scores(batch, net, unstable_mask) + decision = strategy.select(scores) + children, parent_index = decision.apply(batch, net) Subclass contract ~~~~~~~~~~~~~~~~~ @@ -139,7 +222,27 @@ def select(self, scores: torch.Tensor | BranchingScores) -> torch.Tensor | Split if scores.flat is None: raise ValueError("Base BranchingStrategy requires flat scores") scores = scores.flat - return scores.argmax(dim=-1) + return _argmax_splittable(scores) + + +def _argmax_splittable(scores: torch.Tensor) -> torch.Tensor: + """Row-wise ``argmax`` that refuses a masked-out (``-inf``) winner. + + Input-axis scorers mask zero-width dimensions to ``-inf``. A lane whose best + score is still ``-inf`` has nothing to split: bisecting a zero-width + dimension yields two children identical to the parent, so the frontier keeps + re-expanding them and BaB livelocks. Failing loudly is the sound option — + silently emitting no children would let ``bab.py`` drain the pool and report + CERTIFIED. + """ + best = scores.argmax(dim=-1) + best_score = scores.gather(-1, best.unsqueeze(-1)).squeeze(-1) + if not bool(torch.isfinite(best_score).all().item()): + raise ValueError( + "no splittable input dimension: every candidate dimension is " + "zero-width or masked out in at least one lane" + ) + return best # --------------------------------------------------------------------------- @@ -181,10 +284,11 @@ def compute_scores( embedding_mask = _perturbed_embedding_input_mask(net, batch) widths = batch.widths() # (N, D) if embedding_mask is None: - scores = torch.rand(N, D, device=device) * (widths > 0).float() + scores = torch.rand(N, D, device=device) else: mask = embedding_mask.unsqueeze(0).expand(N, -1) scores = widths.masked_fill(~mask, float("-inf")) + scores = scores.masked_fill(widths <= 0, float("-inf")) return scores @@ -205,11 +309,12 @@ def compute_scores( unstable_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: widths = batch.widths() + scores = widths.masked_fill(widths <= 0, float("-inf")) embedding_mask = _perturbed_embedding_input_mask(net, batch) if embedding_mask is not None: mask = embedding_mask.unsqueeze(0).expand_as(widths) - return widths.masked_fill(~mask, float("-inf")) - return widths + return scores.masked_fill(~mask, float("-inf")) + return scores # --------------------------------------------------------------------------- @@ -300,6 +405,7 @@ def compute_scores( flat=None, per_layer=per_layer, intercept_per_layer=intercept_per_layer, + input_fallback=self._baseline_scores(batch, None, net), ) def _baseline_scores( @@ -324,7 +430,8 @@ def _baseline_scores( mask = mask.unsqueeze(0).expand(batch.batch_size, -1) scores = scores * mask - return scores + # Masked last: -inf * 0 from the unstable mask would be NaN. + return scores.masked_fill(widths <= 0, float("-inf")) def _already_split_mask( self, @@ -341,14 +448,25 @@ def _already_split_mask( signs = batch.split_signs[lid] return (signs != 0).any(dim=1) + @staticmethod + def _input_axis_fallback(scores: BranchingScores) -> SplitDecision: + if scores.input_fallback is None: + raise ValueError( + "no splittable input dimension: neuron branching fell back to an " + "input-axis split but no width scores were recorded" + ) + return SplitDecision( + kind="input_axis", input_axis=_argmax_splittable(scores.input_fallback) + ) + def select(self, scores: torch.Tensor | BranchingScores) -> torch.Tensor | SplitDecision: if isinstance(scores, torch.Tensor): - return scores.argmax(dim=-1) + return _argmax_splittable(scores) if scores.flat is not None: - return SplitDecision(kind="input_axis", input_axis=scores.flat.argmax(dim=-1)) + return SplitDecision(kind="input_axis", input_axis=_argmax_splittable(scores.flat)) per_layer = scores.per_layer if not per_layer: - return SplitDecision(kind="input_axis", input_axis=0) + return self._input_axis_fallback(scores) N = next(iter(per_layer.values())).shape[0] device = next(iter(per_layer.values())).device @@ -390,7 +508,7 @@ def select(self, scores: torch.Tensor | BranchingScores) -> torch.Tensor | Split continue self.icp_score_counter = 0 - return SplitDecision(kind="input_axis", input_axis=0) + return self._input_axis_fallback(scores) return SplitDecision(kind="neuron", layer_id=decisions_layer, neuron_idx=decisions_neuron) @@ -398,6 +516,45 @@ def select(self, scores: torch.Tensor | BranchingScores) -> torch.Tensor | Split _WITNESS_GAP_TOL = 1e-6 +def witness_relu_preactivations( + net: Net, + witness_input: torch.Tensor, + input_shape: tuple[int, ...], + dual_config: DualConfig, +) -> Optional[Dict[int, torch.Tensor]]: + """Return concrete per-ReLU pre-activations at a dual witness.""" + if witness_input.numel() == 0: + return None + x = witness_input + if input_shape and x.dim() == 2 and x.shape[1] == int(math.prod(input_shape)): + x = x.reshape(x.shape[0], *input_shape) + bounds_dict = None + try: + bounds_dict = compute_forward_bounds( + net, + x, + x, + forward_lin_max_perturbed=dual_config.forward_lin_max_perturbed, + ) + except ValueError: + # compute_forward_bounds raises ValueError for layer kinds it cannot + # propagate; those nets use BaBSR scoring. + bounds_dict = None + if bounds_dict is None: + return None + preactivations = { + layer_id: bounds.lb.flatten(start_dim=1) + for layer_id, bounds in bounds_dict.items() + if _layer_kind_upper(net.by_id[layer_id]) == LayerKind.RELU.value + } + return preactivations or None + + +def _layer_kind_upper(layer: Layer) -> str: + kind = layer.kind + return kind.value if isinstance(kind, LayerKind) else kind.upper() + + class WitnessResidualBranching(BaBSRBranching): """Branch on the realized ReLU relaxation gap at the dual witness. @@ -489,7 +646,12 @@ def compute_scores( self._record_babsr_disagreement( batch, net, unstable_mask, bounds_dict, nu_per_layer, per_layer ) - return BranchingScores(flat=None, per_layer=per_layer, intercept_per_layer=None) + return BranchingScores( + flat=None, + per_layer=per_layer, + intercept_per_layer=None, + input_fallback=self._baseline_scores(batch, None, net), + ) def _record_babsr_disagreement( self, @@ -548,6 +710,61 @@ def _best_pairs( return pairs, torch.isfinite(best_val) +def _assert_witness_residual_decision_unstable( + decision: SplitDecision, + bounds_dict: Optional[Dict[int, Bounds]], +) -> None: + """Check that every non-fallback witness decision splits an unstable ReLU.""" + if decision.kind != "neuron": + return + layer_ids = decision.layer_id + neuron_indices = decision.neuron_idx + violation: Optional[str] = None + if layer_ids is None or neuron_indices is None: + violation = ( + f"neuron decision missing tensors (has_layers={layer_ids is not None}, " + f"has_neurons={neuron_indices is not None})" + ) + elif bounds_dict is None: + violation = ( + f"missing bounds (lanes={layer_ids.numel()}, " + f"decisions={neuron_indices.numel()})" + ) + else: + layer_ids = layer_ids.reshape(-1) + neuron_indices = neuron_indices.reshape(-1) + for lane, (layer_raw, neuron_raw) in enumerate( + zip(layer_ids.tolist(), neuron_indices.tolist()) + ): + layer_id = int(layer_raw) + neuron = int(neuron_raw) + bounds = bounds_dict.get(layer_id) + if bounds is None: + violation = ( + f"decision has no layer bounds (lane={lane}, layer={layer_id}, " + f"neuron={neuron}, available_layers={len(bounds_dict)})" + ) + break + lb = bounds.lb.flatten(start_dim=1) + ub = bounds.ub.flatten(start_dim=1) + if lane >= lb.shape[0] or neuron >= lb.shape[1] or neuron < 0: + violation = ( + f"decision index out of range (lane={lane}, layer={layer_id}, " + f"neuron={neuron}, bound_lanes={lb.shape[0]}, " + f"bound_neurons={lb.shape[1]})" + ) + break + lower = float(lb[lane, neuron].item()) + upper = float(ub[lane, neuron].item()) + if not lower < 0.0 < upper: + violation = ( + f"non-fallback decision selected stable neuron (lane={lane}, " + f"layer={layer_id}, neuron={neuron}, lb={lower:g}, ub={upper:g})" + ) + break + assert violation is None, f"WITNESS RESIDUAL invariant violated: {violation}" + + def _preact_bias_of(net: Net, lid: int) -> torch.Tensor: layer = net.by_id[lid] n_neurons = len(layer.out_vars) @@ -560,8 +777,8 @@ def _preact_bias_of(net: Net, lid: int) -> torch.Tensor: for value in pred.params.values(): if isinstance(value, torch.Tensor): return torch.zeros(n_neurons, dtype=value.dtype, device=value.device) - dtype = torch.float32 - device = torch.device("cpu") + dtype = get_default_dtype() + device = get_default_device() for value in layer.params.values(): if isinstance(value, torch.Tensor): dtype = value.dtype @@ -582,7 +799,10 @@ def _perturbed_embedding_input_mask(net: Optional[Net], batch: SubproblemBatch) if net is None: return None for layer in net.layers: - if layer.kind != "INPUT_SPEC" or layer.params.get("kind") != InKind.LP_EMBEDDING: + if ( + layer.kind != LayerKind.INPUT_SPEC.value + or layer.params.get("kind") != InKind.LP_EMBEDDING + ): continue center = layer.params.get("center") if not isinstance(center, torch.Tensor) or center.dim() < 2: @@ -614,21 +834,6 @@ def _perturbed_embedding_input_mask(net: Optional[Net], batch: SubproblemBatch) return None -def attention_nu_exposure_candidates( - bounds_dict: Optional[Dict[int, Bounds]], - nu_per_layer: Optional[Dict[int, torch.Tensor]], -) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Deferred hook for future attention-neuron splitting. - - Attention BaBSR needs a stable exposure of catalytic attention ν and - pre-activation intervals. Phase P4 intentionally ships input splitting - only, so this hook documents the integration point and returns no - candidates until that exposure contract exists. - """ - del bounds_dict, nu_per_layer - return None - - class FSBBranching(BaBSRBranching): def __init__( self, @@ -699,30 +904,21 @@ def compute_scores( N = batch.batch_size baseline = batch.parent_margins - try: - result = self._evaluate_hypotheses(net, bounds_dict, hypothesis_list) - stacked_margins = result.margins.to(device=batch.lb.device, dtype=batch.lb.dtype) - if stacked_margins.dim() == 1: - stacked_margins = stacked_margins.reshape(1, -1) - if stacked_margins.dim() > 2: - stacked_margins = stacked_margins.reshape(stacked_margins.shape[0], N, -1).mean(dim=-1) - if stacked_margins.shape[0] != len(hypothesis_list): - raise ValueError("KFSB solver returned wrong candidate dimension") - if stacked_margins.shape[1] != N: - stacked_margins = stacked_margins.reshape(len(hypothesis_list), N, -1).mean(dim=-1) - improvements = ( - stacked_margins - baseline.to(stacked_margins.device).unsqueeze(0) - if baseline is not None - else stacked_margins - ) - except Exception: - improvements = self._evaluate_hypotheses_serial( - net, - bounds_dict, - hypothesis_list, - baseline, - N, - ) + result = self._evaluate_hypotheses(net, bounds_dict, hypothesis_list) + stacked_margins = result.margins.to(device=batch.lb.device, dtype=batch.lb.dtype) + if stacked_margins.dim() == 1: + stacked_margins = stacked_margins.reshape(1, -1) + if stacked_margins.dim() > 2: + stacked_margins = stacked_margins.reshape(stacked_margins.shape[0], N, -1).mean(dim=-1) + if stacked_margins.shape[0] != len(hypothesis_list): + raise ValueError("KFSB solver returned wrong candidate dimension") + if stacked_margins.shape[1] != N: + stacked_margins = stacked_margins.reshape(len(hypothesis_list), N, -1).mean(dim=-1) + improvements = ( + stacked_margins - baseline.to(stacked_margins.device).unsqueeze(0) + if baseline is not None + else stacked_margins + ) if bool(torch.isneginf(improvements).all().item()): return bsr @@ -742,7 +938,12 @@ def compute_scores( final_per_layer[lid] = torch.full_like(source, float("-inf")) final_per_layer[lid][lane, neuron_idx] = improvements[ci, lane] - return BranchingScores(flat=None, per_layer=final_per_layer, intercept_per_layer=bsr.intercept_per_layer) + return BranchingScores( + flat=None, + per_layer=final_per_layer, + intercept_per_layer=bsr.intercept_per_layer, + input_fallback=bsr.input_fallback, + ) def _evaluate_hypotheses( self, @@ -762,54 +963,15 @@ def _evaluate_hypotheses( optimize=False, ) - def _evaluate_hypotheses_serial( - self, - net: Net, - bounds_dict: Optional[Dict[int, Bounds]], - hypothesis_list: List[Dict[int, torch.Tensor]], - baseline: Optional[torch.Tensor], - N: int, - ) -> torch.Tensor: - improvements = torch.full( - (len(hypothesis_list), N), - float("-inf"), - device=next(iter(hypothesis_list[0].values())).device, - dtype=next(iter(hypothesis_list[0].values())).dtype, - ) - for c_idx, hypo in enumerate(hypothesis_list): - try: - result = self._evaluate_hypothesis(net, bounds_dict, hypo) - margins = result.margins.reshape(N, -1).mean(dim=-1).to(improvements) - improvements[c_idx] = margins - baseline.to(margins.device) if baseline is not None else margins - except Exception: - # Recoverable: an individual hypothetical split may be incompatible with a legacy solver/mock. - continue - return improvements - - def _evaluate_hypothesis( - self, - net: Net, - bounds_dict: Optional[Dict[int, Bounds]], - split_signs: Dict[int, torch.Tensor], - ) -> Any: - if bounds_dict is None: - raise ValueError("FSB hypothesis evaluation requires bounds_dict") - c = self._zero_objective(net, bounds_dict) - return self.dual_solver.compute_certified_bound( - net, - bounds_dict, - c, - M=1, - split_signs=split_signs, - optimize=False, - ) def _zero_objective( self, net: Net, bounds_dict: Dict[int, Bounds], ) -> torch.Tensor: - assert_layers = [layer for layer in net.layers if layer.kind == "ASSERT"] + assert_layers = [ + layer for layer in net.layers if layer.kind == LayerKind.ASSERT.value + ] if assert_layers: assert_layer = assert_layers[-1] preds = net.preds.get(assert_layer.id, []) @@ -827,6 +989,7 @@ def _clone_split_signs_with_hypothesis( neuron_idx_per_lane: torch.Tensor, net: Net, ) -> Dict[int, torch.Tensor]: + # This is a one-sided +1 scoring hypothesis, not a partition into children. hypo: Dict[int, torch.Tensor] = {} if batch.split_signs is not None: for key, value in batch.split_signs.items(): @@ -875,343 +1038,3 @@ def _build_branching_strategy( raise ValueError("FSB branching requires a dual_solver instance (inject via factory).") return FSBBranching(dual_solver=dual_solver, branching_candidates=branching_candidates) raise ValueError(f"Unknown branching method: {method!r}") - - -# --------------------------------------------------------------------------- -# Joint multi-neuron splitting (verdict-boundary) -# -# "Mining Verdict Boundaries for Neural Network Verification" -# Jiawei Ren, Guanqin Zhang, Zhenya Zhang, Yulei Sui -# FM 2026 -# --------------------------------------------------------------------------- - - -def _collect_neuron_candidates( - branch_batch: SubproblemBatch, - bounds_dict: Dict[int, Bounds], - nu_per_layer: Dict[int, torch.Tensor], -) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Per-lane BaBSR scores (area x |nu|) over all splittable neurons. - - Returns ``(scores, layer_ids, neuron_ids)``, each ``[K, C]`` over the - concatenated candidate axis; stable or already-split neurons score -inf. - """ - kb = branch_batch.batch_size - device = branch_batch.lb.device - cand_layers: List[torch.Tensor] = [] - cand_neurons: List[torch.Tensor] = [] - cand_scores: List[torch.Tensor] = [] - for lid, nut in nu_per_layer.items(): - b = bounds_dict.get(lid) - if b is None: - continue - lb = b.lb.flatten(start_dim=1) - ub = b.ub.flatten(start_dim=1) - if lb.shape[0] != kb: - continue - n = min(lb.shape[-1], nut.shape[-1]) - amb = (lb[:, :n] < 0) & (ub[:, :n] > 0) - already = branch_batch.split_signs.get(lid) if branch_batch.split_signs else None - if already is not None: - amb &= already[:, 0, :n].to(device) == 0 - area = (-lb[:, :n] * ub[:, :n] / (ub[:, :n] - lb[:, :n]).clamp(min=1e-12)).clamp(min=0) - nv = nut.reshape(kb, -1, nut.shape[-1])[:, :, :n].abs().sum(dim=1) - sc = torch.where(amb, area * nv, torch.full_like(area, float("-inf"))) - cand_scores.append(sc) - cand_layers.append(torch.full((kb, n), lid, device=device, dtype=torch.long)) - cand_neurons.append( - torch.arange(n, device=device, dtype=torch.long).expand(kb, n) - ) - if not cand_scores: - return None - return ( - torch.cat(cand_scores, dim=1), - torch.cat(cand_layers, dim=1), - torch.cat(cand_neurons, dim=1), - ) - - -def enumerate_unstable_candidates( - branch_batch: SubproblemBatch, - bounds_dict: Optional[Dict[int, Bounds]], - nu_per_layer: Optional[Dict[int, torch.Tensor]], - *, - limit: Optional[int] = None, -) -> List[Dict[str, Any]]: - if bounds_dict is None or nu_per_layer is None: - return [] - cand = _collect_neuron_candidates(branch_batch, bounds_dict, nu_per_layer) - if cand is None: - return [] - scores, layers, neurons = cand - finite_mask = torch.isfinite(scores) - total = int(finite_mask.sum().item()) - if total == 0 or (limit is not None and total > int(limit)): - return [] - out: List[Dict[str, Any]] = [] - for lane in range(scores.shape[0]): - flat_cache: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} - for col in torch.where(finite_mask[lane])[0].tolist(): - lid = int(layers[lane, col].item()) - nidx = int(neurons[lane, col].item()) - b = bounds_dict.get(lid) - if b is not None and lid not in flat_cache: - flat_cache[lid] = (b.lb.flatten(start_dim=1), b.ub.flatten(start_dim=1)) - lb = float(flat_cache[lid][0][lane, nidx].item()) if lid in flat_cache else 0.0 - ub = float(flat_cache[lid][1][lane, nidx].item()) if lid in flat_cache else 0.0 - denom = ub - lb - area = float(max(0.0, (-lb * ub) / denom)) if denom > 1e-12 else 0.0 - score = float(scores[lane, col].item()) - out.append({ - "lane": lane, "layer_id": lid, "neuron_idx": nidx, - "score": score, "lb": lb, "ub": ub, - "nu": (score / area) if area > 1e-12 else None, "area": area, - }) - return out - - -def _layer_neuron_count(layer: Layer) -> int: - """Width of a layer's output block. - - Derived from the variable-id span rather than ``len(out_vars)`` so that the - sign tensors stay aligned with the solver's flat variable indexing even if a - layer ever declares a non-contiguous ``out_vars`` list. - """ - span = int(layer.out_vars[-1] - layer.out_vars[0] + 1) - assert span == len(layer.out_vars), ( - f"layer {layer.id} has non-contiguous out_vars: span {span} != " - f"{len(layer.out_vars)} declared variables" - ) - return span - - -def _index_subproblem_batch( - batch: SubproblemBatch, lanes: torch.Tensor -) -> SubproblemBatch: - """Restrict every field of ``batch`` to the given lane indices.""" - - def _pick(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: - return None if t is None else t.index_select(0, lanes.to(t.device)) - - def _pick_dict( - state: Optional[Dict[int, torch.Tensor]] - ) -> Optional[Dict[int, torch.Tensor]]: - if state is None: - return None - return {l: t.index_select(0, lanes.to(t.device)) for l, t in state.items()} - - return SubproblemBatch( - lb=batch.lb.index_select(0, lanes.to(batch.lb.device)), - ub=batch.ub.index_select(0, lanes.to(batch.ub.device)), - depths=batch.depths.index_select(0, lanes.to(batch.depths.device)), - incremental_alpha=_pick_dict(batch.incremental_alpha), - incremental_eta=_pick_dict(batch.incremental_eta), - split_signs=_pick_dict(batch.split_signs), - parent_margins=_pick(batch.parent_margins), - lower_bound=_pick(batch.lower_bound), - node_id=_pick(batch.node_id), - parent_id=_pick(batch.parent_id), - ) - - -def _concat_subproblem_batches( - a: SubproblemBatch, b: SubproblemBatch -) -> SubproblemBatch: - """Stack two child batches that were split at different depths. - - The per-layer state dicts may disagree on their key sets: a layer touched - only by one group's selection exists there and not in the other. A missing - key means "no neuron of this layer is constrained", which is exactly what an - all-zero sign block encodes, so padding with zeros preserves each side's - feasible region rather than widening or narrowing it. - """ - - def _cat(x: Optional[torch.Tensor], y: Optional[torch.Tensor]) -> Optional[torch.Tensor]: - if x is None or y is None: - return None - return torch.cat([x, y.to(device=x.device, dtype=x.dtype)], dim=0) - - def _cat_dict( - x: Optional[Dict[int, torch.Tensor]], - y: Optional[Dict[int, torch.Tensor]], - n_x: int, - n_y: int, - ) -> Optional[Dict[int, torch.Tensor]]: - if not x and not y: - return None - x = x or {} - y = y or {} - out: Dict[int, torch.Tensor] = {} - for lid in sorted(set(x) | set(y)): - present = x.get(lid) - if present is None: - present = y[lid] - trailing = present.shape[1:] - left = x.get(lid) - if left is None: - left = torch.zeros( - (n_x, *trailing), dtype=present.dtype, device=present.device - ) - right = y.get(lid) - if right is None: - right = torch.zeros( - (n_y, *trailing), dtype=present.dtype, device=present.device - ) - out[lid] = torch.cat([left, right.to(device=left.device)], dim=0) - return out - - n_a, n_b = a.batch_size, b.batch_size - return SubproblemBatch( - lb=torch.cat([a.lb, b.lb.to(a.lb.device)], dim=0), - ub=torch.cat([a.ub, b.ub.to(a.ub.device)], dim=0), - depths=torch.cat([a.depths, b.depths.to(a.depths.device)], dim=0), - incremental_alpha=_cat_dict(a.incremental_alpha, b.incremental_alpha, n_a, n_b), - incremental_eta=_cat_dict(a.incremental_eta, b.incremental_eta, n_a, n_b), - split_signs=_cat_dict(a.split_signs, b.split_signs, n_a, n_b), - parent_margins=_cat(a.parent_margins, b.parent_margins), - lower_bound=_cat(a.lower_bound, b.lower_bound), - node_id=_cat(a.node_id, b.node_id), - parent_id=_cat(a.parent_id, b.parent_id), - ) - - -def _multi_split_from_groups( - batch: SubproblemBatch, - net: Net, - top_layers: torch.Tensor, - top_neurons: torch.Tensor, - k_eff: int, -) -> Tuple[SubproblemBatch, torch.Tensor]: - # Soundness: the 2^k_eff sign-combination children exactly partition each - # lane's region for ANY (top_layers, top_neurons) — BaBSR- or LLM-chosen. - n_lanes = batch.batch_size - n_children = 2 ** k_eff - device = batch.lb.device - parent_index = torch.arange(n_lanes, device=device).repeat(n_children) - - def _gather(state: Optional[Dict[int, torch.Tensor]]) -> Optional[Dict[int, torch.Tensor]]: - if state is None: - return None - return { - l: t.index_select(0, parent_index.to(t.device)) for l, t in state.items() - } - - m_specs = 1 - if batch.incremental_alpha: - m_specs = int(next(iter(batch.incremental_alpha.values())).shape[1]) - elif batch.split_signs: - m_specs = int(next(iter(batch.split_signs.values())).shape[1]) - - signs = _gather(batch.split_signs) or {} - for lid_val in torch.unique(top_layers).tolist(): - lid_int = int(lid_val) - layer = net.by_id[lid_int] - n_neurons = _layer_neuron_count(layer) - if lid_int not in signs: - signs[lid_int] = torch.zeros( - n_children * n_lanes, m_specs, n_neurons, - device=device, dtype=batch.lb.dtype, - ) - else: - signs[lid_int] = signs[lid_int].clone() - for bit in range(k_eff): - lane_sel = torch.where(top_layers[:, bit] == lid_val)[0] - if lane_sel.numel() == 0: - continue - neuron_sel = top_neurons[lane_sel, bit].to(device=device, dtype=torch.long) - for j in range(n_children): - sign_val = 1.0 if (j >> bit) & 1 else -1.0 - rows = j * n_lanes + lane_sel - # Two index tensors separated by a `:` slice broadcast against - # EACH OTHER to shape [len(rows), m_specs] — positional pairing - # rows[i] <-> neuron_sel[i], not an outer product. Do not rewrite - # as signs[lid_int][rows][:, :, neuron_sel] = ...: that indexes a - # copy, so the assignment silently no-ops. - signs[lid_int][rows, :, neuron_sel] = sign_val - - children = SubproblemBatch( - lb=batch.lb.index_select(0, parent_index), - ub=batch.ub.index_select(0, parent_index), - depths=batch.depths.index_select(0, parent_index) + k_eff, - incremental_alpha=_gather(batch.incremental_alpha), - incremental_eta=_gather(batch.incremental_eta), - split_signs=signs, - parent_margins=( - batch.parent_margins.index_select(0, parent_index) - if batch.parent_margins is not None - else None - ), - lower_bound=( - batch.lower_bound.index_select(0, parent_index) - if batch.lower_bound is not None - else None - ), - ) - return children, parent_index - - -def _multi_split_from_decision( - batch: SubproblemBatch, - net: Net, - bounds_dict: Optional[Dict[int, Bounds]], - nu_per_layer: Optional[Dict[int, torch.Tensor]], - k_levels: int, -) -> Optional[Tuple[SubproblemBatch, torch.Tensor]]: - """Joint top-k neuron split: each lane emits all 2^k sign combinations. - - Verdict-boundary multi-neuron splitting — "Mining Verdict Boundaries for - Neural Network Verification", Jiawei Ren, Guanqin Zhang, Zhenya Zhang, - Yulei Sui, FM 2026. Candidates are scored by the BaBSR heuristic - (``_collect_neuron_candidates``). - - The 2^k children exactly partition each lane's region (every selected - neuron is constrained to >=0 or <=0 in both directions across the - combination set), so replacing the lane by its children is sound. Joint - splits are super-additive in bound gain versus greedy single splits. - - Lanes are grouped by how many finite candidates they actually have, so a - single candidate-starved lane no longer collapses ``k_eff`` for the whole - wave: lanes with >=2 candidates take the joint split, lanes with exactly one - take a plain 2-child split, and the two groups are concatenated. Returns - None only when some lane has no splittable neuron at all, since neuron - branching cannot cover such a lane and the caller must fall back to the - decision path (which can still split an input axis). - """ - if bounds_dict is None or nu_per_layer is None: - return None - cand = _collect_neuron_candidates(batch, bounds_dict, nu_per_layer) - if cand is None: - return None - all_scores, all_layers, all_neurons = cand - finite_per_lane = torch.isfinite(all_scores).sum(dim=1) - if int(finite_per_lane.min().item()) < 1: - return None - k_lane = finite_per_lane.clamp(max=k_levels) - joint_mask = k_lane >= 2 - if not bool(joint_mask.any().item()): - return None - - def _group(lanes: torch.Tensor, k: int) -> Tuple[SubproblemBatch, torch.Tensor]: - sub = _index_subproblem_batch(batch, lanes) - scores = all_scores.index_select(0, lanes) - top = torch.topk(scores, k=k, dim=1).indices - children, local_parent = _multi_split_from_groups( - sub, - net, - all_layers.index_select(0, lanes).gather(1, top), - all_neurons.index_select(0, lanes).gather(1, top), - k, - ) - return children, lanes.to(local_parent.device).index_select(0, local_parent) - - k_eff = min(k_levels, int(finite_per_lane[joint_mask].min().item())) - lanes_joint = torch.where(joint_mask)[0] - lanes_single = torch.where(~joint_mask)[0] - joint_children, joint_parent = _group(lanes_joint, k_eff) - if lanes_single.numel() == 0: - return joint_children, joint_parent - single_children, single_parent = _group(lanes_single, 1) - return ( - _concat_subproblem_batches(joint_children, single_children), - torch.cat([joint_parent, single_parent]), - ) diff --git a/act/back_end/bab/branching/multi_split.py b/act/back_end/bab/branching/multi_split.py new file mode 100644 index 000000000..402e05197 --- /dev/null +++ b/act/back_end/bab/branching/multi_split.py @@ -0,0 +1,391 @@ +# ===- act/back_end/bab/branching/multi_split.py - Joint Multi-Neuron Split -====# +# ACT: Abstract Constraint Transformer +# Copyright (C) 2025– ACT Team +# +# Licensed under the GNU Affero General Public License v3.0 or later (AGPLv3+). +# Distributed without any warranty; see . +# ===---------------------------------------------------------------------====# +# +# Purpose: +# Joint multi-neuron splitting (verdict-boundary) for Branch-and-Bound. +# +# Implements the 2^k sign-combination fan-out described in "Mining Verdict +# Boundaries for Neural Network Verification" (Jiawei Ren, Guanqin Zhang, +# Zhenya Zhang, Yulei Sui, FM 2026). Each lane's top-k BaBSR-scored neurons +# are split together into all 2^k sign combinations in a single wave; the +# joint gain is super-additive versus k greedy single splits. +# +# Entry points: +# ``presplit_root`` — choose root neurons for pre-splitting. +# ``gain_tested_decision`` — choose a measured-gain neuron split. +# ``groups_to_tensors`` — normalize controller neuron groups. +# ``enumerate_unstable_candidates`` — list splittable neurons with scores. +# ``_collect_neuron_candidates`` — per-lane BaBSR scores (internal). +# ``_multi_split_from_decision`` — top-level entry: score, select, split. +# +# ===---------------------------------------------------------------------====# + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple + +import torch + +from act.config.config import BaBConfig, DualConfig +from act.back_end.bab.branching.branching import SplitDecision +from act.back_end.bab.node import ( + SubproblemBatch, + _gather_optional_dict, + _layer_neuron_count, + split_neurons, +) +from act.back_end.core import Bounds, Layer, Net + +if TYPE_CHECKING: + from act.back_end.solver.solver_dual import DualBatchResult + + +# --------------------------------------------------------------------------- +# Joint multi-neuron splitting (verdict-boundary) +# +# "Mining Verdict Boundaries for Neural Network Verification" +# Jiawei Ren, Guanqin Zhang, Zhenya Zhang, Yulei Sui +# FM 2026 +# --------------------------------------------------------------------------- + + +def presplit_root( + root: SubproblemBatch, + net: Net, + bounds_dict: Dict[int, Bounds], + nu_per_layer: Dict[int, torch.Tensor], + k: int, +) -> Optional[SubproblemBatch]: + """Materialize the root descendants from one best layer's top-k neurons.""" + best: Optional[tuple[int, torch.Tensor, torch.Tensor]] = None + for layer_id, nu in nu_per_layer.items(): + bounds = bounds_dict.get(layer_id) + if bounds is None: + continue + lb = bounds.lb.flatten(start_dim=1)[0] + ub = bounds.ub.flatten(start_dim=1)[0] + n = min(lb.shape[-1], nu.shape[-1]) + lb, ub = lb[:n], ub[:n] + ambiguous = (lb < 0) & (ub > 0) + if not bool(ambiguous.any().item()): + continue + area = (-lb * ub / (ub - lb).clamp(min=1e-12)).clamp(min=0) + score = area * nu.reshape(-1, nu.shape[-1])[:, :n].abs().sum(dim=0) + score = torch.where(ambiguous, score, torch.zeros_like(score)) + if best is None or float(score.max()) > float(best[1].max()): + best = (layer_id, score, lb) + if best is None: + return None + layer_id, score, _ = best + k = min(k, int((score > 0).sum().item())) + if k < 1: + return None + top_idx = torch.topk(score, k=k).indices + assert score.shape[-1] == _layer_neuron_count(net.by_id[layer_id]) + children, _ = split_neurons( + root, + net, + torch.full( + (root.batch_size, k), + layer_id, + dtype=torch.long, + device=root.lb.device, + ), + top_idx.unsqueeze(0), + k, + ) + return children + + +def _collect_neuron_candidates( + branch_batch: SubproblemBatch, + bounds_dict: Dict[int, Bounds], + nu_per_layer: Dict[int, torch.Tensor], +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Per-lane BaBSR scores (area x |nu|) over all splittable neurons. + + Returns ``(scores, layer_ids, neuron_ids)``, each ``[K, C]`` over the + concatenated candidate axis; stable or already-split neurons score -inf. + """ + kb = branch_batch.batch_size + device = branch_batch.lb.device + cand_layers: List[torch.Tensor] = [] + cand_neurons: List[torch.Tensor] = [] + cand_scores: List[torch.Tensor] = [] + for lid, nut in nu_per_layer.items(): + b = bounds_dict.get(lid) + if b is None: + continue + lb = b.lb.flatten(start_dim=1) + ub = b.ub.flatten(start_dim=1) + if lb.shape[0] != kb: + continue + n = min(lb.shape[-1], nut.shape[-1]) + amb = (lb[:, :n] < 0) & (ub[:, :n] > 0) + already = branch_batch.split_signs.get(lid) if branch_batch.split_signs else None + if already is not None: + amb &= already[:, 0, :n].to(device) == 0 + area = (-lb[:, :n] * ub[:, :n] / (ub[:, :n] - lb[:, :n]).clamp(min=1e-12)).clamp(min=0) + nv = nut.reshape(kb, -1, nut.shape[-1])[:, :, :n].abs().sum(dim=1) + sc = torch.where(amb, area * nv, torch.full_like(area, float("-inf"))) + cand_scores.append(sc) + cand_layers.append(torch.full((kb, n), lid, device=device, dtype=torch.long)) + cand_neurons.append( + torch.arange(n, device=device, dtype=torch.long).expand(kb, n) + ) + if not cand_scores: + return None + return ( + torch.cat(cand_scores, dim=1), + torch.cat(cand_layers, dim=1), + torch.cat(cand_neurons, dim=1), + ) + + +def enumerate_unstable_candidates( + branch_batch: SubproblemBatch, + bounds_dict: Optional[Dict[int, Bounds]], + nu_per_layer: Optional[Dict[int, torch.Tensor]], + *, + limit: Optional[int] = None, +) -> List[Dict[str, Any]]: + if bounds_dict is None or nu_per_layer is None: + return [] + cand = _collect_neuron_candidates(branch_batch, bounds_dict, nu_per_layer) + if cand is None: + return [] + scores, layers, neurons = cand + finite_mask = torch.isfinite(scores) + total = int(finite_mask.sum().item()) + if total == 0 or (limit is not None and total > int(limit)): + return [] + out: List[Dict[str, Any]] = [] + for lane in range(scores.shape[0]): + flat_cache: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {} + for col in torch.where(finite_mask[lane])[0].tolist(): + lid = int(layers[lane, col].item()) + nidx = int(neurons[lane, col].item()) + b = bounds_dict.get(lid) + if b is not None and lid not in flat_cache: + flat_cache[lid] = (b.lb.flatten(start_dim=1), b.ub.flatten(start_dim=1)) + lb = float(flat_cache[lid][0][lane, nidx].item()) if lid in flat_cache else 0.0 + ub = float(flat_cache[lid][1][lane, nidx].item()) if lid in flat_cache else 0.0 + denom = ub - lb + area = float(max(0.0, (-lb * ub) / denom)) if denom > 1e-12 else 0.0 + score = float(scores[lane, col].item()) + out.append({ + "lane": lane, "layer_id": lid, "neuron_idx": nidx, + "score": score, "lb": lb, "ub": ub, + "nu": (score / area) if area > 1e-12 else None, "area": area, + }) + return out + + +def groups_to_tensors( + groups: Dict[int, Any], batch: SubproblemBatch +) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]: + # Duplicate pairs are sound but waste 2^k rows, so preserve order and dedupe. + batch_size = batch.batch_size + if len(groups) != batch_size: + return None, None, 0 + deduped: Dict[int, List[Tuple[int, int]]] = {} + for lane in range(batch_size): + seen: Dict[Tuple[int, int], None] = {} + for layer_id, neuron_index in groups.get(lane, []): + seen.setdefault((int(layer_id), int(neuron_index)), None) + deduped[lane] = list(seen) + k_eff = min((len(entries) for entries in deduped.values()), default=0) + if k_eff < 1: + return None, None, 0 + device = batch.lb.device + top_layers = torch.zeros( + batch_size, k_eff, dtype=torch.long, device=device + ) + top_neurons = torch.zeros( + batch_size, k_eff, dtype=torch.long, device=device + ) + for lane in range(batch_size): + for bit, (layer_id, neuron_index) in enumerate(deduped[lane][:k_eff]): + top_layers[lane, bit] = layer_id + top_neurons[lane, bit] = neuron_index + return top_layers, top_neurons, k_eff + + +def gain_tested_decision( + branch_batch: SubproblemBatch, + net: Net, + assert_layer: Layer, + config: BaBConfig, + dual_config: DualConfig, + keep_rows: Optional[torch.Tensor], + root_bounds_dict: Optional[Dict[int, Bounds]], + bounds_dict: Optional[Dict[int, Bounds]], + nu_per_layer: Optional[Dict[int, torch.Tensor]], + input_shape: tuple[int, ...], + solve_dual: Callable[..., DualBatchResult], + n_candidates: int = 3, +) -> Optional[SplitDecision]: + """Pick each lane's split by measured child bounds, not by score proxy.""" + if bounds_dict is None or nu_per_layer is None: + return None + batch_size = branch_batch.batch_size + device = branch_batch.lb.device + + candidates = _collect_neuron_candidates( + branch_batch, bounds_dict, nu_per_layer + ) + if candidates is None: + return None + all_scores, all_layers, all_neurons = candidates + n_selected = min(n_candidates, all_scores.shape[1]) + top = torch.topk(all_scores, k=n_selected, dim=1).indices + top_layers = all_layers.gather(1, top) + top_neurons = all_neurons.gather(1, top) + + # Keep this probe builder separate: rows are lane-major, depth is unchanged, + # and M comes from alpha (else 1) because solve_dual reads the probe depth. + repeated_indices = torch.arange(batch_size, device=device).repeat_interleave( + 2 * n_selected + ) + m_specs = 1 + if branch_batch.incremental_alpha: + m_specs = int( + next(iter(branch_batch.incremental_alpha.values())).shape[1] + ) + signs = _gather_optional_dict( + branch_batch.split_signs, repeated_indices + ) or {} + for layer_id_value in torch.unique(top_layers).tolist(): + layer_id = int(layer_id_value) + n_neurons = _layer_neuron_count(net.by_id[layer_id]) + if layer_id not in signs: + signs[layer_id] = torch.zeros( + 2 * n_selected * batch_size, + m_specs, + n_neurons, + device=device, + dtype=branch_batch.lb.dtype, + ) + else: + signs[layer_id] = signs[layer_id].clone() + for lane in range(batch_size): + for candidate in range(n_selected): + if int(top_layers[lane, candidate]) != layer_id: + continue + row = lane * 2 * n_selected + 2 * candidate + neuron = int(top_neurons[lane, candidate]) + signs[layer_id][row, :, neuron] = 1.0 + signs[layer_id][row + 1, :, neuron] = -1.0 + + probe = SubproblemBatch( + lb=branch_batch.lb.index_select(0, repeated_indices), + ub=branch_batch.ub.index_select(0, repeated_indices), + depths=branch_batch.depths.index_select(0, repeated_indices), + incremental_alpha=_gather_optional_dict( + branch_batch.incremental_alpha, repeated_indices + ), + incremental_eta=_gather_optional_dict( + branch_batch.incremental_eta, repeated_indices + ), + split_signs=signs, + ) + n_probe = probe.batch_size + probe_bounds = Bounds( + probe.lb.reshape(n_probe, *input_shape) if input_shape else probe.lb, + probe.ub.reshape(n_probe, *input_shape) if input_shape else probe.ub, + ) + result = solve_dual( + net=net, + assert_layer=assert_layer, + batched_bounds=probe_bounds, + k_actual=n_probe, + batch=probe, + config=config, + dual_config=dual_config, + optimize=False, + keep_rows=keep_rows, + root_bounds_dict=root_bounds_dict, + ) + child_lbs = (-result.solution.max_viol).view( + batch_size, n_selected, 2 + ) + pair_gain = child_lbs.min(dim=2).values + best_candidate = pair_gain.argmax(dim=1) + lane_indices = torch.arange(batch_size, device=device) + return SplitDecision( + kind="neuron", + layer_id=top_layers[lane_indices, best_candidate], + neuron_idx=top_neurons[lane_indices, best_candidate], + ) + + +def _multi_split_from_decision( + batch: SubproblemBatch, + net: Net, + bounds_dict: Optional[Dict[int, Bounds]], + nu_per_layer: Optional[Dict[int, torch.Tensor]], + k_levels: int, +) -> Optional[Tuple[SubproblemBatch, torch.Tensor]]: + """Joint top-k neuron split: each lane emits all 2^k sign combinations. + + Verdict-boundary multi-neuron splitting — "Mining Verdict Boundaries for + Neural Network Verification", Jiawei Ren, Guanqin Zhang, Zhenya Zhang, + Yulei Sui, FM 2026. Candidates are scored by the BaBSR heuristic + (``_collect_neuron_candidates``). + + The 2^k children exactly partition each lane's region (every selected + neuron is constrained to >=0 or <=0 in both directions across the + combination set), so replacing the lane by its children is sound. Joint + splits are super-additive in bound gain versus greedy single splits. + + Lanes are grouped by how many finite candidates they actually have, so a + single candidate-starved lane no longer collapses ``k_eff`` for the whole + wave: lanes with >=2 candidates take the joint split, lanes with exactly one + take a plain 2-child split, and the two groups are concatenated. Returns + None only when some lane has no splittable neuron at all, since neuron + branching cannot cover such a lane and the caller must fall back to the + decision path (which can still split an input axis). + """ + if bounds_dict is None or nu_per_layer is None: + return None + cand = _collect_neuron_candidates(batch, bounds_dict, nu_per_layer) + if cand is None: + return None + all_scores, all_layers, all_neurons = cand + finite_per_lane = torch.isfinite(all_scores).sum(dim=1) + if int(finite_per_lane.min().item()) < 1: + return None + k_lane = finite_per_lane.clamp(max=k_levels) + joint_mask = k_lane >= 2 + if not bool(joint_mask.any().item()): + return None + + def _group(lanes: torch.Tensor, k: int) -> Tuple[SubproblemBatch, torch.Tensor]: + sub = batch.select(lanes) + scores = all_scores.index_select(0, lanes) + top = torch.topk(scores, k=k, dim=1).indices + children, local_parent = split_neurons( + sub, + net, + all_layers.index_select(0, lanes).gather(1, top), + all_neurons.index_select(0, lanes).gather(1, top), + k, + ) + return children, lanes.to(local_parent.device).index_select(0, local_parent) + + k_eff = min(k_levels, int(finite_per_lane[joint_mask].min().item())) + lanes_joint = torch.where(joint_mask)[0] + lanes_single = torch.where(~joint_mask)[0] + joint_children, joint_parent = _group(lanes_joint, k_eff) + if lanes_single.numel() == 0: + return joint_children, joint_parent + single_children, single_parent = _group(lanes_single, 1) + return ( + joint_children.concat(single_children), + torch.cat([joint_parent, single_parent]), + ) diff --git a/act/back_end/bab/climb.py b/act/back_end/bab/climb.py new file mode 100644 index 000000000..a6c57ed2c --- /dev/null +++ b/act/back_end/bab/climb.py @@ -0,0 +1,1062 @@ +# ===- act/back_end/bab/climb.py - CLIMB Certificate Reuse ----------------====# +# ACT: Abstract Constraint Transformer +# Copyright (C) 2025– ACT Team +# +# Licensed under the GNU Affero General Public License v3.0 or later (AGPLv3+). +# Distributed without any warranty; see . +# ===---------------------------------------------------------------------====# +# +# Purpose: +# Query-local CLIMB certificate replay, core learning, and phase propagation. +# ClimbSession owns one verification run's CLIMB state and BaB-loop hooks. +# +# ===---------------------------------------------------------------------====# + +from __future__ import annotations + +import bisect +import time +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeAlias + +import torch + +from act.back_end.bab.branching.bounding import TopKBounding +from act.back_end.bab.node import ( + SubproblemBatch, + _layer_neuron_count, + slice_bounds_dict, +) +from act.back_end.core import Bounds, Net +from act.back_end.solver.solver_base import SolveStatus +from act.back_end.solver.solver_dual import DualBatchResult +from act.config.config import ( + BaBConfig, + CLIMB_SOLVER_TIER, + ConfigError, + NEURON_BRANCHING_METHODS, + NO_REFINEMENT_MODE, + TOP_K_BOUNDINGS, +) +from act.front_end.specs import OutKind +from act.util.device_manager import get_default_device, get_default_dtype + + +# --------------------------------------------------------------------------- +# CLIMB certificate reuse +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class LiteralCodec: + """Immutable DIMACS-style codec for the network's ReLU phase literals.""" + + layer_ids: Tuple[int, ...] + widths: Tuple[int, ...] + offsets: Tuple[int, ...] + + @classmethod + def from_net(cls, net: Net) -> LiteralCodec: + layers = sorted( + (layer for layer in net.layers if layer.kind == "RELU"), + key=lambda layer: layer.id, + ) + layer_ids = tuple(layer.id for layer in layers) + widths = tuple(_layer_neuron_count(layer) for layer in layers) + offsets: List[int] = [] + total = 0 + for width in widths: + offsets.append(total) + total += width + return cls(layer_ids, widths, tuple(offsets)) + + def __post_init__(self) -> None: + if len(self.layer_ids) != len(self.widths) or len(self.widths) != len(self.offsets): + raise ValueError("literal codec tables must have equal lengths") + if tuple(sorted(self.layer_ids)) != self.layer_ids: + raise ValueError("literal codec layer ids must be increasing") + expected = 0 + for width, offset in zip(self.widths, self.offsets): + if width < 0 or offset != expected: + raise ValueError("literal codec offsets must be contiguous") + expected += width + + def width(self, layer_id: int) -> int: + index = bisect.bisect_left(self.layer_ids, layer_id) + if index == len(self.layer_ids) or self.layer_ids[index] != layer_id: + raise KeyError(layer_id) + return self.widths[index] + + def encode(self, layer_id: int, neuron: int, sign: int) -> int: + index = bisect.bisect_left(self.layer_ids, layer_id) + if index == len(self.layer_ids) or self.layer_ids[index] != layer_id: + raise KeyError(layer_id) + if neuron < 0 or neuron >= self.widths[index]: + raise IndexError(neuron) + if sign not in (-1, 1): + raise ValueError("literal sign must be -1 or 1") + return sign * (self.offsets[index] + neuron + 1) + + def decode(self, lit: int) -> tuple[int, int, int]: + if lit == 0: + raise ValueError("zero is literal padding") + variable = abs(lit) - 1 + index = bisect.bisect_right(self.offsets, variable) - 1 + if index < 0 or variable >= self.offsets[index] + self.widths[index]: + raise IndexError(variable) + return self.layer_ids[index], variable - self.offsets[index], 1 if lit > 0 else -1 + + def encode_tensor( + self, layer_ids: torch.Tensor, neurons: torch.Tensor, signs: torch.Tensor + ) -> torch.Tensor: + table_ids = torch.tensor(self.layer_ids, dtype=torch.long, device=layer_ids.device) + table_offsets = torch.tensor(self.offsets, dtype=torch.long, device=layer_ids.device) + indices = torch.searchsorted(table_ids, layer_ids.to(torch.long)) + return signs.to(torch.long) * ( + table_offsets.index_select(0, indices) + neurons.to(torch.long) + 1 + ) + + def decode_tensor( + self, literals: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + variables = literals.abs() - 1 + ends = torch.tensor( + tuple(offset + width for offset, width in zip(self.offsets, self.widths)), + dtype=torch.long, + device=literals.device, + ) + offsets = torch.tensor(self.offsets, dtype=torch.long, device=literals.device) + layer_ids = torch.tensor(self.layer_ids, dtype=torch.long, device=literals.device) + indices = torch.searchsorted(ends, variables.clamp(min=0), right=True) + flat_indices = indices.reshape(-1) + return ( + layer_ids.index_select(0, flat_indices).reshape_as(literals), + variables - offsets.index_select(0, flat_indices).reshape_as(literals), + literals.sign(), + ) + + +Core: TypeAlias = frozenset[int] + + +def _core_has_no_complements(core: Core) -> bool: + return len(core) == len({abs(literal) for literal in core}) + + +def _counters_non_decreasing( + before: tuple[int, int, int], after: tuple[int, int, int] +) -> bool: + return all(current >= previous for previous, current in zip(before, after)) + + +def _retained_literals_are_subset( + original: torch.Tensor, retained: torch.Tensor +) -> bool: + return original.shape == retained.shape and bool( + ((retained == 0) | (retained == original)).all() + ) + + +def _only_fills_unassigned(initial: torch.Tensor, final: torch.Tensor) -> bool: + return initial.shape == final.shape and bool( + ((initial == 0) | (final == initial)).all() + ) + + +@dataclass +class ClimbMetrics: + enabled: bool = False + generated_children: int = 0 + main_bound_row_passes: int = 0 + main_bound_calls: int = 0 + peak_frontier_rows: int = 0 + prebound_discharged: int = 0 + presplit_discharged: int = 0 + recheck_row_passes: int = 0 + recheck_calls: int = 0 + recheck_time_s: float = 0.0 + coarsen_time_s: float = 0.0 + propagate_time_s: float = 0.0 + cores_inserted: int = 0 + core_literals_original: int = 0 + core_literals_retained: int = 0 + library: Optional[CoreLibrary] = None + + def observe_frontier(self, pending_plus_active_rows: int) -> None: + self.peak_frontier_rows = max(self.peak_frontier_rows, pending_plus_active_rows) + + def metadata(self) -> Dict[str, int | float]: + if not self.enabled: + return {} + return { + "climb_generated_children": self.generated_children, + "climb_main_bound_row_passes": self.main_bound_row_passes, + "climb_main_bound_calls": self.main_bound_calls, + "climb_peak_frontier_rows": self.peak_frontier_rows, + "climb_prebound_discharged": self.prebound_discharged, + "climb_presplit_discharged": self.presplit_discharged, + "climb_recheck_row_passes": self.recheck_row_passes, + "climb_recheck_calls": self.recheck_calls, + "climb_recheck_time_s": self.recheck_time_s, + "climb_coarsen_time_s": self.coarsen_time_s, + "climb_propagate_time_s": self.propagate_time_s, + "climb_cores_inserted": self.cores_inserted, + "climb_cores_subsumed": self.library.subsumed if self.library is not None else 0, + "climb_cores_resolved": self.library.resolved if self.library is not None else 0, + "climb_cores_evicted": self.library.evicted if self.library is not None else 0, + "climb_core_literals_original": self.core_literals_original, + "climb_core_literals_retained": self.core_literals_retained, + } + + +def pack_split_matrix( + split_signs: Optional[Dict[int, torch.Tensor]], + n_rows: int, + codec: LiteralCodec, + *, + device: Optional[torch.device] = None, +) -> torch.Tensor: + """Pack split-state dictionaries into zero-padded signed literal rows.""" + target_device = ( + next(iter(split_signs.values())).device + if split_signs + else (device if device is not None else get_default_device()) + ) + row_parts: List[torch.Tensor] = [] + literal_parts: List[torch.Tensor] = [] + count_parts: List[torch.Tensor] = [] + if split_signs: + for layer_id in sorted(split_signs): + values = split_signs[layer_id] + if values.shape[0] != n_rows: + raise ValueError("split_signs leading dimension mismatch") + flattened = values.reshape(n_rows, values.shape[1], -1) + first = flattened[:, :1, :] + if not torch.equal(flattened, first.expand_as(flattened)): + raise ValueError("CLIMB requires spec-invariant split signs") + width = flattened.shape[-1] + if codec.width(layer_id) != width: + raise ValueError("split_signs width does not match literal codec") + lane_values = first[:, 0, :].to(target_device) + positions = torch.nonzero(lane_values, as_tuple=False) + counts = torch.bincount(positions[:, 0], minlength=n_rows) + count_parts.append(counts) + row_parts.append(positions[:, 0]) + literal_parts.append( + codec.encode_tensor( + torch.full_like(positions[:, 0], layer_id), + positions[:, 1], + lane_values[positions[:, 0], positions[:, 1]].sign(), + ) + ) + if not count_parts: + return torch.zeros((n_rows, 0), dtype=torch.long, device=target_device) + total_counts = torch.stack(count_parts).sum(dim=0) + packed = torch.zeros( + (n_rows, int(total_counts.max())), dtype=torch.long, device=target_device + ) + layer_bases = torch.zeros(n_rows, dtype=torch.long, device=target_device) + for rows, literals, counts in zip(row_parts, literal_parts, count_parts): + starts = counts.cumsum(0) - counts + local_slots = torch.arange(rows.numel(), device=target_device) - torch.repeat_interleave( + starts, counts + ) + packed[rows, layer_bases[rows] + local_slots] = literals + layer_bases += counts + return packed + + +class CoreLibrary: + """Query-local certified phase cores with subsumption and resolution.""" + + def __init__(self, max_cores: int) -> None: + if max_cores < 1: + raise ValueError("max_cores must be positive") + self.max_cores = max_cores + self.subsumed = 0 + self.resolved = 0 + self.evicted = 0 + self._cores: List[Core] = [] + + @property + def core_count(self) -> int: + return len(self._cores) + + @property + def cores(self) -> Tuple[Core, ...]: + return tuple(self._cores) + + def _admit(self, literals: Core) -> int: + """Insert unless subsumed; drop strict supersets. Returns removed count.""" + assert _core_has_no_complements(literals), ( + "CLIMB invariant violated: admitted core contains complementary literals " + f"(core_size={len(literals)})" + ) + before = len(self._cores) + self._cores = [core for core in self._cores if not literals < core] + self._cores.append(literals) + return before + 1 - len(self._cores) + + def _is_subsumed(self, literals: Core) -> bool: + return any(core <= literals for core in self._cores) + + def insert(self, packed: torch.Tensor) -> int: + counters_before = (self.subsumed, self.resolved, self.evicted) + inserted = 0 + for lane in range(packed.shape[0]): + literals = frozenset(packed[lane, packed[lane] != 0].tolist()) + if self._is_subsumed(literals): + self.subsumed += 1 + continue + self.subsumed += self._admit(literals) + inserted += 1 + + while (resolvent := self._find_resolvent()) is not None: + self._admit(resolvent) + self.resolved += 1 + + self._cores.sort(key=len) + self.evicted += max(0, len(self._cores) - self.max_cores) + self._cores = self._cores[: self.max_cores] + counters_after = (self.subsumed, self.resolved, self.evicted) + assert _counters_non_decreasing(counters_before, counters_after), ( + "CLIMB invariant violated: core-library counters decreased " + f"(before={counters_before}, after={counters_after})" + ) + return inserted + + def _find_resolvent(self) -> Optional[Core]: + """First pair differing only by one opposite pivot sign, if any.""" + snapshot = list(self._cores) + for left_index, left in enumerate(snapshot): + for right in snapshot[left_index + 1 :]: + difference = left ^ right + if len(left) != len(right) or len(difference) != 2 or sum(difference) != 0: + continue + resolvent = left & right + if not self._is_subsumed(resolvent): + return resolvent + return None + + def pack(self, device: torch.device) -> torch.Tensor: + width = max((len(core) for core in self._cores), default=0) + packed = torch.zeros((len(self._cores), width), dtype=torch.long, device=device) + for row, core in enumerate(self._cores): + values = sorted(core, key=lambda literal: (abs(literal), literal)) + packed[row, : len(values)] = torch.tensor(values, dtype=torch.long, device=device) + return packed + + +def apply_literal_assignments( + batch: SubproblemBatch, + assignments: torch.Tensor, + active_variables: torch.Tensor, + codec: LiteralCodec, +) -> SubproblemBatch: + """Decode compact inferred phases back to a lossless subproblem batch.""" + signs = {key: value.clone() for key, value in (batch.split_signs or {}).items()} + specs = 1 + if signs: + specs = next(iter(signs.values())).shape[1] + elif batch.incremental_alpha: + specs = next(iter(batch.incremental_alpha.values())).shape[1] + elif batch.incremental_eta: + specs = next(iter(batch.incremental_eta.values())).shape[1] + if active_variables.numel() == 0: + batch.split_signs = signs or None + return batch + positive_literals = active_variables + 1 + layer_ids, neurons, _ = codec.decode_tensor(positive_literals) + for layer_id, width in zip(codec.layer_ids, codec.widths): + columns = torch.where(layer_ids == layer_id)[0] + if columns.numel() == 0: + continue + if layer_id not in signs: + signs[layer_id] = torch.zeros( + batch.batch_size, + specs, + width, + dtype=batch.lb.dtype, + device=batch.lb.device, + ) + layer_values = assignments.index_select(1, columns) + lane_indices, local_columns = torch.nonzero(layer_values, as_tuple=True) + if lane_indices.numel() == 0: + continue + layer_neurons = neurons.index_select(0, columns).index_select(0, local_columns) + values = layer_values[lane_indices, local_columns].to(signs[layer_id]) + signs[layer_id][lane_indices, :, layer_neurons] = values.unsqueeze(1) + batch.split_signs = signs or None + return batch + + +@torch.no_grad() +def propagate( + batches: Sequence[SubproblemBatch], + library: CoreLibrary, + codec: LiteralCodec, + *, + core_chunk: int, +) -> Tuple[SubproblemBatch, ...]: + """Apply certified cores jointly to pending and active groups to fixpoint. + + ``core_chunk`` caps the cores per matmul chunk; it bounds memory only, the + fixpoint is independent of it. + """ + if core_chunk < 1: + raise ValueError("core_chunk must be positive") + if not batches: + return () + group_sizes = [batch.batch_size for batch in batches] + packed_batches = [ + pack_split_matrix( + batch.split_signs, + batch.batch_size, + codec, + device=batch.lb.device, + ) + for batch in batches + ] + if library.core_count == 0: + return tuple(batches) + device = batches[0].lb.device + packed_cores = library.pack(device) + batch_values = torch.cat( + [packed[packed != 0] for packed in packed_batches] + ) + core_values = packed_cores[packed_cores != 0] + all_values = torch.cat((batch_values, core_values)) + active_variables, inverse = torch.unique( + all_values.abs() - 1, sorted=True, return_inverse=True + ) + assignments = torch.zeros( + (sum(group_sizes), active_variables.numel()), dtype=torch.int8, device=device + ) + offset = 0 + inverse_offset = 0 + for packed, size in zip(packed_batches, group_sizes): + lanes, slots = torch.nonzero(packed, as_tuple=True) + count = lanes.numel() + assignments[offset + lanes, inverse[inverse_offset : inverse_offset + count]] = packed[ + lanes, slots + ].sign().to(torch.int8) + inverse_offset += count + offset += size + initial_assignments = assignments.clone() + core_dense = torch.zeros( + (library.core_count, active_variables.numel()), dtype=torch.int8, device=device + ) + core_rows, core_slots = torch.nonzero(packed_cores, as_tuple=True) + core_dense[core_rows, inverse[inverse_offset:]] = packed_cores[ + core_rows, core_slots + ].sign().to(torch.int8) + active = torch.ones(assignments.shape[0], dtype=torch.bool, device=assignments.device) + compute_dtype = get_default_dtype() + core_abs = core_dense.abs() + lengths = core_abs.sum(dim=1).to(compute_dtype) + while bool(active.any().item()): + active_rows = torch.where(active)[0] + work = assignments[active] + work_f = work.to(compute_dtype) + work_abs_f = work.abs().to(compute_dtype) + # One snapshot per round: every chunk reads the same `work`, so the + # fixpoint is independent of chunk size (memory is O(N * chunk * W)). + discharge_local = torch.zeros(work.shape[0], dtype=torch.bool, device=work.device) + positive = torch.zeros_like(work, dtype=torch.bool) + negative = torch.zeros_like(positive) + for start in range(0, core_dense.shape[0], core_chunk): + chunk = core_dense[start : start + core_chunk] + chunk_lengths = lengths[start : start + core_chunk] + product = work_f @ chunk.to(compute_dtype).T + overlap = work_abs_f @ core_abs[start : start + core_chunk].to(compute_dtype).T + discharge_local |= (product == chunk_lengths.unsqueeze(0)).any(dim=1) + unit = (product == overlap) & (overlap == chunk_lengths.unsqueeze(0) - 1) + lane_local, core_local = torch.nonzero(unit, as_tuple=True) + if lane_local.numel() != 0: + missing = (chunk.index_select(0, core_local) != 0) & ( + work.index_select(0, lane_local) == 0 + ) + valid = missing.sum(dim=1) == 1 + lane_local = lane_local[valid] + core_local = core_local[valid] + columns = missing[valid].to(torch.int64).argmax(dim=1) + proposed = -chunk[core_local, columns] + positive[lane_local[proposed > 0], columns[proposed > 0]] = True + negative[lane_local[proposed < 0], columns[proposed < 0]] = True + if bool(discharge_local.any().item()): + active[active_rows[discharge_local]] = False + + survivor_local = ~discharge_local + if not bool(survivor_local.any().item()): + continue + survivor_rows = active_rows[survivor_local] + positive = positive[survivor_local] + negative = negative[survivor_local] + conflict = (positive & negative).any(dim=1) + if bool(conflict.any().item()): + active[survivor_rows[conflict]] = False + writable = ~conflict + changed = False + if bool(writable.any().item()): + rows = survivor_rows[writable] + proposals = positive[writable].to(torch.int8) - negative[writable].to(torch.int8) + write = (assignments[rows] == 0) & (proposals != 0) + if bool(write.any().item()): + assignments[rows] = torch.where(write, proposals, assignments[rows]) + changed = True + if not changed and not bool(discharge_local.any().item()) and not bool(conflict.any().item()): + break + + assert _only_fills_unassigned(initial_assignments, assignments), ( + "CLIMB invariant violated: propagation changed an existing split literal " + f"(rows={assignments.shape[0]}, variables={assignments.shape[1]})" + ) + result_batches: List[SubproblemBatch] = [] + offset = 0 + for batch, size in zip(batches, group_sizes): + group_active = active[offset : offset + size] + keep = torch.where(group_active)[0] + restricted = batch.select(keep) + restricted_assignments = assignments[offset : offset + size].index_select(0, keep) + propagated = apply_literal_assignments( + restricted, restricted_assignments, active_variables, codec + ) + assert propagated.batch_size <= size, ( + "CLIMB invariant violated: propagation increased a group batch size " + f"(input_rows={size}, output_rows={propagated.batch_size})" + ) + result_batches.append(propagated) + offset += size + return tuple(result_batches) + + +def is_certified(slack: torch.Tensor, out_kind: str) -> torch.Tensor: + """Per-lane certification mask over ``[N, M]`` slack rows. + + Byte-consistent with the ``DualSolver.solve_spec_batch`` lane statuses: an + ``UNSAFE_LINEAR`` lane certifies when any finite row is strictly positive; + every other kind certifies only when all rows are finite and non-negative. + """ + if out_kind == OutKind.UNSAFE_LINEAR: + return (torch.isfinite(slack) & (slack > 0)).any(dim=1) + return torch.isfinite(slack).all(dim=1) & (slack >= 0).all(dim=1) + + +@torch.no_grad() +def certificate_replay( + *, + net: Net, + batch: SubproblemBatch, + reference_bounds: Dict[int, Bounds], + c_rows: torch.Tensor, + thresholds: torch.Tensor, + m_specs: int, + out_kind: str, + literals: Optional[torch.Tensor] = None, + codec: Optional[LiteralCodec] = None, + with_costs: bool = True, +) -> Optional[tuple[torch.Tensor, Optional[torch.Tensor]]]: + """Replay a frozen alpha/eta certificate directly through ``DualSolver``. + + Returns ``(slack, costs)`` — for ``UNSAFE_LINEAR`` both are gathered on the + pinned strict row (``slack [N, 1]``, ``costs [N, 1, W]``); otherwise + ``slack`` is ``[N, M]`` and ``costs`` is ``[N, M, W]``. ``costs`` is + ``None`` when ``with_costs=False``. Returns ``None`` when the batch lacks + replayable state or the solver emits no usable per-layer nu. + """ + if with_costs and (literals is None or codec is None): + raise ValueError("with_costs=True requires packed literals") + if ( + batch.incremental_alpha is None + or batch.incremental_eta is None + or batch.split_signs is None + or m_specs < 1 + ): + return None + from act.back_end.solver.solver_dual import DualSolver + + eta = { + layer_id: value.detach().clone().clamp(min=0) + for layer_id, value in batch.incremental_eta.items() + } + result = DualSolver().compute_certified_bound( + net, + reference_bounds, + c_rows, + M=m_specs, + alpha=batch.incremental_alpha, + eta=eta, + split_signs=batch.split_signs, + optimize=False, + return_nu_per_layer=True, + local_phase_clamp=True, + ) + n_lanes = batch.batch_size + slack = result.margins.reshape(n_lanes, m_specs) + slack = slack - thresholds.reshape(n_lanes, m_specs).to(slack) + pinned_rows: Optional[torch.Tensor] = None + if out_kind == OutKind.UNSAFE_LINEAR: + passing = torch.isfinite(slack) & (slack > 0) + pinned_rows = passing.to(torch.int64).argmax(dim=1) + if result.nu_per_layer is None: + return None + nu: Dict[int, torch.Tensor] = {} + for layer_id, value in result.nu_per_layer.items(): + if value.shape[0] == n_lanes * m_specs: + nu[layer_id] = value.reshape(n_lanes, m_specs, -1) + elif value.shape[0] == n_lanes: + nu[layer_id] = value.reshape(n_lanes, -1, value.shape[-1]) + else: + return None + costs: Optional[torch.Tensor] = None + if with_costs: + assert literals is not None + assert codec is not None + costs = _literal_costs( + literals, codec, slack, eta, nu, reference_bounds, pinned_rows + ) + if pinned_rows is not None: + slack = slack.gather(1, pinned_rows[:, None]) + return slack, costs + + +@torch.no_grad() +def _literal_costs( + literals: torch.Tensor, + codec: LiteralCodec, + slack: torch.Tensor, + eta: Dict[int, torch.Tensor], + nu: Dict[int, torch.Tensor], + bounds: Dict[int, Bounds], + pinned_rows: Optional[torch.Tensor], +) -> torch.Tensor: + """Full ReLU Eq. 12 costs ``(eta + relu(-nu)) * m`` in ``slack.dtype``, ``[N, M, W]``.""" + n_lanes, width = literals.shape + m_specs = slack.shape[1] + costs = torch.zeros( + (n_lanes, m_specs, width), + dtype=slack.dtype, + device=literals.device, + ) + mask = literals != 0 + if bool(mask.any()): + layer_ids, neurons, signs = codec.decode_tensor(literals.masked_fill(~mask, 1)) + for layer_id in codec.layer_ids: + lane, slot = torch.where(mask & (layer_ids == layer_id)) + if lane.numel() == 0: + continue + if layer_id not in eta or layer_id not in nu or layer_id not in bounds: + costs[lane, :, slot] = torch.inf + continue + neuron = neurons[lane, slot] + lane_eta = eta[layer_id][lane, :, neuron].clamp(min=0) + lane_nu = nu[layer_id][lane, :, neuron] + layer_bounds = bounds[layer_id] + lower = layer_bounds.lb.flatten(start_dim=1)[lane, neuron] + upper = layer_bounds.ub.flatten(start_dim=1)[lane, neuron] + magnitude = torch.where( + signs[lane, slot] > 0, + (-lower).clamp(min=0), + upper.clamp(min=0), + ) + costs[lane, :, slot] = ( + lane_eta + (-lane_nu).clamp(min=0) + ) * magnitude.unsqueeze(1) + if pinned_rows is not None: + rows = pinned_rows.to(device=costs.device, dtype=torch.long) + costs = costs.gather(1, rows[:, None, None].expand(-1, 1, width)) + return costs + + +@torch.no_grad() +def coarsen( + literals: torch.Tensor, + costs: torch.Tensor, + slack: torch.Tensor, + *, + theta: float, + delta_abs: float, + delta_rel: float, + recheck_k: int, + recheck: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], +) -> torch.Tensor: + """Alg. 2 vector-budget deletion plus at most ``recheck_k`` one-step rechecks. + + Per row ``m`` the budget is ``theta * slack_m - delta_m`` with + ``delta_m = delta_abs + delta_rel * |slack_m|``; rows with + ``slack_m <= delta_m`` never fund a deletion. The budget arithmetic runs + in the dtype of ``costs``/``slack``. Literals are ordered by + ``(max_m cost_m / budget_m, column, slot)`` and the longest feasible prefix + is deleted. A recheck retries the next literal in that same order for the + first ``recheck_k`` lanes; ``recheck`` returns the ``BoolTensor[K]`` + certification mask that confirms each candidate deletion. + """ + if not 0.0 <= theta <= 1.0: + raise ValueError("theta must be in [0, 1]") + if not (delta_abs >= 0.0 and delta_rel >= 0.0): + raise ValueError("delta_abs and delta_rel must be non-negative") + if recheck_k < 0: + raise ValueError("recheck_k must be non-negative") + n_lanes, width = literals.shape + if costs.shape[0] != n_lanes or costs.shape[2] != width: + raise ValueError("cost tensor shape must be [N, M, W]") + if slack.shape != costs.shape[:2]: + raise ValueError("slack tensor shape must be [N, M]") + retained = literals != 0 + next_in_cost_order: Dict[int, int] = {} + for lane in range(n_lanes): + valid_slots = torch.where(literals[lane] != 0)[0] + if valid_slots.numel() == 0: + continue + lane_costs = costs[lane].index_select(1, valid_slots) + lane_slack = slack[lane] + if not bool(torch.isfinite(lane_costs).all().item()) or not bool( + torch.isfinite(lane_slack).all().item() + ): + continue + delta = delta_abs + delta_rel * lane_slack.abs() + if bool((lane_slack <= delta).any().item()): + continue + budget = theta * lane_slack - delta + denominator = budget.clamp(min=torch.finfo(budget.dtype).tiny) + ratios = (lane_costs / denominator[:, None]).amax(dim=0) + ordered = sorted( + range(valid_slots.numel()), + key=lambda pos: ( + float(ratios[pos].item()), + int(literals[lane, valid_slots[pos]].abs().item()) - 1, + int(valid_slots[pos].item()), + ), + ) + running = torch.zeros_like(lane_slack) + for position in ordered: + candidate = running + lane_costs[:, position] + slot = int(valid_slots[position].item()) + if not bool((candidate <= budget).all().item()): + next_in_cost_order[lane] = slot + break + retained[lane, slot] = False + running = candidate + + candidates: List[int] = [] + next_slots: List[int] = [] + if recheck_k > 0: + for lane, slot in next_in_cost_order.items(): + candidates.append(lane) + next_slots.append(slot) + if len(candidates) == recheck_k: + break + if candidates: + rows = torch.tensor(candidates, dtype=torch.long, device=literals.device) + candidate_mask = retained.index_select(0, rows).clone() + for local, slot in enumerate(next_slots): + candidate_mask[local, slot] = False + candidate = literals.index_select(0, rows).masked_fill(~candidate_mask, 0) + passing = recheck(candidate, rows) + for local, passed in enumerate(passing.tolist()): + if passed: + retained[candidates[local], next_slots[local]] = False + + return literals.masked_fill(~retained, 0) + + +def validate_climb_config(config: BaBConfig) -> None: + """Fail before solving when CLIMB's soundness prerequisites are absent.""" + if not config.climb_enabled: + return + failures: List[str] = [] + if config.solver_tier != CLIMB_SOLVER_TIER: + failures.append(f"solver_tier must be {CLIMB_SOLVER_TIER!r}") + if config.bounding not in TOP_K_BOUNDINGS: + failures.append(f"bounding must be one of {TOP_K_BOUNDINGS}") + if config.branching_method not in NEURON_BRANCHING_METHODS: + failures.append(f"branching_method must be one of {NEURON_BRANCHING_METHODS}") + if config.root_bounds_reuse != "none": + failures.append("root_bounds_reuse must be 'none'") + if config.intermediate_refine != NO_REFINEMENT_MODE: + failures.append(f"intermediate_refine must be {NO_REFINEMENT_MODE!r}") + if config.per_subproblem_refine != NO_REFINEMENT_MODE: + failures.append(f"per_subproblem_refine must be {NO_REFINEMENT_MODE!r}") + if failures: + raise ConfigError("CLIMB configuration error: " + "; ".join(failures)) + + +@dataclass(frozen=True) +class _ReplayContext: + batch: SubproblemBatch + bounds: Dict[int, Bounds] + c_lanes: torch.Tensor + thresholds: torch.Tensor + m_specs: int + + +class ClimbSession: + """Own the query-local CLIMB library, metrics, and BaB-loop hooks.""" + + def __init__(self, config: BaBConfig, net: Net, out_kind: str) -> None: + self._net = net + self._out_kind = out_kind + self._theta = config.climb_theta + self._delta_abs = config.climb_delta_abs + self._delta_rel = config.climb_delta_rel + self._recheck_k = config.climb_recheck_k + self._propagate_core_chunk = config.climb_propagate_core_chunk + self.codec = LiteralCodec.from_net(net) + self.library = CoreLibrary(config.climb_max_cores) + self.metrics = ClimbMetrics(enabled=True, library=self.library) + + @classmethod + def from_config( + cls, config: BaBConfig, net: Net, out_kind: str + ) -> Optional[ClimbSession]: + validate_climb_config(config) + return cls(config, net, out_kind) if config.climb_enabled else None + + def before_pop(self, pool: TopKBounding) -> None: + if self.library.core_count > 0: + # H1: eager discharge of the whole pending pool before any lane is + # bounded. view_all/replace_all never score, cool, or probe. + pending_before = pool.view_all() + propagate_started = time.perf_counter() + (pending_after,) = propagate( + [pending_before], + self.library, + self.codec, + core_chunk=self._propagate_core_chunk, + ) + self.metrics.propagate_time_s += time.perf_counter() - propagate_started + self.metrics.prebound_discharged += ( + pending_before.batch_size - pending_after.batch_size + ) + pool.replace_all(pending_after) + if pool.empty: + return + self.metrics.observe_frontier(len(pool)) + + def after_pop(self, frontier_rows: int) -> None: + self.metrics.observe_frontier(frontier_rows) + + def after_main_bound(self, k_actual: int) -> None: + self.metrics.main_bound_calls += 1 + self.metrics.main_bound_row_passes += k_actual + + def learn( + self, + batch: SubproblemBatch, + statuses: Sequence[str], + dual_solve_result: DualBatchResult, + k_actual: int, + ) -> None: + # H2: replay -> coarsen -> insert. Certified lanes only; replay reads + # the immutable forward snapshot, the theta budget coarsens the + # validated vector, and the survivor is inserted as a core. + unsat_idx = torch.tensor( + [i for i, status in enumerate(statuses) if status == SolveStatus.UNSAT], + device=batch.lb.device, + dtype=torch.long, + ) + if ( + int(unsat_idx.numel()) == 0 + or dual_solve_result.c_rows is None + or dual_solve_result.thresholds is None + or dual_solve_result.reference_bounds is None + or dual_solve_result.m_specs <= 0 + ): + return + replay_batch = batch.select(unsat_idx) + replay_m_specs = dual_solve_result.m_specs + # The replay reference is the exact unhardened forward snapshot the + # main solve consumed, sliced by certified lane; it is never recomputed + # so Eq. 12 magnitudes and the local clamps read the same immutable bounds. + replay_bounds = slice_bounds_dict( + dual_solve_result.reference_bounds, unsat_idx + ) + replay_c_lanes = dual_solve_result.c_rows.reshape( + k_actual, replay_m_specs, -1 + ).index_select(0, unsat_idx.to(dual_solve_result.c_rows.device)) + replay_c = replay_c_lanes.reshape(-1, replay_c_lanes.shape[-1]) + replay_thresholds = dual_solve_result.thresholds.index_select( + 0, unsat_idx.to(dual_solve_result.thresholds.device) + ) + packed = pack_split_matrix( + replay_batch.split_signs, + replay_batch.batch_size, + self.codec, + device=replay_batch.lb.device, + ) + replay = certificate_replay( + net=self._net, + batch=replay_batch, + reference_bounds=replay_bounds, + c_rows=replay_c, + thresholds=replay_thresholds, + m_specs=replay_m_specs, + out_kind=self._out_kind, + literals=packed, + codec=self.codec, + ) + if replay is not None: + slack, costs = replay + valid = is_certified(slack, self._out_kind) + else: + slack = costs = valid = None + if costs is None or valid is None or not bool(valid.any().item()): + return + budget_slack = slack + assert budget_slack is not None + costs[~valid] = torch.inf + context = _ReplayContext( + replay_batch, replay_bounds, replay_c_lanes, replay_thresholds, replay_m_specs + ) + coarsen_started = time.perf_counter() + retained = coarsen( + packed, + costs, + budget_slack, + theta=self._theta, + delta_abs=self._delta_abs, + delta_rel=self._delta_rel, + recheck_k=self._recheck_k, + recheck=lambda candidate, rows: self._instrumented_recheck( + context, candidate, rows + ), + ) + self.metrics.coarsen_time_s += time.perf_counter() - coarsen_started + original_literals = int((packed != 0).sum()) + retained_literals = int((retained != 0).sum()) + assert _retained_literals_are_subset(packed, retained), ( + "CLIMB invariant violated: coarsening added or changed a literal " + f"(rows={packed.shape[0]}, width={packed.shape[1]}, " + f"original_literals={original_literals}, retained_literals={retained_literals})" + ) + self.metrics.core_literals_original += original_literals + self.metrics.core_literals_retained += retained_literals + assert ( + self.metrics.core_literals_retained + <= self.metrics.core_literals_original + ), ( + "CLIMB invariant violated: retained literal accounting exceeds original " + f"(original_total={self.metrics.core_literals_original}, " + f"retained_total={self.metrics.core_literals_retained})" + ) + valid_rows = torch.where(valid)[0] + valid_cores = retained.index_select(0, valid_rows) + self.metrics.cores_inserted += self.library.insert(valid_cores) + + def _recheck( + self, + context: _ReplayContext, + candidate: torch.Tensor, + rows: torch.Tensor, + ) -> torch.Tensor: + candidate_batch = context.batch.select(rows) + # The solver indexes split_signs[lid] for every incremental_eta layer, so a + # layer whose literals were all coarsened away stays as an all-zero block. + if candidate_batch.split_signs is not None: + candidate_batch.split_signs = { + lid: torch.zeros_like(signs) + for lid, signs in candidate_batch.split_signs.items() + } + candidate_values = candidate[candidate != 0] + active_variables = torch.unique(candidate_values.abs() - 1, sorted=True) + dense = torch.zeros( + (candidate.shape[0], active_variables.numel()), + dtype=torch.int8, + device=candidate.device, + ) + lanes, slots = torch.nonzero(candidate, as_tuple=True) + columns = torch.searchsorted(active_variables, candidate[lanes, slots].abs() - 1) + dense[lanes, columns] = candidate[lanes, slots].sign().to(torch.int8) + apply_literal_assignments(candidate_batch, dense, active_variables, self.codec) + eta_layers = set(candidate_batch.incremental_eta or {}) + sign_layers = set(candidate_batch.split_signs or {}) + missing_layers = sorted(eta_layers - sign_layers) + assert not missing_layers, ( + "CLIMB invariant violated: recheck split_signs missing incremental_eta layer " + f"(KeyError: {missing_layers[0]}, eta_layers={len(eta_layers)}, " + f"split_sign_layers={len(sign_layers)})" + ) + row_bounds = slice_bounds_dict(context.bounds, rows) + row_c_lanes = context.c_lanes.index_select( + 0, rows.to(context.c_lanes.device) + ) + row_c = row_c_lanes.reshape(-1, row_c_lanes.shape[-1]) + row_thresholds = context.thresholds.index_select( + 0, rows.to(context.thresholds.device) + ) + replayed = certificate_replay( + net=self._net, + batch=candidate_batch, + reference_bounds=row_bounds, + c_rows=row_c, + thresholds=row_thresholds, + m_specs=context.m_specs, + out_kind=self._out_kind, + with_costs=False, + ) + if replayed is None: + return torch.zeros( + int(rows.numel()), dtype=torch.bool, device=rows.device + ) + replayed_slack, _ = replayed + return is_certified(replayed_slack, self._out_kind) + + def _instrumented_recheck( + self, + context: _ReplayContext, + candidate: torch.Tensor, + rows: torch.Tensor, + ) -> torch.Tensor: + recheck_started = time.perf_counter() + passing = self._recheck(context, candidate, rows) + self.metrics.recheck_time_s += time.perf_counter() - recheck_started + self.metrics.recheck_calls += 1 + self.metrics.recheck_row_passes += int(rows.numel()) + return passing + + def before_split( + self, pool: TopKBounding, unresolved: SubproblemBatch + ) -> SubproblemBatch: + if self.library.core_count == 0: + return unresolved + # H3: one joint fixpoint over (pending, active); pending removals were + # never bounded (pre-bound), active removals are parents spared a split + # (pre-split). + self.metrics.observe_frontier(len(pool) + unresolved.batch_size) + propagation_groups: List[SubproblemBatch] = [] + has_pool = len(pool) > 0 + if has_pool: + propagation_groups.append(pool.view_all()) + propagation_groups.append(unresolved) + presplit_before = unresolved.batch_size + propagate_started = time.perf_counter() + propagation = propagate( + propagation_groups, + self.library, + self.codec, + core_chunk=self._propagate_core_chunk, + ) + self.metrics.propagate_time_s += time.perf_counter() - propagate_started + if has_pool: + pool.replace_all(propagation[0]) + self.metrics.prebound_discharged += ( + propagation_groups[0].batch_size - propagation[0].batch_size + ) + unresolved = propagation[-1] + self.metrics.presplit_discharged += presplit_before - unresolved.batch_size + self.metrics.observe_frontier(len(pool) + unresolved.batch_size) + return unresolved + + def reject_input_axis_split(self) -> None: + raise ValueError( + "CLIMB requires ReLU phase literals only: the " + "brancher fell back to an input-axis split, which " + "would break the shared-box premise of core reuse" + ) + + def after_children(self, generated: int, frontier_rows: int) -> None: + self.metrics.generated_children += generated + self.metrics.observe_frontier(frontier_rows) + + def metadata(self) -> Dict[str, int | float]: + return self.metrics.metadata() diff --git a/act/back_end/bab/node.py b/act/back_end/bab/node.py index eacb9fcd1..b1a287a5e 100644 --- a/act/back_end/bab/node.py +++ b/act/back_end/bab/node.py @@ -13,11 +13,9 @@ # with leading batch dimension (N, …) so that branching, bounding, and # (future) batched solving operate in pure tensor arithmetic. # -# BabNode is a thin single-subproblem wrapper that carries one priority -# score and one optional candidate counterexample. ``BabNode.to_batch`` -# converts a single node into a ``SubproblemBatch`` of size 1 so that -# call sites holding a single node can dispatch through the batched -# code path. +# ``SubproblemBatch.select`` and ``SubproblemBatch.concat`` are the lane +# gather / stack primitives shared by the BaB loop and the branchers; +# child construction for input and neuron splits lives here. # # ===---------------------------------------------------------------------====# @@ -27,10 +25,11 @@ import math from typing import Dict, Optional -import numpy as np import torch -from act.back_end.core import Bounds +from act.back_end.core import Bounds, Layer, Net, ParamValue +from act.back_end.layer_schema import LayerKind +from act.front_end.specs import InKind # --------------------------------------------------------------------------- @@ -101,11 +100,52 @@ def from_bounds(bounds: Bounds, depth: int = 0) -> SubproblemBatch: depths = torch.full((b,), depth, dtype=torch.long, device=lb.device) return SubproblemBatch(lb=lb, ub=ub, depths=depths) - # -- conversions -------------------------------------------------------- + # -- lane selection / concatenation ------------------------------------- + + def select(self, indices: torch.Tensor) -> SubproblemBatch: + """Return the lanes at ``indices`` with every per-lane field preserved.""" + return SubproblemBatch( + lb=self.lb.index_select(0, indices.to(self.lb.device)), + ub=self.ub.index_select(0, indices.to(self.ub.device)), + depths=self.depths.index_select(0, indices.to(self.depths.device)), + incremental_alpha=_gather_optional_dict(self.incremental_alpha, indices), + incremental_eta=_gather_optional_dict(self.incremental_eta, indices), + split_signs=_gather_optional_dict(self.split_signs, indices), + parent_margins=_gather_optional_tensor(self.parent_margins, indices), + lower_bound=_gather_optional_tensor(self.lower_bound, indices), + node_id=_gather_optional_tensor(self.node_id, indices), + parent_id=_gather_optional_tensor(self.parent_id, indices), + ) - def to_bounds_list(self) -> list[Bounds]: - """Convert back to a list of ``Bounds`` for solver compatibility.""" - return [Bounds(self.lb[i], self.ub[i]) for i in range(self.batch_size)] + def concat(self, other: SubproblemBatch) -> SubproblemBatch: + """Stack ``other``'s lanes after this batch's lanes. + + ``other`` is moved to this batch's device (optional per-lane tensors + also to its dtype). The per-layer state dicts may disagree on their key + sets: a layer touched only by one side's splits exists there and not in + the other. Zero padding is exact for ``split_signs`` and + ``incremental_eta`` (0 = unconstrained / no multiplier), so it + preserves each side's feasible region. ``incremental_alpha`` is only + ever concatenated between children of the same parent batch, whose key + sets agree. + """ + n_self, n_other = self.batch_size, other.batch_size + return SubproblemBatch( + lb=torch.cat([self.lb, other.lb.to(self.lb.device)], dim=0), + ub=torch.cat([self.ub, other.ub.to(self.ub.device)], dim=0), + depths=torch.cat([self.depths, other.depths.to(self.depths.device)], dim=0), + incremental_alpha=_concat_padded_dict( + self.incremental_alpha, other.incremental_alpha, n_self, n_other + ), + incremental_eta=_concat_padded_dict( + self.incremental_eta, other.incremental_eta, n_self, n_other + ), + split_signs=_concat_padded_dict(self.split_signs, other.split_signs, n_self, n_other), + parent_margins=_concat_optional_tensor(self.parent_margins, other.parent_margins), + lower_bound=_concat_optional_tensor(self.lower_bound, other.lower_bound), + node_id=_concat_optional_tensor(self.node_id, other.node_id), + parent_id=_concat_optional_tensor(self.parent_id, other.parent_id), + ) # -- geometry ----------------------------------------------------------- @@ -113,60 +153,271 @@ def widths(self) -> torch.Tensor: """Per-dimension widths: ``(N, D)``.""" return self.ub - self.lb - def total_width(self) -> torch.Tensor: - """Sum of widths per subproblem: ``(N,)``.""" - return self.widths().sum(dim=-1) - # --------------------------------------------------------------------------- -# Batch splitting (tensor-native) +# Lane gather / bounds slicing / concat helpers # --------------------------------------------------------------------------- +def slice_bounds_dict( + bounds_dict: Dict[int, Bounds], rows: torch.Tensor +) -> Dict[int, Bounds]: + """Select lanes from a batched bounds dictionary without mutating it.""" + return { + layer_id: Bounds( + bounds.lb.index_select(0, rows.to(bounds.lb.device)), + bounds.ub.index_select(0, rows.to(bounds.ub.device)), + ) + for layer_id, bounds in bounds_dict.items() + } + + +def slice_branching_state( + bounds_dict: Optional[Dict[int, Bounds]], + nu_per_layer: Optional[Dict[int, torch.Tensor]], + lane_idx: torch.Tensor, + k_actual: int, +) -> tuple[Optional[Dict[int, Bounds]], Optional[Dict[int, torch.Tensor]]]: + # ν/bounds are computed over the full k_actual wave; the brancher runs on the + # sub-batch actually being split. Bounds are [k_actual, *]; ν is [k_actual*M, n] + # packed sample-major (row b*M+j), so ν rows expand per selected lane. + bd_out = ( + slice_bounds_dict(bounds_dict, lane_idx) if bounds_dict is not None else None + ) + nu_out: Optional[Dict[int, torch.Tensor]] = None + if nu_per_layer is not None: + nu_out = {} + for lid, tensor in nu_per_layer.items(): + total = int(tensor.shape[0]) + if k_actual > 0 and total != k_actual and total % k_actual == 0: + m = total // k_actual + rows = ( + lane_idx.to(tensor.device).unsqueeze(1) * m + + torch.arange(m, device=tensor.device) + ).reshape(-1) + else: + rows = lane_idx.to(tensor.device) + nu_out[lid] = tensor.index_select(0, rows) + return bd_out, nu_out + + +def _layer_neuron_count(layer: Layer) -> int: + """Width of a layer's contiguous output-variable block.""" + span = int(layer.out_vars[-1] - layer.out_vars[0] + 1) + assert span == len(layer.out_vars), ( + f"layer {layer.id} has non-contiguous out_vars: span {span} != " + f"{len(layer.out_vars)} declared variables" + ) + return span + + def _gather_optional_dict( d: Optional[Dict[int, torch.Tensor]], idx: torch.Tensor, ) -> Optional[Dict[int, torch.Tensor]]: + """Gather rows ``idx`` from every tensor of an optional per-layer dict.""" if d is None: return None return {k: t.index_select(0, idx.to(t.device)) for k, t in d.items()} +def _gather_optional_tensor( + t: Optional[torch.Tensor], + idx: torch.Tensor, +) -> Optional[torch.Tensor]: + """Gather rows ``idx`` from an optional per-lane tensor.""" + return None if t is None else t.index_select(0, idx.to(t.device)) + + +def _concat_optional_tensor( + x: Optional[torch.Tensor], + y: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + """Stack two optional per-lane tensors; ``None`` unless both are present.""" + if x is None or y is None: + return None + return torch.cat([x, y.to(device=x.device, dtype=x.dtype)], dim=0) + + +def _concat_padded_dict( + x: Optional[Dict[int, torch.Tensor]], + y: Optional[Dict[int, torch.Tensor]], + n_x: int, + n_y: int, +) -> Optional[Dict[int, torch.Tensor]]: + """Stack two per-layer dicts, zero-filling a layer block missing on one side.""" + if not x and not y: + return None + x = x or {} + y = y or {} + out: Dict[int, torch.Tensor] = {} + for lid in sorted(set(x) | set(y)): + present = x.get(lid) + if present is None: + present = y[lid] + trailing = present.shape[1:] + left = x.get(lid) + if left is None: + left = torch.zeros((n_x, *trailing), dtype=present.dtype, device=present.device) + right = y.get(lid) + if right is None: + right = torch.zeros((n_y, *trailing), dtype=present.dtype, device=present.device) + out[lid] = torch.cat([left, right.to(device=left.device)], dim=0) + return out + + +# --------------------------------------------------------------------------- +# Batch splitting (tensor-native) +# --------------------------------------------------------------------------- + + +def _assert_splittable(batch: SubproblemBatch, dims2: torch.Tensor) -> None: + """Refuse a zero-width cut: its children are copies of the parent.""" + widths_at = batch.widths().gather(1, dims2) + if not bool((widths_at > 0).all().item()): + raise ValueError( + f"split dimension has zero width (min {float(widths_at.min().item()):g}); " + "bisecting it returns two children identical to the parent" + ) + + +def _child_lanes( + batch: SubproblemBatch, + parent_index: torch.Tensor, + depth_inc: int, +) -> SubproblemBatch: + """Gather child lanes while leaving provenance for the BaB engine to assign.""" + return SubproblemBatch( + lb=batch.lb.index_select(0, parent_index.to(batch.lb.device)), + ub=batch.ub.index_select(0, parent_index.to(batch.ub.device)), + depths=( + batch.depths.index_select(0, parent_index.to(batch.depths.device)) + + depth_inc + ), + incremental_alpha=_gather_optional_dict( + batch.incremental_alpha, parent_index + ), + incremental_eta=_gather_optional_dict(batch.incremental_eta, parent_index), + split_signs=_gather_optional_dict(batch.split_signs, parent_index), + parent_margins=_gather_optional_tensor(batch.parent_margins, parent_index), + lower_bound=_gather_optional_tensor(batch.lower_bound, parent_index), + ) + + def split_input( batch: SubproblemBatch, split_dims: torch.Tensor, ) -> tuple[SubproblemBatch, torch.Tensor]: n = batch.batch_size device = batch.lb.device + dims2 = split_dims.unsqueeze(1) + _assert_splittable(batch, dims2) mid = (batch.lb + batch.ub) / 2 - split_vals = mid.gather(1, split_dims.unsqueeze(1)) + split_vals = mid.gather(1, dims2) parent_index = torch.arange(n, device=device).repeat(2) - child_lb = batch.lb.index_select(0, parent_index) - child_ub = batch.ub.index_select(0, parent_index) - dims2 = split_dims.unsqueeze(1) - child_ub[:n].scatter_(1, dims2, split_vals) - child_lb[n:].scatter_(1, dims2, split_vals) - child_depths = batch.depths.index_select(0, parent_index) + 1 - - children = SubproblemBatch( - lb=child_lb, - ub=child_ub, - depths=child_depths, - incremental_alpha=_gather_optional_dict(batch.incremental_alpha, parent_index), - incremental_eta=_gather_optional_dict(batch.incremental_eta, parent_index), - split_signs=_gather_optional_dict(batch.split_signs, parent_index), - parent_margins=( - batch.parent_margins.index_select(0, parent_index.to(batch.parent_margins.device)) - if batch.parent_margins is not None - else None - ), - lower_bound=( - batch.lower_bound.index_select(0, parent_index.to(batch.lower_bound.device)) - if batch.lower_bound is not None - else None - ), + children = _child_lanes(batch, parent_index, 1) + children.ub[:n].scatter_(1, dims2, split_vals) + children.lb[n:].scatter_(1, dims2, split_vals) + return children, parent_index + + +def _selected_parent_sign_conflicts( + batch: SubproblemBatch, + top_layers: torch.Tensor, + top_neurons: torch.Tensor, + k: int, +) -> int: + """Count selected neurons that already carry a parent split literal.""" + if not batch.split_signs: + return 0 + conflicts = 0 + for bit in range(k): + for layer_id_value in torch.unique(top_layers[:, bit]).tolist(): + layer_id = int(layer_id_value) + parent_signs = batch.split_signs.get(layer_id) + if parent_signs is None: + continue + lanes = torch.where(top_layers[:, bit] == layer_id_value)[0] + neurons = top_neurons[lanes, bit].to( + device=parent_signs.device, dtype=torch.long + ) + selected = parent_signs[ + lanes.to(parent_signs.device), :, neurons + ] + conflicts += int((selected != 0).any(dim=1).sum().item()) + return conflicts + + +def split_neurons( + batch: SubproblemBatch, + net: Net, + top_layers: torch.Tensor, + top_neurons: torch.Tensor, + k: int, + *, + first_sign: float = -1.0, +) -> tuple[SubproblemBatch, torch.Tensor]: + """Build every sign-combination child for each lane's selected neurons.""" + n_lanes = batch.batch_size + n_children = 2**k + device = batch.lb.device + assert ( + conflicts := _selected_parent_sign_conflicts( + batch, top_layers, top_neurons, k + ) + ) == 0, ( + "NEURON SPLIT invariant violated: selected neuron already split in parent " + f"(lanes={n_lanes}, k={k}, selections={n_lanes * k}, " + f"conflicts={conflicts})" ) + parent_index = torch.arange(n_lanes, device=device).repeat(n_children) + children = _child_lanes(batch, parent_index, k) + expected_children = n_lanes * n_children + assert ( + children.batch_size == parent_index.numel() == expected_children + ), ( + "NEURON SPLIT invariant violated: child/provenance count mismatch " + f"(lanes={n_lanes}, k={k}, expected={expected_children}, " + f"children={children.batch_size}, parent_indices={parent_index.numel()})" + ) + + m_specs = 1 + if batch.incremental_alpha: + m_specs = int(next(iter(batch.incremental_alpha.values())).shape[1]) + elif batch.split_signs: + m_specs = int(next(iter(batch.split_signs.values())).shape[1]) + + signs = children.split_signs or {} + for layer_id_value in torch.unique(top_layers).tolist(): + layer_id = int(layer_id_value) + n_neurons = _layer_neuron_count(net.by_id[layer_id]) + if layer_id not in signs: + signs[layer_id] = torch.zeros( + n_children * n_lanes, + m_specs, + n_neurons, + device=device, + dtype=batch.lb.dtype, + ) + else: + signs[layer_id] = signs[layer_id].clone() + for bit in range(k): + lane_selection = torch.where(top_layers[:, bit] == layer_id_value)[0] + if lane_selection.numel() == 0: + continue + neuron_selection = top_neurons[lane_selection, bit].to( + device=device, dtype=torch.long + ) + for child in range(n_children): + sign_value = ( + -first_sign if (child >> bit) & 1 else first_sign + ) + rows = child * n_lanes + lane_selection + # Paired advanced indices write rows[i] to neuron_selection[i]. + signs[layer_id][rows, :, neuron_selection] = sign_value + + children.split_signs = signs return children, parent_index @@ -216,6 +467,74 @@ def rederive_embedding_block_eps( return torch.where(mask, radii, torch.zeros_like(radii)) +def _finite_embedding_spec(net: Net) -> Optional[Layer]: + for layer in net.layers: + if ( + layer.kind == LayerKind.INPUT_SPEC.value + and layer.params.get("kind") == InKind.LP_EMBEDDING + ): + p_norm = layer.params.get("p_norm", float("inf")) + if isinstance(p_norm, torch.Tensor): + p_value = float(p_norm.reshape(-1)[0].item()) + elif isinstance(p_norm, (int, float, bool)): + p_value = float(p_norm) + else: + continue + if p_value != float("inf"): + return layer + return None + + +def _install_embedding_child_block_eps( + net: Net, + batched_bounds: Bounds, + batch: SubproblemBatch, +) -> list[tuple[Layer, ParamValue]]: + """Install finite-p embedding radii derived from a split child box.""" + spec = _finite_embedding_spec(net) + if ( + spec is None + or batch.depths.numel() == 0 + or int(batch.depths.max().item()) == 0 + ): + return [] + p_raw = spec.params.get("p_norm", float("inf")) + if isinstance(p_raw, torch.Tensor): + p_norm = float(p_raw.reshape(-1)[0].item()) + elif isinstance(p_raw, (int, float, bool)): + p_norm = float(p_raw) + else: + return [] + input_shape = tuple(batched_bounds.lb.shape[1:]) + positions_raw = spec.params.get("perturbed_positions") + positions = positions_raw if isinstance(positions_raw, torch.Tensor) else None + block_eps = rederive_embedding_block_eps( + batched_bounds.lb.flatten(start_dim=1), + batched_bounds.ub.flatten(start_dim=1), + input_shape, + positions, + p_norm, + ) + old_values: list[tuple[Layer, ParamValue]] = [] + for layer in net.layers: + kind = layer.kind.upper() if isinstance(layer.kind, str) else layer.kind + if kind in (LayerKind.INPUT.value, LayerKind.INPUT_SPEC.value): + old_values.append((layer, layer.params.get("bab_block_eps", None))) + layer.params["bab_block_eps"] = block_eps + return old_values + + +def _restore_embedding_child_block_eps( + updates: list[tuple[Layer, ParamValue]], +) -> None: + """Restore input-layer embedding radii after a child solve.""" + for layer, old in updates: + if old is None: + layer.params.pop("bab_block_eps", None) + else: + layer.params["bab_block_eps"] = old + + def split_input_nary( batch: SubproblemBatch, cut_dim: torch.Tensor, @@ -225,235 +544,25 @@ def split_input_nary( raise ValueError(f"fanout must be >= 2, got {k}") n = batch.batch_size device = batch.lb.device + _assert_splittable( + batch, cut_dim.to(device=device, dtype=torch.long).reshape(-1).unsqueeze(1) + ) parent_index = torch.arange(n, device=device).repeat(k) section = torch.arange(k, device=device).repeat_interleave(n) - child_lb = batch.lb.index_select(0, parent_index) - child_ub = batch.ub.index_select(0, parent_index) + children = _child_lanes(batch, parent_index, math.ceil(math.log2(k))) cut_c = cut_dim.to(device=device, dtype=torch.long).index_select(0, parent_index) cut_c2 = cut_c.unsqueeze(1) - lb_at = child_lb.gather(1, cut_c2) - ub_at = child_ub.gather(1, cut_c2) + lb_at = children.lb.gather(1, cut_c2) + ub_at = children.ub.gather(1, cut_c2) seg = (ub_at - lb_at) / k - section_f = section.to(dtype=child_lb.dtype).unsqueeze(1) + section_f = section.to(dtype=children.lb.dtype).unsqueeze(1) new_lb = torch.where(section.unsqueeze(1) == 0, lb_at, lb_at + section_f * seg) new_ub = torch.where( section.unsqueeze(1) == k - 1, ub_at, lb_at + (section_f + 1) * seg, ) - child_lb.scatter_(1, cut_c2, new_lb) - child_ub.scatter_(1, cut_c2, new_ub) - - depth_inc = math.ceil(math.log2(k)) - child_depths = batch.depths.index_select(0, parent_index) + depth_inc - - children = SubproblemBatch( - lb=child_lb, - ub=child_ub, - depths=child_depths, - incremental_alpha=_gather_optional_dict(batch.incremental_alpha, parent_index), - incremental_eta=_gather_optional_dict(batch.incremental_eta, parent_index), - split_signs=_gather_optional_dict(batch.split_signs, parent_index), - parent_margins=( - batch.parent_margins.index_select(0, parent_index.to(batch.parent_margins.device)) - if batch.parent_margins is not None - else None - ), - lower_bound=( - batch.lower_bound.index_select(0, parent_index.to(batch.lower_bound.device)) - if batch.lower_bound is not None - else None - ), - ) + children.lb.scatter_(1, cut_c2, new_lb) + children.ub.scatter_(1, cut_c2, new_ub) return children, parent_index - - -def concat_children(a: SubproblemBatch, b: SubproblemBatch) -> SubproblemBatch: - def _concat_optional_dicts( - da: Optional[Dict[int, torch.Tensor]], - db: Optional[Dict[int, torch.Tensor]], - ) -> Optional[Dict[int, torch.Tensor]]: - if da is None or db is None: - assert da is None and db is None, "child dict fields must both be present or absent" - return None - assert set(da) == set(db), "child dict fields must have equal keys" - return {key: torch.cat([da[key], db[key]], dim=0) for key in sorted(da)} - - def _concat_optional_tensor( - ta: Optional[torch.Tensor], - tb: Optional[torch.Tensor], - name: str, - ) -> Optional[torch.Tensor]: - if ta is None or tb is None: - assert ta is None and tb is None, f"{name} must be present in both children or neither" - return None - return torch.cat([ta, tb], dim=0) - - return SubproblemBatch( - lb=torch.cat([a.lb, b.lb], dim=0), - ub=torch.cat([a.ub, b.ub], dim=0), - depths=torch.cat([a.depths, b.depths], dim=0), - incremental_alpha=_concat_optional_dicts(a.incremental_alpha, b.incremental_alpha), - incremental_eta=_concat_optional_dicts(a.incremental_eta, b.incremental_eta), - split_signs=_concat_optional_dicts(a.split_signs, b.split_signs), - parent_margins=_concat_optional_tensor(a.parent_margins, b.parent_margins, "parent_margins"), - lower_bound=_concat_optional_tensor(a.lower_bound, b.lower_bound, "lower_bound"), - ) - - -def split_subproblems( - batch: SubproblemBatch, - split_dims: torch.Tensor, -) -> tuple[SubproblemBatch, SubproblemBatch]: - """Bisect each subproblem along the chosen input dimension. - - This is a pure tensor operation — no Python loops over the batch. Incremental-state - fields are deep-copied into both children so they do not alias the parent or - each other. - - Args: - batch: ``(N, D)`` subproblems. - split_dims: ``(N,)`` long tensor — dimension to bisect per subproblem. - - Returns: - ``(left, right)`` — two ``SubproblemBatch`` of the same shape, - where ``left.ub[i, d] == right.lb[i, d] == midpoint``. - """ - mid = (batch.lb + batch.ub) / 2 # (N, D) - split_vals = mid.gather(1, split_dims.unsqueeze(1)) # (N, 1) - - # Left child: upper bound clamped at midpoint - left_ub = batch.ub.clone() - left_ub.scatter_(1, split_dims.unsqueeze(1), split_vals) - - # Right child: lower bound raised to midpoint - right_lb = batch.lb.clone() - right_lb.scatter_(1, split_dims.unsqueeze(1), split_vals) - - new_depths = batch.depths + 1 - - def _clone_dict_tensors( - tensors: Optional[Dict[int, torch.Tensor]], - ) -> Optional[Dict[int, torch.Tensor]]: - return ( - {key: tensor.clone() for key, tensor in tensors.items()} - if tensors is not None - else None - ) - - left_incremental_alpha = _clone_dict_tensors(batch.incremental_alpha) - left_incremental_eta = _clone_dict_tensors(batch.incremental_eta) - left_split_signs = _clone_dict_tensors(batch.split_signs) - left_parent_margins = ( - batch.parent_margins.clone() if batch.parent_margins is not None else None - ) - - right_incremental_alpha = _clone_dict_tensors(batch.incremental_alpha) - right_incremental_eta = _clone_dict_tensors(batch.incremental_eta) - right_split_signs = _clone_dict_tensors(batch.split_signs) - right_parent_margins = ( - batch.parent_margins.clone() if batch.parent_margins is not None else None - ) - - left = SubproblemBatch( - lb=batch.lb.clone(), - ub=left_ub, - depths=new_depths.clone(), - incremental_alpha=left_incremental_alpha, - incremental_eta=left_incremental_eta, - split_signs=left_split_signs, - parent_margins=left_parent_margins, - lower_bound=(batch.lower_bound.clone() if batch.lower_bound is not None else None), - ) - right = SubproblemBatch( - lb=right_lb, - ub=batch.ub.clone(), - depths=new_depths.clone(), - incremental_alpha=right_incremental_alpha, - incremental_eta=right_incremental_eta, - split_signs=right_split_signs, - parent_margins=right_parent_margins, - lower_bound=(batch.lower_bound.clone() if batch.lower_bound is not None else None), - ) - return left, right - - -def split_neuron_subproblems( - batch: SubproblemBatch, - *, - layer_id: int, - neuron_idx: int, - n_neurons: int, - n_specs: int, -) -> tuple[SubproblemBatch, SubproblemBatch]: - if neuron_idx < 0 or neuron_idx >= n_neurons: - raise IndexError(f"neuron_idx {neuron_idx} out of range for n_neurons={n_neurons}") - - def _clone_dict_tensors( - tensors: Optional[Dict[int, torch.Tensor]], - ) -> Optional[Dict[int, torch.Tensor]]: - return {key: tensor.clone() for key, tensor in tensors.items()} if tensors is not None else None - - def _clone_signs(sign: float) -> Dict[int, torch.Tensor]: - signs = _clone_dict_tensors(batch.split_signs) or {} - if layer_id not in signs: - signs[layer_id] = torch.zeros( - (batch.batch_size, n_specs, n_neurons), - dtype=batch.lb.dtype, - device=batch.lb.device, - ) - signs[layer_id][:, :, neuron_idx] = sign - return signs - - new_depths = batch.depths + 1 - on = SubproblemBatch( - lb=batch.lb.clone(), - ub=batch.ub.clone(), - depths=new_depths.clone(), - incremental_alpha=_clone_dict_tensors(batch.incremental_alpha), - incremental_eta=_clone_dict_tensors(batch.incremental_eta), - split_signs=_clone_signs(+1.0), - parent_margins=batch.parent_margins.clone() if batch.parent_margins is not None else None, - lower_bound=batch.lower_bound.clone() if batch.lower_bound is not None else None, - ) - off = SubproblemBatch( - lb=batch.lb.clone(), - ub=batch.ub.clone(), - depths=new_depths.clone(), - incremental_alpha=_clone_dict_tensors(batch.incremental_alpha), - incremental_eta=_clone_dict_tensors(batch.incremental_eta), - split_signs=_clone_signs(-1.0), - parent_margins=batch.parent_margins.clone() if batch.parent_margins is not None else None, - lower_bound=batch.lower_bound.clone() if batch.lower_bound is not None else None, - ) - return on, off - - -# --------------------------------------------------------------------------- -# Single-subproblem wrapper -# --------------------------------------------------------------------------- - - -@dataclass -class BabNode: - """Single-subproblem record: one bounds box, one priority score, one - optional candidate counterexample. - - Prefer :class:`SubproblemBatch` for batch-parallel processing; use this - record when a call site only ever holds one subproblem at a time and - wants priority-queue ordering. ``to_batch`` lifts a single node into a - ``SubproblemBatch`` of size 1 for dispatch through the batched path. - """ - - box: Bounds - depth: int - score: float - candidate_ce: Optional[np.ndarray] = None - - def __lt__(self, other: BabNode) -> bool: # max-heap by score - return self.score > other.score - - def to_batch(self) -> SubproblemBatch: - """Upgrade to tensor batch of size 1.""" - return SubproblemBatch.from_bounds(self.box, depth=self.depth) diff --git a/act/back_end/bab/violation.py b/act/back_end/bab/violation.py new file mode 100644 index 000000000..948d40f39 --- /dev/null +++ b/act/back_end/bab/violation.py @@ -0,0 +1,349 @@ +# ===- act/back_end/bab/violation.py - Counterexample Checking -----------====# +# ACT: Abstract Constraint Transformer +# Copyright (C) 2025– ACT Team +# +# Licensed under the GNU Affero General Public License v3.0 or later (AGPLv3+). +# Distributed without any warranty; see . +# ===---------------------------------------------------------------------====# +# +# Purpose: +# Concrete counterexample checks for BaB candidates. A candidate batch is +# run through the reconstructed PyTorch model (cached per net until +# ``clear_violation_check_module_cache``) and tested against the ASSERT +# parameters; ``_check_input_specs_batched`` confirms each candidate lies +# inside the INPUT_SPEC region. +# +# ===---------------------------------------------------------------------====# + +from __future__ import annotations + +from typing import List, cast + +import torch + +from act.back_end.core import Layer, Net +from act.front_end.specs import InKind, OutKind, normalize_position_mask +from act.util.model_inference import infer_single_model + + +def _as_batched_vector( + value: object, + n_batch: int, + width: int, + device: torch.device, + dtype: torch.dtype, + name: str, +) -> torch.Tensor: + t = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + t = t.to(device=device, dtype=dtype) + if t.dim() == 0: + return t.expand(n_batch, width).contiguous() + if t.dim() == 1: + if t.numel() == width: + return t.unsqueeze(0).expand(n_batch, -1).contiguous() + if width == 1 and t.numel() == n_batch: + return t.reshape(n_batch, 1).contiguous() + if t.dim() == 2: + if t.shape == (1, width): + return t.expand(n_batch, -1).contiguous() + if t.shape == (n_batch, width): + return t.contiguous() + raise ValueError( + f"{name}: expected scalar, ({width},), (1,{width}), or " + f"({n_batch},{width}); got {tuple(t.shape)}" + ) + + +def _as_batched_index( + value: object, + n_batch: int, + n_out: int, + device: torch.device, + name: str, +) -> torch.Tensor: + t = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + t = t.to(device=device, dtype=torch.long).reshape(-1) + if t.numel() == 1: + t = t.expand(n_batch) + if t.numel() != n_batch: + raise ValueError( + f"{name}: expected 1 or {n_batch} indices, got {t.numel()}" + ) + if bool(((t < 0) | (t >= n_out)).any().item()): + raise ValueError(f"{name}: index out of range for n_out={n_out}: {t.tolist()}") + return t.contiguous() + + +# Per-net cache of reconstructed PyTorch nn.Module used for CE validation. +# Without this, every check_violations_batched call rebuilds the module from +# scratch via ACTToTorch.run() — costly under K-batched BaB which can invoke +# CE validation dozens of times for a single net. Cleared per top-level +# verify-all dispatch via clear_violation_check_module_cache(). +_VIOLATION_CHECK_MODULE_CACHE: dict[int, torch.nn.Module] = {} + + +def clear_violation_check_module_cache() -> None: + _VIOLATION_CHECK_MODULE_CACHE.clear() + + +def _module_float_dtype(module: torch.nn.Module, default: torch.dtype) -> torch.dtype: + for tensor in module.parameters(): + if tensor.is_floating_point(): + return tensor.dtype + for tensor in module.buffers(): + if tensor.is_floating_point(): + return tensor.dtype + return default + + +def _forward_for_violation_check(net: object, x_batch: torch.Tensor) -> torch.Tensor: + if isinstance(net, torch.nn.Module): + module = net + else: + key = id(net) + cached = _VIOLATION_CHECK_MODULE_CACHE.get(key) + if cached is None: + from act.pipeline.verification.act2torch import ACTToTorch + + cached = ACTToTorch(cast(Net, net)).run() + # ACTToTorch emits mixed float32 weights + float64 buffers; unify to + # the analysis dtype or the internal forward clashes float32/float64. + if x_batch.is_floating_point(): + cached = cached.to(dtype=x_batch.dtype) + _VIOLATION_CHECK_MODULE_CACHE[key] = cached + module = cached + _ = module.eval() + target_dtype = _module_float_dtype(module, x_batch.dtype) + if x_batch.dtype != target_dtype: + x_batch = x_batch.to(dtype=target_dtype) + success, output, error = infer_single_model("ce_validate_batched", module, x_batch) + if not success or output is None: + raise RuntimeError(f"check_violations_batched: model forward failed: {error}") + if output.dim() < 2: + raise ValueError( + f"check_violations_batched: model output must be batched, got " + f"shape={tuple(output.shape)}" + ) + return output.reshape(output.shape[0], -1) + + +@torch.no_grad() +def check_violations_batched(net: object, x_batch: torch.Tensor, assert_layer: Layer) -> torch.Tensor: + """[BATCHED-API] Return a ``[N]`` bool tensor for concrete ASSERT violations. + + ``x_batch`` is always treated as a tensor-view batch ``[N, *input_shape]``; + N=1 is represented by a length-one leading dimension. ASSERT parameters are + read directly from ``assert_layer.params`` in their batch-native form. + """ + if x_batch.dim() < 2: + raise ValueError( + f"check_violations_batched: x_batch must be [N, *input_shape], " + f"got shape={tuple(x_batch.shape)}" + ) + y_batch = _forward_for_violation_check(net, x_batch) + n_batch = int(x_batch.shape[0]) + if int(y_batch.shape[0]) != n_batch: + raise ValueError( + f"check_violations_batched: output batch {int(y_batch.shape[0])} " + f"!= input batch {n_batch}" + ) + n_out = int(y_batch.shape[1]) + device = y_batch.device + dtype = y_batch.dtype + params = assert_layer.params + kind = params.get("kind") + eps = 1e-8 + + if kind == OutKind.TOP1_ROBUST: + y_true = _as_batched_index(params["y_true"], n_batch, n_out, device, "y_true") + y_true_scores = y_batch.gather(1, y_true.unsqueeze(1)).squeeze(1) + mask = torch.ones_like(y_batch, dtype=torch.bool) + _ = mask.scatter_(1, y_true.unsqueeze(1), False) + other_scores = y_batch.masked_fill(~mask, -float("inf")) + return (other_scores.max(dim=1).values - y_true_scores) >= 0 + + if kind == OutKind.MARGIN_ROBUST: + y_true = _as_batched_index(params["y_true"], n_batch, n_out, device, "y_true") + margin = _as_batched_vector( + params["margin"], n_batch, 1, device, dtype, "margin" + ).reshape(n_batch) + y_true_scores = y_batch.gather(1, y_true.unsqueeze(1)).squeeze(1) + mask = torch.ones_like(y_batch, dtype=torch.bool) + _ = mask.scatter_(1, y_true.unsqueeze(1), False) + other_scores = y_batch.masked_fill(~mask, -float("inf")) + # ``-margin``, not ``+margin``: ``encode_linear`` emits rows ``e_j - e_t`` + # with ``thresholds = -margin``, so a lane is certified iff + # ``max_j(z_j - z_t) < -margin``. The negative sign is deliberate. + return (other_scores.max(dim=1).values - y_true_scores) >= -margin + + if kind == OutKind.LINEAR_LE: + c_raw = params["c"] + c_t = (c_raw if isinstance(c_raw, torch.Tensor) else torch.as_tensor(c_raw)).to(device=device, dtype=dtype) + if c_t.dim() <= 1: + rows = 1 + elif c_t.dim() == 2: + if c_t.shape[1] != n_out: + raise ValueError(f"LINEAR_LE: c cols {c_t.shape[1]} != n_out {n_out}") + rows = int(c_t.shape[0]) + else: + rows = int(c_t.shape[1]) + if rows <= 1: + coeff = _as_batched_vector(params["c"], n_batch, n_out, device, dtype, "c") + bound = _as_batched_vector(params["d"], n_batch, 1, device, dtype, "d").reshape(n_batch) + return (coeff * y_batch).sum(dim=1) >= bound + eps + if c_t.dim() == 2: + c_view = c_t.unsqueeze(0).expand(n_batch, -1, -1) + else: + c_view = c_t if c_t.shape[0] == n_batch else c_t.expand(n_batch, -1, -1) + d_raw = params["d"] + d_t = (d_raw if isinstance(d_raw, torch.Tensor) else torch.as_tensor(d_raw)).to(device=device, dtype=dtype).flatten() + if d_t.numel() == rows: + d_view = d_t.unsqueeze(0).expand(n_batch, -1) + elif d_t.numel() == n_batch * rows: + d_view = d_t.reshape(n_batch, rows) + else: + raise ValueError(f"LINEAR_LE: d numel {d_t.numel()} incompatible with rows {rows}") + lhs = torch.einsum("bmo,bo->bm", c_view.contiguous(), y_batch) + return (lhs > d_view + eps).any(dim=1) + + if kind == OutKind.RANGE: + result = torch.zeros(n_batch, dtype=torch.bool, device=device) + lb_raw = params.get("lb") + ub_raw = params.get("ub") + if lb_raw is not None: + lb = _as_batched_vector(lb_raw, n_batch, n_out, device, dtype, "lb") + result = result | (y_batch < lb - eps).any(dim=1) + if ub_raw is not None: + ub = _as_batched_vector(ub_raw, n_batch, n_out, device, dtype, "ub") + result = result | (y_batch > ub + eps).any(dim=1) + return result + + if kind == OutKind.UNSAFE_LINEAR: + m_raw = params.get("M", 1) + if isinstance(m_raw, torch.Tensor): + m_rows = int(m_raw.item()) + elif isinstance(m_raw, int): + m_rows = m_raw + else: + raise ValueError(f"UNSAFE_LINEAR: M must be int, got {m_raw!r}") + c_raw = params.get("C", params.get("c")) + if c_raw is None: + raise ValueError("UNSAFE_LINEAR requires C or c params") + c_tensor = c_raw if isinstance(c_raw, torch.Tensor) else torch.as_tensor(c_raw) + c_tensor = c_tensor.to(device=device, dtype=dtype) + if c_tensor.dim() == 2: + if c_tensor.shape == (m_rows, n_out): + c_view = c_tensor.unsqueeze(0).expand(n_batch, -1, -1).contiguous() + elif c_tensor.shape == (n_batch * m_rows, n_out): + c_view = c_tensor.reshape(n_batch, m_rows, n_out).contiguous() + else: + raise ValueError( + f"UNSAFE_LINEAR: C shape {tuple(c_tensor.shape)} incompatible " + f"with N={n_batch}, M={m_rows}, n_out={n_out}" + ) + elif c_tensor.dim() == 3: + if c_tensor.shape == (1, m_rows, n_out): + c_view = c_tensor.expand(n_batch, -1, -1).contiguous() + elif c_tensor.shape == (n_batch, m_rows, n_out): + c_view = c_tensor.contiguous() + else: + raise ValueError( + f"UNSAFE_LINEAR: c shape {tuple(c_tensor.shape)} incompatible " + f"with N={n_batch}, M={m_rows}, n_out={n_out}" + ) + else: + raise ValueError(f"UNSAFE_LINEAR: unsupported C dim {c_tensor.dim()}") + d_raw = params.get("thresholds", params.get("d")) + if d_raw is None: + raise ValueError("UNSAFE_LINEAR requires thresholds or d params") + d_tensor = d_raw if isinstance(d_raw, torch.Tensor) else torch.as_tensor(d_raw) + d_tensor = d_tensor.to(device=device, dtype=dtype) + if d_tensor.dim() == 1 and d_tensor.numel() == m_rows: + d_view = d_tensor.unsqueeze(0).expand(n_batch, -1).contiguous() + elif d_tensor.shape == (1, m_rows): + d_view = d_tensor.expand(n_batch, -1).contiguous() + elif d_tensor.shape == (n_batch, m_rows): + d_view = d_tensor.contiguous() + else: + raise ValueError( + f"UNSAFE_LINEAR: d shape {tuple(d_tensor.shape)} incompatible " + f"with N={n_batch}, M={m_rows}" + ) + lhs = torch.einsum("bmo,bo->bm", c_view, y_batch) + return (lhs <= d_view + eps).all(dim=1) + + raise NotImplementedError(f"ASSERT kind not supported: {kind}") + + +def _check_input_specs_batched(x_batch: torch.Tensor, spec_layers: List[Layer]) -> torch.Tensor: + result = torch.ones(x_batch.shape[0], device=x_batch.device, dtype=torch.bool) + tol = 1e-7 + for layer in spec_layers: + kind = layer.params.get("kind") + if kind not in (InKind.BOX, InKind.LINF_BALL, InKind.LP_EMBEDDING): + continue + lb = layer.params.get("lb") + ub = layer.params.get("ub") + if isinstance(lb, torch.Tensor) and isinstance(ub, torch.Tensor): + lb_t = lb.to(device=x_batch.device, dtype=x_batch.dtype) + ub_t = ub.to(device=x_batch.device, dtype=x_batch.dtype) + result &= ((x_batch >= lb_t - tol) & (x_batch <= ub_t + tol)).flatten(start_dim=1).all(dim=1) + if kind != InKind.LP_EMBEDDING: + continue + center = layer.params.get("center") + eps = layer.params.get("eps") + p_norm = layer.params.get("p_norm") + if not isinstance(center, torch.Tensor) or eps is None or p_norm is None: + result &= torch.zeros_like(result) + continue + center_t = center.to(device=x_batch.device, dtype=x_batch.dtype) + if isinstance(eps, torch.Tensor): + eps_t = eps.to(device=x_batch.device, dtype=x_batch.dtype) + elif isinstance(eps, (int, float, bool)): + eps_t = center_t.new_tensor(float(eps)) + else: + result &= torch.zeros_like(result) + continue + if isinstance(p_norm, torch.Tensor): + p_value = float(p_norm.reshape(-1)[0].item()) + elif isinstance(p_norm, (int, float, bool)): + p_value = float(p_norm) + else: + result &= torch.zeros_like(result) + continue + positions_raw = layer.params.get("perturbed_positions") + positions = ( + positions_raw + if positions_raw is None or isinstance(positions_raw, (torch.Tensor, list, tuple)) + else None + ) + mask = normalize_position_mask( + positions, + int(center_t.shape[-2]), + batch_shape=tuple(center_t.shape[:-2]), + device=x_batch.device, + ) + if mask.shape[0] == 1 and x_batch.shape[0] != 1: + mask_b = mask.expand(x_batch.shape[0], *mask.shape[1:]) + center_b = center_t.expand_as(x_batch) + else: + mask_b = mask + center_b = center_t.expand_as(x_batch) + delta = x_batch - center_b + clean = (~mask_b).unsqueeze(-1).expand_as(delta) + if bool(clean.any().item()): + clean_ok = torch.where(clean, delta.abs(), torch.zeros_like(delta)).flatten(start_dim=1).amax(dim=1) <= tol + result &= clean_ok + perturbed_delta = delta[mask_b.unsqueeze(-1).expand_as(delta)].reshape(x_batch.shape[0], -1, center_t.shape[-1]) + if perturbed_delta.numel() == 0: + continue + if p_value == float("inf"): + norms = perturbed_delta.abs().amax(dim=-1) + elif p_value == 1.0: + norms = perturbed_delta.abs().sum(dim=-1) + elif p_value == 2.0: + norms = torch.linalg.vector_norm(perturbed_delta, ord=2, dim=-1) + else: + norms = torch.linalg.vector_norm(perturbed_delta, ord=p_value, dim=-1) + result &= (norms <= eps_t.reshape(-1)[0] + tol).all(dim=1) + return result diff --git a/act/back_end/core.py b/act/back_end/core.py index 7aa7f4b68..1df58820b 100644 --- a/act/back_end/core.py +++ b/act/back_end/core.py @@ -60,27 +60,8 @@ def get_bounds_for_var(self, fact: 'Fact', var_id: int, is_output: bool = True) return lb, ub - def get_all_var_bounds(self, fact: 'Fact', is_output: bool = True) -> Dict[int, Tuple[float, float]]: - """Get dict of {var_id: (lb, ub)} for all variables. Use is_output=True for out_vars, False for in_vars.""" - var_list = self.out_vars if is_output else self.in_vars - - bounds_dict = {} - for var_id in var_list: - bounds_dict[var_id] = self.get_bounds_for_var(fact, var_id, is_output) - - return bounds_dict - def get_input_shape(self) -> Optional[Tuple[int, ...]]: - """Get input shape from params, or None if not stored.""" - if self.kind == "INPUT": - return self.params.get("shape") - return self.params.get("input_shape") - def get_output_shape(self) -> Optional[Tuple[int, ...]]: - """Get output shape from params, or None if not stored.""" - if self.kind == "INPUT": - return self.params.get("shape") - return self.params.get("output_shape") def get_tensor(self, key: str, default: Optional[torch.Tensor] = None) -> Optional[torch.Tensor]: """Get tensor parameter. Returns default if missing or not a tensor.""" @@ -107,13 +88,7 @@ def is_tensor(self, key: str) -> bool: """Check if parameter is a tensor.""" return isinstance(self.params.get(key), torch.Tensor) - def get_num_input_vars(self) -> int: - """Get number of input variables (flattened dimension).""" - return len(self.in_vars) - def get_num_output_vars(self) -> int: - """Get number of output variables (flattened dimension).""" - return len(self.out_vars) @dataclass class Net: @@ -132,10 +107,6 @@ def __post_init__(self): validate_wrapper_graph(self.layers) # helpers - def last_validation(self) -> Optional[Layer]: - for L in reversed(self.layers): - if L.is_validation(): return L - return None def assert_last_is_validation(self) -> None: if not self.layers or not self.layers[-1].is_validation(): @@ -158,12 +129,6 @@ def get_all_predecessor_bounds(self, layer_id: int, after: Dict[int, 'Fact'], return [self.get_predecessor_bounds(layer_id, after, before, i) for i in range(len(self.preds[layer_id]))] - def get_layer_shape(self, layer_id: int, facts: Dict[int, 'Fact'], - is_output: bool = True) -> Tuple[int, ...]: - """Get shape from bounds tensor in facts dict. Use is_output=True for output shape, False for input.""" - if layer_id not in facts: - raise KeyError(f"Layer {layer_id} not found in facts dictionary") - return facts[layer_id].bounds.lb.shape def topological_sort(net: Net, reverse: bool = False) -> List[int]: @@ -251,38 +216,17 @@ def add_box(self, layer_id: int, var_ids: List[int], B: Bounds): "lb": B.lb.detach().reshape(-1).clone(), "ub": B.ub.detach().reshape(-1).clone()})) - def add_op(self, tag: str, var_ids: List[int], **meta): - """ - Generic operator constraint container. - - tag: e.g. "dense:12", "relu:5" - - var_ids: ordered exactly as exporter expects - - meta: payload used by cons_exportor.py - """ - op = tag.split(":", 1)[0] - if op and not self._is_op_supported_by_exporter(op): - raise ValueError( - f"Unknown op tag '{op}' (tag='{tag}'). " - "Update act/back_end/layer_schema.py SUPPORTED_EXPORT_OPS " - "and exporter handling if intentional." - ) - m = {"tag": tag} - m.update(meta) - self.replace(Con("INEQ", tuple(var_ids), m)) @staticmethod def _is_op_supported_by_exporter(op: str) -> bool: - """ - Best-effort early validation against exporter registry. - Falls back to allow if exporter cannot be imported. - """ - try: - mod = importlib.import_module("act.back_end.layer_util") - fn = getattr(mod, "is_supported_op", None) - if fn is None: - return True - return bool(fn(op)) - except Exception: - return True + """Validate an operator against the exporter registry.""" + mod = importlib.import_module("act.back_end.layer_util") + fn = getattr(mod, "is_supported_op", None) + if fn is None: + raise RuntimeError( + f"Exporter capability function is_supported_op is missing while checking {op!r}" + ) + return bool(fn(op)) def __iter__(self): """Iterate over constraints (Con objects). Makes ConSet iterable.""" diff --git a/act/back_end/dual_tf/tf_cnn.py b/act/back_end/dual_tf/tf_cnn.py index 73951b286..73c60faff 100644 --- a/act/back_end/dual_tf/tf_cnn.py +++ b/act/back_end/dual_tf/tf_cnn.py @@ -20,8 +20,8 @@ from .tf_forward import ( LinearBound, Frame, - _fwd_conv2d, _fwd_conv2d_interval, _fwd_maxpool2d, _fwd_avgpool2d, - _reset_forward_box, _concretize, + _fwd_conv2d, _fwd_conv2d_interval, _fwd_maxpool2d, _fwd_maxpool2d_lin, + _fwd_avgpool2d, _reset_forward_box, _concretize, _intersect_boxes, ) @@ -46,8 +46,15 @@ def forward_conv2d( Tries linear-relaxation conv via ``_fwd_conv2d``; when it returns ``None`` (kernel/stride shape unsupported), falls back to the interval conv ``_fwd_conv2d_interval`` and resets the dual-track state at the new - concrete box via ``_reset_forward_box``. ``stored`` equals ``out`` - (CONV2D has no activation split). + concrete box via ``_reset_forward_box``. On the symbolic path the + concretized box is intersected with the centre-radius interval image + (mirroring :func:`~act.back_end.dual_tf.tf_mlp.forward_dense`): the raw + concretization's lb/ub are two independently rounded conv accumulations, + so at collapsed box widths they cross by a few ulps of the PARTIAL-sum + magnitude (vgg16 late convs: partials 1e5-1e7, results O(1-100)), which + the backward degenerate check rejects. ``_intersect_boxes`` both tightens + and degrades such crossings to the structurally ordered interval box. + ``stored`` equals ``out`` (CONV2D has no activation split). """ assert len(parent_boxes) == 1, f"CONV2D expects 1 predecessor, got {len(parent_boxes)}" parent_box = parent_boxes[0] @@ -63,7 +70,9 @@ def forward_conv2d( lin, frame = _reset_forward_box(lb, ub, device, dtype) else: lin = new_lin - lb, ub = _concretize(lin, x_L, x_U) + lin_lb, lin_ub = _concretize(lin, x_L, x_U) + int_lb, int_ub = _fwd_conv2d_interval(L, parent_box.lb, parent_box.ub) + lb, ub = _intersect_boxes(lin_lb, lin_ub, int_lb, int_ub) out = Bounds(lb, ub) stored = out return stored, out, lin, frame @@ -190,48 +199,106 @@ def forward_maxpool2d( device: torch.device, dtype: torch.dtype, ) -> Tuple[Bounds, Bounds, LinearBound, Frame]: - """MaxPool2D forward bounds (interval-only; dual-track resets). - - Source: tf_forward.py lines 448-452 (pre-refactor monolithic MAXPOOL2D - branch). - - MaxPool has no linear relaxation here, so we compute interval bounds - via ``_fwd_maxpool2d`` and reset the dual-track state at the resulting - concrete box. ``stored`` equals ``out``. + """MaxPool2D forward bounds (DeepPoly dual-track when the parent is symbolic). + + When the parent carries an explicit linear bound, ``_fwd_maxpool2d_lin`` + gathers each window's argmax-lb row (exact upper bound where that input + dominates, constant interval max elsewhere) and the frame survives the + pool. The argmax indices and dominance mask are cached on ``L.cache`` for + the dual backward pass. On the lazy-identity path (parent ``A`` is None) + the previous behaviour is kept: interval bounds via ``_fwd_maxpool2d`` + and a dual-track reset at the concrete box. ``stored`` equals ``out``. """ assert len(parent_boxes) == 1, f"MAXPOOL2D expects 1 predecessor, got {len(parent_boxes)}" parent_box = parent_boxes[0] - lb, ub = _fwd_maxpool2d(L, parent_box.lb, parent_box.ub) + parent_lin = parent_lins[0] + parent_frame = parent_frames[0] + result = _fwd_maxpool2d_lin(L, parent_lin, parent_box.lb, parent_box.ub, parent_frame) + if result is None: + lb, ub = _fwd_maxpool2d(L, parent_box.lb, parent_box.ub) + out = Bounds(lb, ub) + stored = out + lin, frame = _reset_forward_box(lb, ub, device, dtype) + return stored, out, lin, frame + lin, idx_flat, dominant, box_lb, box_ub = result + # Exact comparisons that tolerate NaN bounds, which degrade to UNKNOWN downstream. + assert torch.isclose( + parent_box.lb.flatten(start_dim=1).gather(1, idx_flat), box_lb, + rtol=0.0, atol=0.0, equal_nan=True, + ).all(), f"MAXPOOL2D layer {L.id}: argmax index does not select the pooled lower bound" + assert torch.isclose( + parent_box.ub.flatten(start_dim=1).gather(1, idx_flat)[dominant], box_ub[dominant], + rtol=0.0, atol=0.0, equal_nan=True, + ).all(), f"MAXPOOL2D layer {L.id}: a dominant window's max upper bound is not at its argmax" + + x_L, x_U = parent_frame + lin_lb, lin_ub = _concretize(lin, x_L, x_U) + lb, ub = _intersect_boxes(lin_lb, lin_ub, box_lb, box_ub) + L.cache["maxpool_argmax_flat"] = idx_flat + L.cache["maxpool_dominant"] = dominant + # The mask is only valid for the box it was computed on. Cache the pooled + # output box actually stored so the backward pass can reject a stale mask + # left by a later forward on a different box (root_bounds_reuse modes run + # backward with refresh_forward=False against earlier bounds_dicts). + L.cache["maxpool_lb"] = lb.clone() + L.cache["maxpool_ub"] = ub.clone() out = Bounds(lb, ub) stored = out - lin, frame = _reset_forward_box(lb, ub, device, dtype) - return stored, out, lin, frame + return stored, out, lin, parent_frame -def backward_maxpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None): - """MaxPool2D backward — conservative constant-bound (sound but loose). +def _cache_box_matches(cached: Any, current: torch.Tensor) -> bool: + """True iff the forward-cached pooled box bit-matches the caller's bounds. - MaxPool is non-linear: ``y = max(x_window)``. For sound dual lower bound - on ``c @ output``, we use the dual decomposition ``nu @ y = nu_pos @ y + - nu_neg @ y >= nu_pos @ LB(y) + nu_neg @ UB(y)``, where: - - - LB(y) = max over window of lb_in = bounds_dict[L].lb (constant, sound: max(x) >= LB(max(x)) = max(lb)) - - UB(y) = max over window of ub_in = bounds_dict[L].ub (constant, sound: max(x) <= UB(max(x)) = max(ub)) - - Both bounds are CONSTANT in x → ``nu_in = 0`` (no propagation), full - contribution to obj via ``contrib = nu_pos @ lb_out + nu_neg @ ub_out``. - - This is the LOOSEST sound approach. Tighter alternatives: - - Linear-in-x at argmax_LB (one-hot): tighter LB but needs cached argmax - indices in forward (future enhancement). + ``bounds_dict`` entries are ``stored.copy()`` clones, so identity checks + cannot work; ``torch.equal`` on the ``[B, n_out]`` box is cheap. Shape and + dtype are guarded first because ``torch.equal`` raises across devices and + dtypes on some builds. + """ + if not isinstance(cached, torch.Tensor): + return False + if cached.dtype != current.dtype or cached.shape != current.shape: + return False + return bool(torch.equal(cached.to(device=current.device), current)) - Sound because: - LB(y) is sound: max(x_window) >= max_{k in window} lb[k] = bounds.lb[output_position] - UB(y) is sound: max(x_window) <= max_{k in window} ub[k] = bounds.ub[output_position] - Both are computed by forward_maxpool2d via F.max_pool2d on lb_in / ub_in. - Lazy M-broadcast: bounds_dict[L.id] is at [B, *shape]; nu is at - [B*M, *shape]. We broadcast the [B, 1, n] bounds against [B, M, n] nu. +def backward_maxpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None): + """MaxPool2D backward — one-hot pass-through at the argmax-lb input. + + Contract (shared by every dual backward handler, cf. the ReLU kernel): + returns ``(pred_nus, contrib)`` with ``nu @ y >= pred_nu @ x + contrib`` + pointwise on the box, so the full backward pass yields a LOWER bound on + ``c @ output``. + + Uses the cache written by :func:`forward_maxpool2d` via + ``_fwd_maxpool2d_lin``: ``maxpool_argmax_flat`` (long ``[B, n_out]``, flat + index into the pool input of each window's argmax-lb entry ``i*``) and + ``maxpool_dominant`` (bool ``[B, n_out]``, true where ``lb_{i*} >= + max_{j != i*} ub_j``). Per window and envelope side: + + - dominant: ``y = x_{i*}`` exactly, so ``nu`` (both signs) passes + through one-hot to ``i*``; zero contribution. + - non-dominant, ``nu >= 0`` (lower envelope): ``y >= x_{i*}`` pointwise + and ``nu >= 0`` give ``nu*y >= nu*x_{i*}`` — pass through to ``i*``, + zero contribution (matches the forward lower relaxation, which is the + argmax-lb input's own linear bound). + - non-dominant, ``nu < 0`` (upper envelope): the constant bound + ``y <= ub_out`` gives ``nu*y >= nu*ub_out`` — contribution only. + + Each window feeds ``v_out`` OR ``contrib`` per envelope side, never both. + Overlapping windows may target the same ``i*``; ``scatter_add_`` + accumulates them. + + Fallback (cache absent — forward ran the lazy-identity/interval path — + cached shapes do not match ``[B, n_out]``, or the cached pooled box + ``maxpool_lb``/``maxpool_ub`` is not bit-equal to ``bounds_dict[L.id]``, + i.e. the mask is stale from a forward on a DIFFERENT box under + root_bounds_reuse): the constant relaxation ``nu @ y = nu_pos @ y + + nu_neg @ y >= nu_pos @ LB(y) + nu_neg @ UB(y)`` with ``contrib = + nu_pos @ lb_out + nu_neg @ ub_out`` and ``v_out = 0``. + + Lazy M-broadcast: bounds_dict[L.id] and the cache are at [B, *shape]; nu + is at [B*M, *shape]. We broadcast the [B, 1, n] tensors against [B, M, n]. """ bounds = bounds_dict.get(L.id) if bounds is None: @@ -254,11 +321,6 @@ def backward_maxpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None): lb_b = lb_out_flat.unsqueeze(1) ub_b = ub_out_flat.unsqueeze(1) - nu_pos = v.clamp(min=0) - nu_neg = v.clamp(max=0) - contrib_BMn = nu_pos * lb_b + nu_neg * ub_b - contrib = contrib_BMn.sum(dim=-1).view(BM) - input_shape = L.params.get("input_shape") if not isinstance(input_shape, (list, tuple)): raise ValueError(f"backward_maxpool2d: layer {L.id} missing/invalid 'input_shape' param") @@ -270,6 +332,32 @@ def backward_maxpool2d(L, nu, bounds_dict, preds, M: int = 1, alpha=None): n_in = int(c_in) * int(iH) * int(iW) v_out = torch.zeros(BM, n_in, dtype=nu.dtype, device=nu.device) + idx_flat = L.cache.get("maxpool_argmax_flat") + dominant = L.cache.get("maxpool_dominant") + cached_lb = L.cache.get("maxpool_lb") + cached_ub = L.cache.get("maxpool_ub") + cache_ok = ( + isinstance(idx_flat, torch.Tensor) + and isinstance(dominant, torch.Tensor) + and idx_flat.shape == (B_actual, n) + and dominant.shape == (B_actual, n) + and _cache_box_matches(cached_lb, bounds.lb.flatten(start_dim=1)) + and _cache_box_matches(cached_ub, bounds.ub.flatten(start_dim=1)) + ) + if cache_ok: + dom = dominant.to(device=nu.device).unsqueeze(1) # [B, 1, n] + idx = idx_flat.to(device=nu.device).unsqueeze(1).expand(B_actual, M, n) + pass_nu = torch.where(dom, v, v.clamp(min=0)) # [B, M, n] + v_out.view(B_actual, M, n_in).scatter_add_(2, idx, pass_nu) + contrib_BMn = torch.where(dom, torch.zeros_like(v), v.clamp(max=0)) * ub_b + contrib = contrib_BMn.sum(dim=-1).view(BM) + else: + # A missing or stale cache is legitimate (lazy forward, root_bounds_reuse): no assert. + nu_pos = v.clamp(min=0) + nu_neg = v.clamp(max=0) + contrib_BMn = nu_pos * lb_b + nu_neg * ub_b + contrib = contrib_BMn.sum(dim=-1).view(BM) + assert len(preds) == 1, f"MAXPOOL2D expects 1 predecessor, got {len(preds)}" return [v_out], contrib diff --git a/act/back_end/dual_tf/tf_forward.py b/act/back_end/dual_tf/tf_forward.py index 871500945..f156be1c4 100644 --- a/act/back_end/dual_tf/tf_forward.py +++ b/act/back_end/dual_tf/tf_forward.py @@ -18,17 +18,21 @@ # pyright: reportMissingImports=false, reportUnknownVariableType=false, reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownArgumentType=false, reportAttributeAccessIssue=false, reportOptionalMemberAccess=false, reportMissingParameterType=false, reportUntypedFunctionDecorator=false, reportDeprecated=false +import logging from dataclasses import dataclass import torch import torch.nn.functional as F -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, cast from act.back_end.core import Bounds, Layer, Net, topological_sort from act.back_end.layer_schema import LayerKind -from act.back_end.utils import affine_bounds, split_weight +from act.back_end.utils import pair_2d +from act.config.config import DualConfig from act.util.device_manager import get_default_device, get_default_dtype +log = logging.getLogger(__name__) + @dataclass class LinearBound: @@ -47,6 +51,107 @@ class LinearBound: _DENSE_LIN_BOUND_MAX_DIM: int = 10_000 +_LIN_BOUND_FREE_FRACTION: float = 0.8 + +# Peak transient working set of one lane of the symbolic composition, as a +# multiple of that lane's persistent coefficient pair +# ``2 * n_out * n_sym * itemsize``. Measured as ``max_memory_allocated()`` +# deltas over the three helpers on toy CUDA layers (float32, B=4): CONV2D 2.38 +# (3->64 channels, VGG conv1 shape) and 3.45 (64->64), DENSE 2.25 (widening) +# and 4.02 (square), MAXPOOL2D 2.51 — each rounded up to the next integer. +# Allocations that do NOT scale with the lane count sit outside this model and +# no lane split can shrink them: ``W.abs()`` plus the einsum's weight copies +# for DENSE (they follow ``n_out * n_in``, and dominate whenever +# ``n_in >> n_out``), and the im2col buffer for spatially tiny convolutions. +_LIN_BOUND_CHUNK_TRANSIENT_CONV2D: int = 4 +_LIN_BOUND_CHUNK_TRANSIENT_DENSE: int = 4 +_LIN_BOUND_CHUNK_TRANSIENT_MAXPOOL2D: int = 3 + +_LIN_BOUND_REPORTED: set[Tuple[int, str]] = set() + + +def _report_lin_bound_once(layer_id: int, mode: str, level: int, msg: str, *args: object) -> None: + """Emit ``msg`` the first time a ``(layer id, mode)`` pair degrades. + + Forward bounds are recomputed for every BaB node, so an unconditional log + call here produces one line per node per layer. + """ + key = (layer_id, mode) + if key in _LIN_BOUND_REPORTED: + return + _LIN_BOUND_REPORTED.add(key) + log.log(level, msg, *args) + + +def _lin_bound_lane_chunk(B: int, n_out: int, n_sym: int, ref: torch.Tensor, + k_transient: int, layer_id: int, mode: str) -> Optional[int]: + """Plan how many lanes a symbolic ``[B, n_out, n_sym]`` layer may do at once. + + Returns ``B`` when the whole batch fits ~80% of the free CUDA memory (the + unguarded path, bit-for-bit unchanged), a chunk size ``1 <= c < B`` when the + batch must be split lane-wise, or ``None`` when even a single lane's working + set overruns the budget or the concatenated result cannot fit at all — the + caller then returns ``None`` and ``forward_X`` degrades that layer to the + interval box. CPU allocations are never split. + + Chunking only moves the *transient* peak: the result ``[B, n_out, n_sym]`` + pair is concatenated back and stays resident either way, so it is checked + against the full free memory separately. + """ + if not ref.is_cuda: + return B + pair_bytes = 2 * n_out * n_sym * ref.element_size() + if pair_bytes <= 0: + return B + free_bytes, _ = torch.cuda.mem_get_info(ref.device) + budget = free_bytes * _LIN_BOUND_FREE_FRACTION + persistent = B * pair_bytes + if persistent <= budget: + return B + + chunk = min(B, int(budget // (pair_bytes * k_transient))) + if chunk >= 1 and persistent <= free_bytes: + _report_lin_bound_once( + layer_id, mode, logging.INFO, + "forward linear bound splits %d lanes into chunks of %d for a " + "[%d, %d, %d] coefficient pair (%.2f GiB resident, %.2f GiB CUDA free)", + B, chunk, B, n_out, n_sym, persistent / 2 ** 30, free_bytes / 2 ** 30, + ) + return chunk + + _report_lin_bound_once( + layer_id, mode, logging.WARNING, + "forward linear bound needs ~%.2f GiB for a [%d, %d, %d] coefficient pair " + "(%.3f GiB for a single lane) with only %.2f GiB CUDA free; falling back " + "to interval for this layer", + persistent / 2 ** 30, B, n_out, n_sym, + pair_bytes * k_transient / 2 ** 30, free_bytes / 2 ** 30, + ) + return None + + +def _lin_lanes(lin: LinearBound, start: int, end: int) -> LinearBound: + """View of lanes ``[start, end)`` — slices, so no copy is made.""" + return LinearBound( + A_lb=None if lin.A_lb is None else lin.A_lb[start:end], + b_lb=lin.b_lb[start:end], + A_ub=None if lin.A_ub is None else lin.A_ub[start:end], + b_ub=lin.b_ub[start:end], + ) + + +def _cat_lanes(parts: List[LinearBound]) -> LinearBound: + """Reassemble lane chunks along the batch axis.""" + return LinearBound( + A_lb=None if parts[0].A_lb is None else torch.cat( + [cast(torch.Tensor, part.A_lb) for part in parts], dim=0), + b_lb=torch.cat([part.b_lb for part in parts], dim=0), + A_ub=None if parts[0].A_ub is None else torch.cat( + [cast(torch.Tensor, part.A_ub) for part in parts], dim=0), + b_ub=torch.cat([part.b_ub for part in parts], dim=0), + ) + + def _concretize(lin: LinearBound, x_L: torch.Tensor, x_U: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Concretize dual-track affine bounds over a batched input box. @@ -105,6 +210,33 @@ def _reset_lin(lb: torch.Tensor, ub: torch.Tensor, device, dtype B, n = lb.shape[0], lb.shape[1] return _identity_lin(B, n, device, dtype), lb.clone(), ub.clone() +def _entry_lin_frame(lb_in: torch.Tensor, ub_in: torch.Tensor, + max_perturbed: int, device, dtype + ) -> Tuple[LinearBound, Frame]: + """Entry LinearBound/Frame, symbolic only in the perturbed input dims. + + ``pert`` is the union over the batch of dims with ``ub > lb``. When + ``0 < n_pert <= max_perturbed``, the entry bound is the explicit one-hot + selection ``A[:, i, k] = 1`` iff input dim ``i`` is the ``k``-th perturbed + dim; unperturbed dims fold into the constant term ``b`` and the frame + shrinks to the perturbed columns. Downstream ``_fwd_X`` handlers then see + ``A is not None`` and never take the ``_DENSE_LIN_BOUND_MAX_DIM`` bail, + so e.g. VGG16 with <=100 perturbed pixels keeps linear bounds at every + conv. Otherwise (``n_pert == 0`` or above the cap) fall back to the lazy + identity over all dims, preserving the previous behaviour exactly. + """ + B, input_dim = lb_in.shape + pert = (ub_in > lb_in).any(dim=0) + n_pert = int(pert.sum().item()) + if n_pert == 0 or n_pert > max_perturbed: + return _identity_lin(B, input_dim, device, dtype), (lb_in, ub_in) + selection = torch.zeros(input_dim, n_pert, device=device, dtype=dtype) + selection[pert.nonzero(as_tuple=True)[0], torch.arange(n_pert, device=device)] = 1.0 + A = selection.unsqueeze(0).expand(B, input_dim, n_pert) + b = torch.where(pert, torch.zeros_like(lb_in), lb_in) + lin = LinearBound(A_lb=A, b_lb=b, A_ub=A, b_ub=b.clone()) + return lin, (lb_in[:, pert], ub_in[:, pert]) + def _match_lin_input_dim(lin: LinearBound, n_in: int) -> LinearBound: """Pad or truncate the current output-feature axis to size n_in. @@ -158,7 +290,31 @@ def _match_lin_input_dim(lin: LinearBound, n_in: int) -> LinearBound: def _intersect_boxes(lb_a: torch.Tensor, ub_a: torch.Tensor, lb_b: torch.Tensor, ub_b: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - return torch.maximum(lb_a, lb_b), torch.minimum(ub_a, ub_b) + """Intersect two sound boxes; EVERY crossing degrades to box ``b`` per entry. + + Both boxes enclose the same true value set, so in exact arithmetic they + overlap and ``max(lbs) <= min(ubs)``. In float32 the linear-track + concretization (box ``a``) and the interval track (box ``b``) are + independent accumulations; when the true width is below the accumulation + noise the pair can cross (RC2) — on layers whose partial sums dwarf the + cancelled result (late VGG convs / FC, partials 1e5-1e7 vs results + O(1-100)) by a small multiple of the PARTIAL-sum ulp. Any crossed entry, + whatever its magnitude, falls back to box ``b`` verbatim (the + centre-radius interval track, ordered by construction), which is sound + and merely looser. A tolerance-gated "repair" that swaps the crossed pair + to ``[min(ubs), max(lbs)]`` is UNSOUND: that interval lies between the two + boxes and may contain points of neither — e.g. true range [1.0, 2.0], + interval box [1.0001, 2.0001], linear box [0.0, 1.0] (1 ulp low at 1e5 + partial sums) swapped to [1.0, 1.0001], excluding the true value 2.0. + Genuine producer bugs that emit inverted boxes directly remain visible to + the downstream degenerate-interval check. + """ + lb = torch.maximum(lb_a, lb_b) + ub = torch.minimum(ub_a, ub_b) + crossed = lb > ub + out_lb = torch.where(crossed, lb_b, lb) + out_ub = torch.where(crossed, ub_b, ub) + return out_lb, out_ub def _align_batch(a: torch.Tensor, n: int) -> torch.Tensor: if a.shape[1] == n: @@ -178,14 +334,23 @@ def _int_param(value: object, default: int) -> int: def _box_dense(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Interval DENSE as ``m = c Wᵀ + b``, ``rho = r |W|ᵀ``, ``[m-rho, m+rho]``. + + ``rho >= 0`` structurally, so ``ub - lb = 2*rho >= 0`` per neuron whatever + the rounding; the ``W⁺l + W⁻u`` / ``W⁺u + W⁻l`` form sums two independent + accumulations that float32 cancellation can order the wrong way. + """ W = layer.params["weight"] b = layer.params.get("bias") lb = _align_batch(lb, W.shape[1]) ub = _align_batch(ub, W.shape[1]) - W_pos, W_neg = split_weight(W) - bias_vec = lb.new_zeros(W.shape[0]) if b is None else _align(b.flatten(), W.shape[0]) - out = affine_bounds(W_pos, W_neg, bias_vec, Bounds(lb, ub)) - return out.lb, out.ub + centre = (lb + ub) * 0.5 + radius = (ub - lb) * 0.5 + mid = centre @ W.T + spread = radius @ W.abs().T + if b is not None: + mid = mid + _align(b.flatten(), W.shape[0]) + return mid - spread, mid + spread def _box_bias(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: @@ -298,11 +463,22 @@ def _materialize(A: Optional[torch.Tensor], example: torch.Tensor) -> torch.Tens @torch.no_grad() def compute_forward_bounds(net: Net, input_lb: torch.Tensor, input_ub: torch.Tensor, post_activation: bool = False, - alphas: Optional[Dict[int, torch.Tensor]] = None) -> Dict[int, Bounds]: - """Forward bounds, natively batched with singleton auto-promotion.""" + alphas: Optional[Dict[int, torch.Tensor]] = None, + forward_lin_max_perturbed: Optional[int] = None, + ) -> Dict[int, Bounds]: + """Forward bounds, natively batched with singleton auto-promotion. + + ``forward_lin_max_perturbed`` caps the number of perturbed input dims for + which the linear track stays symbolic (see :func:`_entry_lin_frame`); + ``None`` resolves to the DualConfig default at call time rather than being + frozen at import time. + """ # Lazy import to break circular dep (dual_tf imports compute_forward_bounds) from .dual_tf import DualTF + if forward_lin_max_perturbed is None: + forward_lin_max_perturbed = DualConfig().forward_lin_max_perturbed + device, dtype = get_default_device(), get_default_dtype() if ( input_lb.dtype != dtype or input_lb.device != device @@ -318,7 +494,6 @@ def compute_forward_bounds(net: Net, input_lb: torch.Tensor, input_ub: torch.Ten B = input_lb.shape[0] lb_in = input_lb.reshape(B, -1) ub_in = input_ub.reshape(B, -1) - input_dim = lb_in.shape[1] bounds_dict: Dict[int, Bounds] = {} box_state: Dict[int, Bounds] = {} @@ -327,8 +502,9 @@ def compute_forward_bounds(net: Net, input_lb: torch.Tensor, input_ub: torch.Ten topo_order = topological_sort(net) by_id = getattr(net, "by_id", {layer.id: layer for layer in net.layers}) entry_box = Bounds(lb_in, ub_in) - entry_lin = _identity_lin(B, input_dim, device, dtype) - entry_frame = (lb_in, ub_in) + entry_lin, entry_frame = _entry_lin_frame( + lb_in, ub_in, forward_lin_max_perturbed, device, dtype, + ) for lid in topo_order: layer = by_id[lid] @@ -444,32 +620,64 @@ def _fwd_dense(layer: Layer, lin: LinearBound) -> Optional[LinearBound]: the caller to fall back to interval-only forward via :func:`_reset_forward_box`. Same pattern as :func:`_fwd_conv2d` and :func:`_fwd_relu`. + + Under CUDA memory pressure the batch is composed in lane chunks (see + :func:`_lin_bound_lane_chunk`) and reassembled; interval fallback is kept + only for the case where a single lane does not fit. """ W = layer.params["weight"] - b = layer.params.get("bias") lin = _match_lin_input_dim(lin, W.shape[1]) - W_pos = W.clamp(min=0) - W_neg = W.clamp(max=0) + B = lin.b_lb.shape[0] + n_sym = W.shape[1] if lin.A_lb is None else lin.A_lb.shape[2] + chunk = _lin_bound_lane_chunk(B, W.shape[0], n_sym, lin.b_lb, + _LIN_BOUND_CHUNK_TRANSIENT_DENSE, layer.id, "dense") + if chunk is None: + return None + if chunk >= B: + return _fwd_dense_lanes(layer, lin) + + parts: List[LinearBound] = [] + for start in range(0, B, chunk): + part = _fwd_dense_lanes(layer, _lin_lanes(lin, start, min(start + chunk, B))) + if part is None: + return None + parts.append(part) + return _cat_lanes(parts) + + +def _fwd_dense_lanes(layer: Layer, lin: LinearBound) -> Optional[LinearBound]: + """Dense composition for one lane chunk; ``lin`` must already be dim-matched.""" + W = layer.params["weight"] + b = layer.params.get("bias") B = lin.b_lb.shape[0] bias_vec = torch.zeros(W.shape[0], device=lin.b_lb.device, dtype=lin.b_lb.dtype) if b is not None: bias_vec = _align(b.flatten(), W.shape[0]) + # Centre-radius composition; see _fwd_conv2d for the identity and the + # float32 cancellation rationale. + b_centre = (lin.b_lb + lin.b_ub) * 0.5 + b_radius = (lin.b_ub - lin.b_lb) * 0.5 + b_mid = torch.einsum("oc,bc->bo", W, b_centre) + bias_vec + b_spread = torch.einsum("oc,bc->bo", W.abs(), b_radius) + if lin.A_lb is None and lin.A_ub is None: if W.shape[1] > _DENSE_LIN_BOUND_MAX_DIM: return None W_broadcast_lb = W.unsqueeze(0).expand(B, -1, -1).contiguous() W_broadcast_ub = W.unsqueeze(0).expand(B, -1, -1).contiguous() - b_lb_new = torch.einsum("oc,bc->bo", W_pos, lin.b_lb) + torch.einsum("oc,bc->bo", W_neg, lin.b_ub) + bias_vec - b_ub_new = torch.einsum("oc,bc->bo", W_pos, lin.b_ub) + torch.einsum("oc,bc->bo", W_neg, lin.b_lb) + bias_vec - return LinearBound(A_lb=W_broadcast_lb, b_lb=b_lb_new, - A_ub=W_broadcast_ub, b_ub=b_ub_new) + return LinearBound(A_lb=W_broadcast_lb, b_lb=b_mid - b_spread, + A_ub=W_broadcast_ub, b_ub=b_mid + b_spread) + A_centre = (lin.A_lb + lin.A_ub) * 0.5 + A_radius = (lin.A_ub - lin.A_lb) * 0.5 + A_mid = torch.einsum("oc,bci->boi", W, A_centre) + A_spread = torch.einsum("oc,bci->boi", W.abs(), A_radius) return LinearBound( - A_lb=torch.einsum("oc,bci->boi", W_pos, lin.A_lb) + torch.einsum("oc,bci->boi", W_neg, lin.A_ub), - b_lb=torch.einsum("oc,bc->bo", W_pos, lin.b_lb) + torch.einsum("oc,bc->bo", W_neg, lin.b_ub) + bias_vec, - A_ub=torch.einsum("oc,bci->boi", W_pos, lin.A_ub) + torch.einsum("oc,bci->boi", W_neg, lin.A_lb), - b_ub=torch.einsum("oc,bc->bo", W_pos, lin.b_ub) + torch.einsum("oc,bc->bo", W_neg, lin.b_lb) + bias_vec, + A_lb=A_mid - A_spread, + b_lb=b_mid - b_spread, + A_ub=A_mid + A_spread, + b_ub=b_mid + b_spread, ) @@ -647,8 +855,13 @@ def _fwd_lrelu(lin: LinearBound, lb: torch.Tensor, ub: torch.Tensor, alpha: floa def _fwd_conv2d(layer: Layer, lin: LinearBound) -> Optional[LinearBound]: - """Propagate dual-track affine bounds through Conv2D via batched F.conv2d.""" - weight = layer.params["weight"] + """Propagate dual-track affine bounds through Conv2D via batched F.conv2d. + + Under CUDA memory pressure the batch is convolved in lane chunks (see + :func:`_lin_bound_lane_chunk`) and reassembled; interval fallback is kept + only for the case where a single lane does not fit. + """ + weight = cast(torch.Tensor, layer.params["weight"]) bias = layer.params.get("bias") stride = layer.params.get("stride", 1) padding = layer.params.get("padding", 0) @@ -697,36 +910,85 @@ def _fwd_conv2d(layer: Layer, lin: LinearBound) -> Optional[LinearBound]: if in_h == 0 or in_w == 0: return None - W_pos = weight.clamp(min=0) - W_neg = weight.clamp(max=0) + k_h, k_w = int(weight.shape[2]), int(weight.shape[3]) + stride_i = cast(int, stride) + padding_i = cast(int, padding) + dilation_i = cast(int, dilation) + out_h = (in_h + 2 * padding_i - dilation_i * (k_h - 1) - 1) // stride_i + 1 + out_w = (in_w + 2 * padding_i - dilation_i * (k_w - 1) - 1) // stride_i + 1 + chunk = _lin_bound_lane_chunk(B, out_c * out_h * out_w, input_dim, lin.b_lb, + _LIN_BOUND_CHUNK_TRANSIENT_CONV2D, layer.id, "conv2d") + if chunk is None: + return None def conv_A(A_mat: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor: - A_t = A_mat.transpose(1, 2).contiguous().view(B * input_dim, in_c, in_h, in_w) + lanes = A_mat.shape[0] + A_t = A_mat.transpose(1, 2).contiguous().view(lanes * input_dim, in_c, in_h, in_w) out = F.conv2d(A_t, kernel, None, stride, padding, dilation, groups) - return out.flatten(start_dim=1).reshape(B, input_dim, -1).transpose(1, 2).contiguous() + return out.flatten(start_dim=1).reshape(lanes, input_dim, -1).transpose(1, 2).contiguous() def conv_b(vec: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor: - b_4d = vec.view(B, in_c, in_h, in_w) + b_4d = vec.view(vec.shape[0], in_c, in_h, in_w) return F.conv2d(b_4d, kernel, None, stride, padding, dilation, groups).flatten(start_dim=1) - A_lb_new = conv_A(lin.A_lb, W_pos) + conv_A(lin.A_ub, W_neg) - A_ub_new = conv_A(lin.A_ub, W_pos) + conv_A(lin.A_lb, W_neg) - b_lb_new = conv_b(lin.b_lb, W_pos) + conv_b(lin.b_ub, W_neg) - b_ub_new = conv_b(lin.b_ub, W_pos) + conv_b(lin.b_lb, W_neg) - - if bias is not None: - out_spatial = b_lb_new.shape[1] // out_c - bias_bc = bias.view(out_c, 1).expand(out_c, out_spatial).reshape(-1) - b_lb_new = b_lb_new + bias_bc - b_ub_new = b_ub_new + bias_bc - - return LinearBound(A_lb=A_lb_new, b_lb=b_lb_new, A_ub=A_ub_new, b_ub=b_ub_new) + # Centre-radius composition, algebraically identical to the split-sign + # form: W+ @ A_lb + W- @ A_ub == W @ Ac - |W| @ Ar with Ac=(A_lb+A_ub)/2, + # Ar=(A_ub-A_lb)/2 (and the mirrored identity for the upper track). The + # split-sign form runs four INDEPENDENT accumulations whose float32 + # rounding crosses by ~n*eps*sum|W||b| once the tracks nearly coincide + # (RC2's linear-track sibling; vgg spec2 inverted by 1e2 at ~1e5 + # magnitudes). Sharing the centre term keeps bit-equal input tracks + # bit-equal on output, so concretized boxes cannot invert along affine + # chains. cuDNN is disabled for the radius convs: Winograd's subtractive + # transforms emit small negatives for a zero/nonnegative radius (see + # _fwd_conv2d_interval). + def compose(sub: LinearBound) -> LinearBound: + A_centre = (sub.A_lb + sub.A_ub) * 0.5 + A_radius = (sub.A_ub - sub.A_lb) * 0.5 + b_centre = (sub.b_lb + sub.b_ub) * 0.5 + b_radius = (sub.b_ub - sub.b_lb) * 0.5 + A_mid = conv_A(A_centre, weight) + b_mid = conv_b(b_centre, weight) + with torch.backends.cudnn.flags(enabled=False): + A_spread = conv_A(A_radius, weight.abs()) + b_spread = conv_b(b_radius, weight.abs()) + A_lb_new = A_mid - A_spread + A_ub_new = A_mid + A_spread + b_lb_new = b_mid - b_spread + b_ub_new = b_mid + b_spread + + if bias is not None: + out_spatial = b_lb_new.shape[1] // out_c + bias_bc = bias.view(out_c, 1).expand(out_c, out_spatial).reshape(-1) + b_lb_new = b_lb_new + bias_bc + b_ub_new = b_ub_new + bias_bc + + return LinearBound(A_lb=A_lb_new, b_lb=b_lb_new, A_ub=A_ub_new, b_ub=b_ub_new) + + if chunk >= B: + return compose(lin) + return _cat_lanes([ + compose(_lin_lanes(lin, start, min(start + chunk, B))) + for start in range(0, B, chunk) + ]) def _fwd_conv2d_interval(layer: Layer, lb: torch.Tensor, ub: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: - """Fallback interval Conv2D for when the linear-relaxation path cannot infer shape.""" - weight = layer.params["weight"] + """Fallback interval Conv2D for when the linear-relaxation path cannot infer shape. + + Centre-radius form: ``m = c ⊛ W + b``, ``rho = r ⊛ |W|``, ``[m-rho, m+rho]``. + ``rho >= 0`` structurally, so ``ub - lb = 2*rho >= 0`` per neuron whatever + the rounding; the ``W⁺l + W⁻u`` / ``W⁺u + W⁻l`` form sums two independent + 4608-term accumulations that float32 cancellation can order the wrong way + (VGG16 L30: 11 515 inverted neurons, worst excess 3.9e-1). Two convs, not four. + + The radius conv runs with cuDNN disabled: cuDNN picks Winograd for 3x3, whose + transforms contain subtractions, so it returns small *negative* values for a + non-negative input/kernel pair (measured -1.3 where neighbouring tiles carry a + large radius). Plain accumulation keeps ``rho >= 0`` exact, hence ``ub >= lb``. + """ + weight = cast(torch.Tensor, layer.params["weight"]) bias = layer.params.get("bias") stride = layer.params.get("stride", 1) padding = layer.params.get("padding", 0) @@ -761,23 +1023,21 @@ def _fwd_conv2d_interval(layer: Layer, lb: torch.Tensor, ub: torch.Tensor in_h = in_w = side try: - lb_4d = lb.view(B, in_c, in_h, in_w) - ub_4d = ub.view(B, in_c, in_h, in_w) + centre = ((lb + ub) * 0.5).view(B, in_c, in_h, in_w) + radius = ((ub - lb) * 0.5).view(B, in_c, in_h, in_w) except RuntimeError as e: raise ValueError( f"_fwd_conv2d_interval: reshape to [B={B}, {in_c}, {in_h}, {in_w}] " f"failed for lb.shape={tuple(lb.shape)}" ) from e - W_pos, W_neg = split_weight(weight) conv_kw = dict(stride=stride, padding=padding, dilation=dilation, groups=groups) - lb_out = F.conv2d(lb_4d, W_pos, None, **conv_kw) + F.conv2d(ub_4d, W_neg, None, **conv_kw) - ub_out = F.conv2d(ub_4d, W_pos, None, **conv_kw) + F.conv2d(lb_4d, W_neg, None, **conv_kw) + mid = F.conv2d(centre, weight, None, **conv_kw) + with torch.backends.cudnn.flags(enabled=False): + spread = F.conv2d(radius, weight.abs(), None, **conv_kw) if bias is not None: - bias_4d = bias.view(1, -1, 1, 1) - lb_out = lb_out + bias_4d - ub_out = ub_out + bias_4d - return lb_out.flatten(start_dim=1), ub_out.flatten(start_dim=1) + mid = mid + bias.view(1, -1, 1, 1) + return (mid - spread).flatten(start_dim=1), (mid + spread).flatten(start_dim=1) def _fwd_maxpool2d(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: @@ -803,6 +1063,135 @@ def _fwd_maxpool2d(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[to return lb_out.flatten(start_dim=1), ub_out.flatten(start_dim=1) +def _fwd_maxpool2d_lin( + layer: Layer, lin: LinearBound, lb: torch.Tensor, ub: torch.Tensor, + frame: Frame, +) -> Optional[Tuple[LinearBound, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + """DeepPoly MaxPool2D on the dual-track affine bounds (explicit ``A`` only). + + Per window let ``i*`` be the argmax of the interval lower bounds. Since + ``max(x_window) >= x_{i*}`` always, input ``i*``'s affine lower bound is a + sound lower bound for the pool output. When ``i*`` dominates — + ``lb_{i*} >= max_{j != i*} ub_j`` — the pool output *is* ``x_{i*}``, so its + affine upper bound is exact; otherwise the constant ``max_j ub_j`` is used, + lifted to also cover the concretized max of the gathered lower row over + ``frame``. ``ub_max`` comes from the interval track while the lower row + comes from the linear track; once the box width collapses (deep input + splits) the two independent float32 accumulations can cross by a few ulps, + and a crossed (lower_fn, upper_const) pair is amplified by every following + affine layer's |W| mass until the backward's degenerate-interval check + raises. Raising the upper constant to ``max(ub_max, max_frame lower_fn)`` + is sound (uppers may only grow) and restores the exact-arithmetic + invariant ``upper >= lower`` pointwise on the frame. + ``ub_second`` (largest ub among the window's other inputs) is computed via + ``F.unfold`` with the argmax slot and zero-filled padding slots masked to + ``-inf``. + + Under CUDA memory pressure the gathers run in lane chunks (see + :func:`_lin_bound_lane_chunk`); the window statistics above them do not + scale with ``n_sym`` and stay batched. + + Returns ``None`` on the lazy-identity path (``A is None``), on a feature + width mismatch, or when a single lane's gathered ``[1, n_out, n_sym]`` + coefficients would overrun CUDA memory — the caller then keeps the + interval-reset behaviour. Otherwise returns ``(lin_out, idx_flat, dominant, lb_max, + ub_max)`` where ``idx_flat`` is the flat argmax index per output (long + ``[B, n_out]``), ``dominant`` is a bool ``[B, n_out]`` mask, and + ``lb_max``/``ub_max`` are the interval pool box ``[B, n_out]``. + """ + if lin.A_lb is None or lin.A_ub is None: + return None + kernel_size = pair_2d(layer.params.get("kernel_size", 2)) + stride = pair_2d(layer.params.get("stride", layer.params.get("kernel_size", 2))) + padding = pair_2d(layer.params.get("padding", 0)) + dilation = pair_2d(layer.params.get("dilation", 1)) + input_shape = _shape_list(layer.params.get("input_shape")) + if input_shape is None: + raise ValueError( + f"_fwd_maxpool2d_lin: layer {layer.id} missing required 'input_shape' param" + ) + shape = input_shape + if len(shape) == 4: + _, c, h, w = shape + else: + c, h, w = shape[-3], shape[-2], shape[-1] + B, n_in = lb.shape + if n_in != c * h * w or lin.A_lb.shape[1] != n_in or lin.b_lb.shape[1] != n_in: + return None + + lb_4d = lb.view(B, c, h, w) + ub_4d = ub.view(B, c, h, w) + lb_max_4d, idx_plane = F.max_pool2d( + lb_4d, kernel_size, stride, padding, dilation, return_indices=True + ) + ub_max_4d = F.max_pool2d(ub_4d, kernel_size, stride, padding, dilation) + n_out = c * lb_max_4d.shape[2] * lb_max_4d.shape[3] + n_sym = lin.A_lb.shape[2] + chunk = _lin_bound_lane_chunk(B, n_out, n_sym, lin.b_lb, + _LIN_BOUND_CHUNK_TRANSIENT_MAXPOOL2D, layer.id, "maxpool2d") + if chunk is None: + return None + + # Window membership via im2col: map each kernel slot to its input plane + # index. Unfold zero-fills padding, so shift indices by one to tell a + # padded slot (0) from plane position 0; float64 keeps indices exact. + plane = torch.arange( + 1, h * w + 1, device=lb.device, dtype=torch.float64 + ).view(1, 1, h, w) + plane_unf = F.unfold( + plane, kernel_size, dilation=dilation, padding=padding, stride=stride + ) + pad_slot = plane_unf == 0 + slot_plane = plane_unf.long() - 1 + k2, spatial_out = plane_unf.shape[1], plane_unf.shape[2] + + # Largest ub among the window's OTHER inputs: mask the argmax-lb slot and + # the padded slots (unfold zero-fill) to -inf, then reduce over the kernel. + ub_unf = F.unfold( + ub_4d, kernel_size, dilation=dilation, padding=padding, stride=stride + ).view(B, c, k2, spatial_out) + is_argmax = slot_plane.view(1, 1, k2, spatial_out) == idx_plane.view(B, c, 1, spatial_out) + ub_second = ub_unf.masked_fill( + is_argmax | pad_slot.view(1, 1, k2, spatial_out), float("-inf") + ).amax(dim=2) + + lb_max = lb_max_4d.flatten(start_dim=1) + ub_max = ub_max_4d.flatten(start_dim=1) + dominant = (lb_max_4d.view(B, c, spatial_out) >= ub_second).view(B, n_out) + + channel_base = (torch.arange(c, device=lb.device) * (h * w)).view(1, c, 1) + idx_flat = (idx_plane.view(B, c, spatial_out) + channel_base).view(B, n_out) + + x_L, x_U = frame + + def gather_lanes(start: int, end: int) -> LinearBound: + idx = idx_flat[start:end] + gather_idx = idx.unsqueeze(-1).expand(end - start, n_out, n_sym) + dom = dominant[start:end] + A_lb_out = cast(torch.Tensor, lin.A_lb)[start:end].gather(1, gather_idx) + b_lb_out = lin.b_lb[start:end].gather(1, idx) + A_ub_out = cast(torch.Tensor, lin.A_ub)[start:end].gather(1, gather_idx).masked_fill( + ~dom.unsqueeze(-1), 0.0) + lower_row_max = ( + torch.einsum("boi,bi->bo", A_lb_out.clamp(min=0), x_U[start:end]) + + torch.einsum("boi,bi->bo", A_lb_out.clamp(max=0), x_L[start:end]) + + b_lb_out + ) + b_ub_out = torch.where( + dom, lin.b_ub[start:end].gather(1, idx), + torch.maximum(ub_max[start:end], lower_row_max), + ) + return LinearBound(A_lb=A_lb_out, b_lb=b_lb_out, A_ub=A_ub_out, b_ub=b_ub_out) + + if chunk >= B: + lin_out = gather_lanes(0, B) + else: + lin_out = _cat_lanes([ + gather_lanes(start, min(start + chunk, B)) for start in range(0, B, chunk) + ]) + return lin_out, idx_flat, dominant, lb_max, ub_max + + def _fwd_avgpool2d(layer: Layer, lb: torch.Tensor, ub: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """AvgPool2D interval propagation with runtime batch size.""" kernel_size = layer.params.get("kernel_size", 2) diff --git a/act/back_end/hybridz_tf/hybridz_tf.py b/act/back_end/hybridz_tf/hybridz_tf.py index afad49158..bdd706316 100644 --- a/act/back_end/hybridz_tf/hybridz_tf.py +++ b/act/back_end/hybridz_tf/hybridz_tf.py @@ -388,20 +388,17 @@ def _sparse_fact(fact: Fact, hz: SparseHZono) -> Fact: def _seed_sparse_cache(self, L: Layer, input_bounds: Bounds) -> None: k = L.kind.upper() - try: - if k in ("INPUT", "INPUT_SPEC"): + if k in ("INPUT", "INPUT_SPEC"): + self._sparse_hz_cache[L.id] = self._sparse_from_bounds(input_bounds) + self._sparse_drop_reasons.pop(L.id, None) + elif k != "ASSERT": + preds = self._net.preds.get(L.id, []) + if preds and preds[0] in self._sparse_hz_cache: + self._sparse_hz_cache[L.id] = self._sparse_hz_cache[preds[0]] + self._sparse_drop_reasons.pop(L.id, None) + elif not preds: self._sparse_hz_cache[L.id] = self._sparse_from_bounds(input_bounds) self._sparse_drop_reasons.pop(L.id, None) - elif k != "ASSERT": - preds = self._net.preds.get(L.id, []) - if preds and preds[0] in self._sparse_hz_cache: - self._sparse_hz_cache[L.id] = self._sparse_hz_cache[preds[0]] - self._sparse_drop_reasons.pop(L.id, None) - elif not preds: - self._sparse_hz_cache[L.id] = self._sparse_from_bounds(input_bounds) - self._sparse_drop_reasons.pop(L.id, None) - except Exception as exc: - self._drop_sparse_hz(L.id, f"sparse_seed_failed:{type(exc).__name__}") def _drop_sparse_hz(self, layer_id: int, reason: str) -> None: lid = int(layer_id) @@ -423,24 +420,21 @@ def _propagate_sparse_hz(self, L: Layer, input_bounds: Bounds, result: Fact) -> if self._sparse_exceeds_limit(hz, result.bounds.lb.numel()): self._drop_sparse_hz(L.id, f"sparse_size_limit:{k}") return result - try: - for apply_sparse in ( - hz_mlp.sparse_hz_apply_layer, - hz_cnn.sparse_hz_apply_layer, - hz_transformer.sparse_hz_apply_layer, - ): - handled, out, drop_reason = apply_sparse(L, hz, input_bounds, result, self) - if not handled: - continue - if out is None: - self._drop_sparse_hz(L.id, drop_reason or f"unsupported_sparse_op:{k}") - return result - self._sparse_hz_cache[L.id] = out - self._sparse_drop_reasons.pop(L.id, None) - return self._sparse_fact(result, out) - self._drop_sparse_hz(L.id, f"unsupported_sparse_op:{k}") - except Exception as exc: - self._drop_sparse_hz(L.id, f"sparse_op_failed:{k}:{type(exc).__name__}") + for apply_sparse in ( + hz_mlp.sparse_hz_apply_layer, + hz_cnn.sparse_hz_apply_layer, + hz_transformer.sparse_hz_apply_layer, + ): + handled, out, drop_reason = apply_sparse(L, hz, input_bounds, result, self) + if not handled: + continue + if out is None: + self._drop_sparse_hz(L.id, drop_reason or f"unsupported_sparse_op:{k}") + return result + self._sparse_hz_cache[L.id] = out + self._sparse_drop_reasons.pop(L.id, None) + return self._sparse_fact(result, out) + self._drop_sparse_hz(L.id, f"unsupported_sparse_op:{k}") return result def apply( diff --git a/act/back_end/hybridz_tf/tf_mlp.py b/act/back_end/hybridz_tf/tf_mlp.py index 213acf56f..e697d4dc4 100644 --- a/act/back_end/hybridz_tf/tf_mlp.py +++ b/act/back_end/hybridz_tf/tf_mlp.py @@ -167,6 +167,7 @@ def _matmul_term_rows(L, batch: int): y_shape = tuple(int(d) for d in L.params["y_shape"]) if len(x_shape) < 2 or len(y_shape) < 2 or x_shape[-1] != y_shape[-2]: return None + rows = None try: batch_shape = np.broadcast_shapes(x_shape[:-2], y_shape[:-2]) m, reduction, n = x_shape[-2], x_shape[-1], y_shape[-1] @@ -185,8 +186,13 @@ def _matmul_term_rows(L, batch: int): np.swapaxes(y_local, -2, -1)[..., None, :, :], (*batch_shape, m, n, reduction), ).reshape(-1, reduction) + rows = (x_terms, y_terms, reduction) except (ValueError, TypeError): + # Non-broadcastable MATMUL metadata is not applicable to sparse lifting. + rows = None + if rows is None: return None + x_terms, y_terms, reduction = rows x_size, y_size = _prod(x_shape), _prod(y_shape) offsets = np.arange(int(batch), dtype=np.int64)[:, None, None] return ( @@ -1169,12 +1175,14 @@ def _row_indices_expand(L, n: int): if per == 0 or int(n) % per != 0: return None batch = int(n) // per - try: - return torch.arange(int(n)).view(batch, *in_shape).broadcast_to( - batch, *out_shape - ).reshape(-1) - except RuntimeError: + source_shape = (batch, *in_shape) + target_shape = (batch, *out_shape) + if len(source_shape) > len(target_shape) or any( + source not in (1, target) + for source, target in zip(reversed(source_shape), reversed(target_shape)) + ): return None + return torch.arange(int(n)).view(source_shape).broadcast_to(target_shape).reshape(-1) def _row_indices_reduce(L, n_in: int, n_out: int, axes_key: str, keepdims_key: str): diff --git a/act/back_end/hybridz_tf/tf_transformer.py b/act/back_end/hybridz_tf/tf_transformer.py index 977cd8b39..2fd4d3ba9 100644 --- a/act/back_end/hybridz_tf/tf_transformer.py +++ b/act/back_end/hybridz_tf/tf_transformer.py @@ -157,13 +157,16 @@ def _attention_score_differences(L, bounds: Bounds, rowsize: int, tf): continue if kind in {"SCALE", "BIAS"}: key = "a" if kind == "SCALE" else "c" - try: - value = ( - _broadcast_flat(score_layer.params[key], n_scores) - .detach().cpu().double().numpy() - ) - except ValueError: + raw_value = score_layer.params[key] + value_count = int(raw_value.numel()) + if value_count == 0 or ( + value_count not in (1, n_scores) and n_scores % value_count != 0 + ): return None + value = ( + _broadcast_flat(raw_value, n_scores) + .detach().cpu().double().numpy() + ) if kind == "SCALE": scale *= value else: diff --git a/act/back_end/layer_util.py b/act/back_end/layer_util.py index 425d0e106..bf5eda83c 100644 --- a/act/back_end/layer_util.py +++ b/act/back_end/layer_util.py @@ -44,11 +44,8 @@ # Will import at runtime when needed to avoid circular import pass -try: - import torch - Tensor = torch.Tensor -except Exception: # typing only - Tensor = "torch.Tensor" # type: ignore +import torch +Tensor = torch.Tensor # ------------------------------ # Strict validation & helpers @@ -250,54 +247,49 @@ def create_layer(id: int, kind: str, params: Dict[str, Any], # Tiny example (run file) # --------------------- if __name__ == "__main__": - try: - import torch # type: ignore - import sys - import os - # Add parent directory to path for absolute imports - from act.util.path_config import get_project_root - project_root = get_project_root() - sys.path.insert(0, project_root) + import sys + # Add parent directory to path for absolute imports + from act.util.path_config import get_project_root + project_root = get_project_root() + sys.path.insert(0, project_root) - from act.back_end.core import Layer - from act.back_end.layer_schema import LayerKind - from typing import List - layers: List[Layer] = [] + from act.back_end.core import Layer + from act.back_end.layer_schema import LayerKind + from typing import List + layers: List[Layer] = [] - # INPUT - layers.append(create_layer( + # INPUT + layers.append(create_layer( id=0, kind=LayerKind.INPUT.value, params={}, in_vars=[0], out_vars=[0], - )) - # SPEC (directly after INPUT - no adapters) - lb_tensor = torch.full((1,3,32,32), -1.0) - ub_tensor = torch.full((1,3,32,32), 1.0) - layers.append(create_layer( + )) + # SPEC (directly after INPUT - no adapters) + lb_tensor = torch.full((1,3,32,32), -1.0) + ub_tensor = torch.full((1,3,32,32), 1.0) + layers.append(create_layer( id=1, kind=LayerKind.INPUT_SPEC.value, params={"kind": InKind.BOX, "lb": lb_tensor, "ub": ub_tensor}, in_vars=[0], out_vars=[0], - )) - # Model toy - layers.append(create_layer( + )) + # Model toy + layers.append(create_layer( id=2, kind=LayerKind.FLATTEN.value, params={}, in_vars=[0], out_vars=[1], - )) - W, b = torch.randn(10, 3072), torch.randn(10) - layers.append(create_layer( + )) + W, b = torch.randn(10, 3072), torch.randn(10) + layers.append(create_layer( id=3, kind=LayerKind.DENSE.value, params={"W": W, "b": b}, in_vars=[1], out_vars=[2], - )) - layers.append(create_layer( + )) + layers.append(create_layer( id=4, kind=LayerKind.ASSERT.value, params={}, in_vars=[2], out_vars=[2], - )) + )) - validate_graph(layers) - validate_wrapper_graph(layers) - print("OK — wrapper model passes with", len(layers), "layers.") - except Exception as e: - print("Example failed:\n", e) \ No newline at end of file + validate_graph(layers) + validate_wrapper_graph(layers) + print("OK — wrapper model passes with", len(layers), "layers.") diff --git a/act/back_end/net_factory.py b/act/back_end/net_factory.py index 38ce0dbec..b9e0c11e0 100644 --- a/act/back_end/net_factory.py +++ b/act/back_end/net_factory.py @@ -296,15 +296,6 @@ def append_binary_op(layers, *, op_kind: str, x_idx: int, y_idx: int) -> None: ) -def append_concat(layers, *, input_indices: List[int], concat_dim: int = 0) -> None: - _append_layer( - layers, - LayerKind.CONCAT.value, - params={"concat_dim": concat_dim}, - preds=input_indices, - ) - - def append_flatten(layers) -> None: _append_layer(layers, LayerKind.FLATTEN.value, params={"start_dim": 1}) @@ -1186,11 +1177,7 @@ def __init__( self.total_generated = 0 def _compute_allowed_layers(self, tf_targets, mode): - try: - return _get_allowed_layers(tf_targets, mode) - except Exception as e: - logger.warning("TF capabilities unavailable: %s. Using defaults.", e) - return frozenset(_DEFAULT_COVERAGE_LAYERS) + return _get_allowed_layers(tf_targets, mode) def _init_coverage(self): skip = { diff --git a/act/back_end/serialization/serialization.py b/act/back_end/serialization/serialization.py index d1d797fdf..2f28c0c11 100644 --- a/act/back_end/serialization/serialization.py +++ b/act/back_end/serialization/serialization.py @@ -142,34 +142,14 @@ def deserialize_layer(layer_dict: Dict[str, Any], target_device: Optional[str] = # This is a regular value cache_decoded[name] = value - try: - layer = Layer( - id=layer_dict["id"], - kind=layer_dict["kind"], - params=params_decoded, - in_vars=layer_dict["in_vars"], - out_vars=layer_dict["out_vars"], - cache=cache_decoded - ) - except ValueError as e: - # If core validation fails during deserialization, create the Layer - # without validation for template/example usage - if "schema violation" in str(e): - print(f"⚠️ Creating example/template layer {layer_dict['kind']}(id={layer_dict['id']}) without validation") - - # Create Layer without validation - layer = object.__new__(Layer) - layer.id = layer_dict["id"] - layer.kind = layer_dict["kind"] - layer.params = params_decoded - layer.in_vars = layer_dict["in_vars"] - layer.out_vars = layer_dict["out_vars"] - layer.cache = cache_decoded - else: - # Re-raise for other types of errors - raise - - return layer + return Layer( + id=layer_dict["id"], + kind=layer_dict["kind"], + params=params_decoded, + in_vars=layer_dict["in_vars"], + out_vars=layer_dict["out_vars"], + cache=cache_decoded + ) class NetSerializer: @@ -213,22 +193,8 @@ def deserialize_net(net_dict: Dict[str, Any], target_device: Optional[str] = Non preds = {int(k): v for k, v in graph.get("preds", {}).items()} succs = {int(k): v for k, v in graph.get("succs", {}).items()} - # Create Net object with validation - try: - return Net(layers=layers, preds=preds, succs=succs) - except ValueError as e: - # If validation fails during deserialization, create the Net - # without validation for template/example usage - if "schema violation" in str(e) or "validation" in str(e).lower(): - print(f"Warning: Creating net without validation") - net = object.__new__(Net) - net.layers = layers - net.preds = preds - net.succs = succs - net.by_id = {L.id: L for L in layers} - return net - else: - raise + # Create Net object with validation; malformed serialized graphs are errors. + return Net(layers=layers, preds=preds, succs=succs) class ACTJSONEncoder(json.JSONEncoder): """Custom JSON encoder for ACT objects.""" diff --git a/act/back_end/solver/solver_dual.py b/act/back_end/solver/solver_dual.py index f61f154cf..4f18a6f3b 100644 --- a/act/back_end/solver/solver_dual.py +++ b/act/back_end/solver/solver_dual.py @@ -18,7 +18,13 @@ from act.back_end.core import Bounds, Layer, Net, get_topo_order from act.back_end.layer_schema import LayerKind from act.back_end.dual_tf.tf_forward import _intersect_boxes -from act.back_end.solver.solver_base import Solver, SolverCaps +from act.config.config import DualConfig +from act.back_end.solver.solver_base import ( + BatchLPSolution, + SolveStatus, + Solver, + SolverCaps, +) from act.front_end.specs import OutputSpec, OutKind from act.util.device_manager import get_default_device, get_default_dtype from act.util.stats import SpecBatchResult @@ -38,6 +44,25 @@ class DualResult: nu_per_layer: Optional[Dict[int, torch.Tensor]] = None +@dataclass(frozen=True) +class DualBatchResult: + """Decoded result of :meth:`DualSolver.solve_spec_batch` for one K-lane batch.""" + + solution: BatchLPSolution + margins: torch.Tensor + lower_bounds: torch.Tensor + bounds_dict: Optional[Dict[int, Bounds]] = None + nu_per_layer: Optional[Dict[int, torch.Tensor]] = None + alpha_state: Optional[Dict[int, torch.Tensor]] = None + eta_state: Optional[Dict[int, torch.Tensor]] = None + witness_input: Optional[torch.Tensor] = None + row_slack: Optional[torch.Tensor] = None + c_rows: Optional[torch.Tensor] = None + thresholds: Optional[torch.Tensor] = None + m_specs: int = 0 + reference_bounds: Optional[Dict[int, Bounds]] = None + + def expand_bounds_dict(bounds_dict: Dict[int, Bounds], M: int) -> Dict[int, Bounds]: """Expand each batched Bounds entry from [B, *shape] to [B*M, *shape]. @@ -170,6 +195,225 @@ def __init__(self, n_iters: int = 0): def capabilities(self) -> SolverCaps: return SolverCaps(supports_gpu=True, supports_csp=False, supports_dual=True) + def solve_spec_batch( + self, + net: Net, + bounds_dict: Dict[int, Bounds], + out_spec: OutputSpec, + *, + optimize: bool, + dual_config: DualConfig, + input_shape: tuple[int, ...], + keep_rows: Optional[torch.Tensor] = None, + split_signs: Optional[Dict[int, torch.Tensor]] = None, + eta: Optional[Dict[int, torch.Tensor]] = None, + incremental_alphas: Optional[Dict[int, torch.Tensor]] = None, + incremental_etas: Optional[Dict[int, torch.Tensor]] = None, + optimize_alpha: bool = True, + refresh_forward: bool = True, + return_nu: bool = False, + reuse_bounds_for_branching: bool = False, + ) -> DualBatchResult: + """Solve and decode one prepared K-lane output-spec batch. + + Bound construction and refinement are caller policy. This method owns + output-row encoding, dual optimization, lane status/witness decoding, + and the optional converged-state nu pass used by neuron branching. + """ + sample_bounds = next(iter(bounds_dict.values())) + device = sample_bounds.lb.device + dtype = sample_bounds.lb.dtype + if sample_bounds.lb.dim() < 2: + raise ValueError( + "DualSolver.solve_spec_batch: bounds_dict entries must be batched " + f"[K, *shape]; got dim={sample_bounds.lb.dim()}" + ) + k_actual = int(sample_bounds.lb.shape[0]) + + assert_layers = [ + layer + for layer in net.layers + if (layer.kind.upper() if isinstance(layer.kind, str) else layer.kind) + == LayerKind.ASSERT.value + ] + if not assert_layers: + raise ValueError("DualSolver.solve_spec_batch: net has no ASSERT layer") + assert_layer = assert_layers[0] + assert_preds = net.preds.get(assert_layer.id, []) + if len(assert_preds) != 1: + raise ValueError( + f"ASSERT layer {assert_layer.id} must have exactly 1 predecessor, " + f"got {len(assert_preds)}" + ) + output_bounds = bounds_dict[assert_preds[0]] + n_out = int(output_bounds.lb.flatten(start_dim=1).shape[-1]) + encoded_spec = out_spec.encode_linear( + B=k_actual, n_out=n_out, device=device, dtype=dtype + ) + m_specs = int(encoded_spec["M"]) + + if out_spec.kind == OutKind.UNSAFE_LINEAR: + c_rows = cast(torch.Tensor, encoded_spec["C"]).contiguous() + thresholds = cast(torch.Tensor, encoded_spec["thresholds"]).contiguous() + else: + c_rows = -cast(torch.Tensor, encoded_spec["C"]).contiguous() + thresholds = -cast(torch.Tensor, encoded_spec["thresholds"]).contiguous() + if keep_rows is not None: + idx = keep_rows.to(device=device, dtype=torch.long) + c_rows = ( + c_rows.view(k_actual, m_specs, n_out) + .index_select(1, idx) + .reshape(k_actual * int(idx.numel()), n_out) + .contiguous() + ) + thresholds = thresholds.index_select(1, idx).contiguous() + m_specs = int(idx.numel()) + active_mask = torch.ones( + k_actual, m_specs, dtype=torch.bool, device=device + ) + + compute_certified_bound = cast(Any, self.compute_certified_bound) + if optimize: + dual_result = compute_certified_bound( + net, + bounds_dict, + c_rows, + M=m_specs, + optimize=True, + optimize_alpha=optimize_alpha, + refresh_forward=refresh_forward, + forward_lin_max_perturbed=dual_config.forward_lin_max_perturbed, + n_iters=dual_config.n_iters, + lr_alpha=dual_config.lr_alpha, + lr_beta=dual_config.lr_beta, + lr_decay=dual_config.lr_decay, + eta=eta, + incremental_alphas=incremental_alphas, + incremental_etas=incremental_etas, + split_signs=split_signs, + return_optimized=True, + return_sce=True, + per_class_alpha=dual_config.per_class_alpha, + **({"return_nu_per_layer": True} if return_nu else {}), + ) + else: + dual_result = compute_certified_bound( + net, + bounds_dict, + c_rows, + M=m_specs, + return_sce=True, + **({"return_nu_per_layer": True} if return_nu else {}), + ) + + margins = dual_result.margins.view(k_actual, m_specs) + sce = cast(Optional[torch.Tensor], dual_result.sce) + slack = margins - thresholds + if out_spec.kind == OutKind.UNSAFE_LINEAR: + certified = ((slack > 0) & active_mask).any(dim=-1) + candidate_rows = torch.zeros( + k_actual, dtype=torch.long, device=device + ) + else: + violations = (slack < 0) & active_mask + certified = ~violations.any(dim=-1) + candidate_rows = torch.where( + violations.any(dim=1), + violations.to(torch.int64).argmax(dim=1), + torch.zeros(k_actual, dtype=torch.long, device=device), + ) + + statuses = tuple( + SolveStatus.UNSAT if bool(is_certified.item()) else SolveStatus.SAT + for is_certified in certified + ) + nvars = ( + max((max(layer.out_vars) for layer in net.layers if layer.out_vars), default=-1) + + 1 + ) + x_candidate = torch.zeros(k_actual, nvars, device=device, dtype=dtype) + input_layers = [ + layer + for layer in net.layers + if (layer.kind.upper() if isinstance(layer.kind, str) else layer.kind) + == LayerKind.INPUT.value + ] + if len(input_layers) != 1: + raise ValueError( + f"Expected exactly one INPUT layer, found {len(input_layers)}." + ) + input_ids_list = list(input_layers[0].out_vars) + input_ids = torch.tensor(input_ids_list, device=device, dtype=torch.long) + if sce is not None: + sce_flat = sce.flatten(start_dim=1).to(device=device) + row_offsets = ( + torch.arange(k_actual, device=device) * m_specs + + candidate_rows.to(device=device) + ) + chosen_sce = sce_flat.index_select(0, row_offsets) + x_candidate[:, input_ids] = chosen_sce.to(device=device, dtype=dtype) + else: + statuses = tuple( + SolveStatus.UNSAT if status == SolveStatus.UNSAT else SolveStatus.UNKNOWN + for status in statuses + ) + solution = BatchLPSolution( + statuses=statuses, + x=x_candidate, + max_viol=-slack.min(dim=1).values.detach(), + ) + + branch_bounds: Optional[Dict[int, Bounds]] = None + branch_nu: Optional[Dict[int, torch.Tensor]] = None + if return_nu and reuse_bounds_for_branching: + branch_bounds = bounds_dict + branch_nu = dual_result.nu_per_layer + if branch_nu is None: + nu_pass = self.compute_certified_bound( + net, + bounds_dict, + c_rows, + M=m_specs, + alpha=dual_result.alpha_state, + eta=dual_result.eta_state if split_signs is not None else None, + split_signs=split_signs, + return_nu_per_layer=True, + ) + branch_nu = nu_pass.nu_per_layer + elif return_nu: + branch_bounds, branch_nu = self.recompute_bounds_and_nu( + net, + bounds_dict, + c_rows, + m_specs, + alpha_state=dual_result.alpha_state, + eta_state=dual_result.eta_state if split_signs is not None else None, + split_signs=split_signs, + per_class_alpha=dual_config.per_class_alpha, + ) + + witness_input = ( + x_candidate[:, input_ids].reshape(k_actual, *input_shape).detach() + if sce is not None + else None + ) + lower_bounds = -solution.max_viol + return DualBatchResult( + solution=solution, + margins=margins, + lower_bounds=lower_bounds, + bounds_dict=branch_bounds, + nu_per_layer=branch_nu, + alpha_state=dual_result.alpha_state, + eta_state=dual_result.eta_state, + witness_input=witness_input, + row_slack=slack.detach(), + c_rows=c_rows.detach(), + thresholds=thresholds.detach(), + m_specs=m_specs, + reference_bounds=bounds_dict, + ) + def compute_certified_bound( self, net: Net, bounds_dict: Dict[int, Bounds], c: torch.Tensor, M: int = 1, @@ -191,6 +435,8 @@ def compute_certified_bound( optimize_alpha: bool = True, refresh_forward: bool = True, start_lid: Optional[int] = None, + local_phase_clamp: bool = False, + forward_lin_max_perturbed: Optional[int] = None, ) -> DualResult: """Batched certified lower bound on c^T @ output (DAG-aware). @@ -227,6 +473,9 @@ def compute_certified_bound( (RELU, LRELU, SIGMOID, TANH, GELU). η ≥ 0 invariant (enforced by clamp). split_signs: Per-layer split direction. {-1: inactive, +1: active, 0: unsplit}. Same key set as eta. + local_phase_clamp: For fixed-parameter replay, shallow-copy and clamp only + the current ReLU handler's bounds. Reference bounds remain immutable; + optimization is rejected because it requires globally hardened bounds. η is applied to the TRUE pre-activation variable (immediately AFTER the activation handler in the reverse-topological backward loop): nu_pre = slope · nu_post − η · sign, so the multiplier acts on the @@ -234,7 +483,12 @@ def compute_certified_bound( scale η by the relaxation slope, forcing the effective multiplier to 0 on inactive-split (slope = 0) neurons and discarding the z ≤ 0 constraint's input-region information. Sound for any η ≥ 0. + + forward_lin_max_perturbed: ``None`` resolves to the DualConfig default + at call time rather than being frozen at import time. """ + if forward_lin_max_perturbed is None: + forward_lin_max_perturbed = DualConfig().forward_lin_max_perturbed if isinstance(split_signs, list): # KFSB path: accept K split hypotheses and return stacked margins # [K, N, ...], evaluating each hypothesis through the single-hypothesis @@ -270,6 +524,8 @@ def compute_certified_bound( return_optimized=False, per_class_alpha=per_class_alpha, return_nu_per_layer=False, + local_phase_clamp=local_phase_clamp, + forward_lin_max_perturbed=forward_lin_max_perturbed, ) margins.append(result.margins) if return_sce: @@ -279,7 +535,10 @@ def compute_certified_bound( stacked_sce = torch.stack(cast(List[torch.Tensor], sce_values), dim=0) return DualResult(margins=torch.stack(margins, dim=0), sce=stacked_sce) - bounds_dict = self._harden_split_bounds(bounds_dict, split_signs) + if optimize and local_phase_clamp: + raise ValueError("local_phase_clamp is only valid for fixed-parameter replay") + if not local_phase_clamp: + bounds_dict = self._harden_split_bounds(bounds_dict, split_signs) if optimize: bound, sce, alpha_state, eta_state = self._optimize_alpha_eta( @@ -299,6 +558,7 @@ def compute_certified_bound( optimize_alpha=optimize_alpha, refresh_forward=refresh_forward, start_lid=start_lid, + forward_lin_max_perturbed=forward_lin_max_perturbed, ) if return_optimized: return DualResult( @@ -384,11 +644,21 @@ def compute_certified_bound( nu_snapshot[lid] = nu_here.detach().clone() preds = list(net.preds.get(lid, [])) + handler_bounds = bounds_dict + if ( + local_phase_clamp + and k == LayerKind.RELU.value + and split_signs is not None + and lid in split_signs + ): + handler_bounds = self._local_phase_bounds( + bounds_dict, lid, split_signs[lid] + ) if alpha is None: - pred_nus, contrib = handler(layer, nu_here, bounds_dict, preds, M) + pred_nus, contrib = handler(layer, nu_here, handler_bounds, preds, M) else: pred_nus, contrib = handler( - layer, nu_here, bounds_dict, preds, M, alpha=alpha.get(lid) + layer, nu_here, handler_bounds, preds, M, alpha=alpha.get(lid) ) if eta is not None and split_signs is not None and lid in eta: @@ -573,10 +843,13 @@ def _init_attention_alpha( from act.back_end.dual_tf.tf_transformer import ( attention_rule_alpha, _attention_input_boxes, ) - try: - x_l, x_u, y_l, y_u, _scale, _mask = _attention_input_boxes(layer, bounds_dict) - except KeyError: + kind = layer.kind.upper() if isinstance(layer.kind, str) else layer.kind + src_keys = ("q_src", "k_src") if kind == "ATT_SCORES" else ("w_src", "v_src") + if any(int(layer.params[key]) not in bounds_dict for key in src_keys): + # Missing predecessor bounds disable learned attention slopes. return None + boxes = _attention_input_boxes(layer, bounds_dict) + x_l, x_u, y_l, y_u, _scale, _mask = boxes x_l = x_l.to(device=device, dtype=dtype) x_u = x_u.to(device=device, dtype=dtype) y_l = y_l.to(device=device, dtype=dtype) @@ -611,6 +884,7 @@ def _optimize_alpha_eta( optimize_alpha: bool = True, refresh_forward: bool = True, start_lid: Optional[int] = None, + forward_lin_max_perturbed: Optional[int] = None, ) -> Tuple[ torch.Tensor, Optional[torch.Tensor], @@ -628,7 +902,12 @@ def _optimize_alpha_eta( ``best_bounds`` has shape ``[B*M]``, ``best_sce`` is optional, ``alpha_state`` maps ReLU layer id to optimized α, and ``eta_state`` maps split layer id to optimized η. + + ``forward_lin_max_perturbed=None`` resolves to the DualConfig default + at call time rather than being frozen at import time. """ + if forward_lin_max_perturbed is None: + forward_lin_max_perturbed = DualConfig().forward_lin_max_perturbed if c.dim() != 2: raise ValueError( f"c must be 2-D [B*M, n_out], got shape {tuple(c.shape)}" @@ -773,6 +1052,7 @@ def _optimize_alpha_eta( input_ub, post_activation=False, alphas=forward_alphas, + forward_lin_max_perturbed=forward_lin_max_perturbed, ) else: # Fixed intermediate bounds (root-reuse mode): alpha/eta @@ -794,6 +1074,21 @@ def _optimize_alpha_eta( bound_bm = result.margins sce = result.sce + if not bound_bm.requires_grad: + # No autograd path reaches any α/η parameter (every ReLU + # stable and every pool window dominant): gradient steps + # cannot move the bound, and backward() would raise. + # Return the fixed-slope bound as optimize=False would. + detached = bound_bm.detach() + improved = detached > best_bounds + best_bounds = torch.where(improved, detached, best_bounds) + if return_sce and sce is not None: + if best_sce is None: + best_sce = sce.detach().clone() + else: + best_sce[improved] = sce[improved].detach() + return best_bounds, best_sce, best_alpha_state, best_eta_state + (-bound_bm.sum()).backward() optimizer.step() if scheduler is not None: @@ -839,6 +1134,86 @@ def _optimize_alpha_eta( return best_bounds.detach(), best_sce, best_alpha_state, best_eta_state + def _interval_refresh_bounds( + self, + net: Net, + base: Dict[int, Bounds], + split_signs: Dict[int, torch.Tensor], + ) -> Optional[Dict[int, Bounds]]: + """Re-propagate split phases with interval arithmetic and intersect base.""" + from act.back_end.dual_tf.tf_forward import _fwd_conv2d_interval + + out = dict(base) + vals: Dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + for layer in net.layers: + k = layer.kind.upper() if isinstance(layer.kind, str) else layer.kind + lid = layer.id + if k == LayerKind.ASSERT.value: + continue + if k in (LayerKind.INPUT.value, LayerKind.INPUT_SPEC.value): + b = out.get(lid) + if b is None: + return None + vals[lid] = (b.lb.flatten(start_dim=1), b.ub.flatten(start_dim=1)) + continue + preds = net.preds.get(lid, []) + if k == LayerKind.CONV2D.value: + plb, pub = vals[preds[0]] + lb, ub = _fwd_conv2d_interval(layer, plb, pub) + lb, ub = lb.flatten(start_dim=1), ub.flatten(start_dim=1) + elif k == LayerKind.DENSE.value: + w = layer.params["weight"] + bias = layer.params.get("bias") + if not isinstance(w, torch.Tensor): + return None + plb, pub = vals[preds[0]] + c = (plb + pub) * 0.5 + r = (pub - plb) * 0.5 + m = c @ w.T + rho = r @ w.abs().T + lb = m - rho + ub = m + rho + if isinstance(bias, torch.Tensor): + lb, ub = lb + bias, ub + bias + elif k == LayerKind.ADD.value: + (alb, aub), (blb, bub) = vals[preds[0]], vals[preds[1]] + lb, ub = alb + blb, aub + bub + elif k in (LayerKind.FLATTEN.value, LayerKind.RESHAPE.value): + lb, ub = vals[preds[0]] + elif k == LayerKind.RELU.value: + lb, ub = vals[preds[0]] + else: + return None + + b = out.get(lid) + if b is not None: + lb = torch.maximum(lb, b.lb.flatten(start_dim=1)) + ub = torch.minimum(ub, b.ub.flatten(start_dim=1)) + ub = torch.maximum(ub, lb) + if k == LayerKind.RELU.value: + s = split_signs.get(lid) + if s is not None: + sl = s[:, 0, :] if s.dim() == 3 else s + n = min(lb.shape[-1], sl.shape[-1]) + sn = sl[..., :n].to(lb.device) + lb, ub = lb.clone(), ub.clone() + lb[..., :n] = torch.where( + sn > 0, lb[..., :n].clamp(min=0.0), lb[..., :n] + ) + ub[..., :n] = torch.where( + sn < 0, ub[..., :n].clamp(max=0.0), ub[..., :n] + ) + ub[..., :n] = torch.maximum(ub[..., :n], lb[..., :n]) + if b is not None: + out[lid] = Bounds( + lb.view_as(b.lb).clone(), ub.view_as(b.ub).clone() + ) + if k == LayerKind.RELU.value: + vals[lid] = (lb.clamp(min=0.0), ub.clamp(min=0.0)) + else: + vals[lid] = (lb, ub) + return out + def _harden_split_bounds( self, bounds_dict: Dict[int, Bounds], @@ -875,6 +1250,34 @@ def _harden_split_bounds( out[lid] = Bounds(lb.view_as(b.lb), ub.view_as(b.ub)) return out + def _local_phase_bounds( + self, + bounds_dict: Dict[int, Bounds], + lid: int, + signs: torch.Tensor, + ) -> Dict[int, Bounds]: + """Clamp one ReLU entry for its handler without mutating reference bounds.""" + bounds = bounds_dict.get(lid) + if bounds is None: + return bounds_dict + phase = signs[:, 0, :] if signs.dim() == 3 else signs + if not bool((phase != 0).any().item()): + return bounds_dict + lb = bounds.lb.flatten(start_dim=1).clone() + ub = bounds.ub.flatten(start_dim=1).clone() + n = min(lb.shape[-1], phase.shape[-1]) + phase = phase[..., :n].to(device=lb.device) + lb[..., :n] = torch.where( + phase > 0, lb[..., :n].clamp(min=0.0), lb[..., :n] + ) + ub[..., :n] = torch.where( + phase < 0, ub[..., :n].clamp(max=0.0), ub[..., :n] + ) + ub[..., :n] = torch.maximum(ub[..., :n], lb[..., :n]) + local = dict(bounds_dict) + local[lid] = Bounds(lb.view_as(bounds.lb), ub.view_as(bounds.ub)) + return local + def refine_intermediate_bounds( self, net: Net, diff --git a/act/back_end/solver/solver_hz.py b/act/back_end/solver/solver_hz.py index e3be99a33..c1ddb25dd 100644 --- a/act/back_end/solver/solver_hz.py +++ b/act/back_end/solver/solver_hz.py @@ -969,40 +969,18 @@ def softmax_ratio_weighted_extreme( ordered_upper = np.take_along_axis(score_upper, order, axis=1).astype( np.longdouble ) - shift = np.max(ordered_upper, axis=1, keepdims=True) - exp_lower = np.exp(ordered_lower - shift) - exp_upper = np.exp(ordered_upper - shift) - prefix_num = np.concatenate( - [ - np.zeros((values.shape[0], 1), dtype=np.longdouble), - np.cumsum(ordered_objective * exp_upper, axis=1), - ], - axis=1, - ) - prefix_den = np.concatenate( - [ - np.zeros((values.shape[0], 1), dtype=np.longdouble), - np.cumsum(exp_upper, axis=1), - ], - axis=1, - ) - lower_num = np.concatenate( - [ - np.zeros((values.shape[0], 1), dtype=np.longdouble), - np.cumsum(ordered_objective * exp_lower, axis=1), - ], - axis=1, - ) - lower_den = np.concatenate( - [ - np.zeros((values.shape[0], 1), dtype=np.longdouble), - np.cumsum(exp_lower, axis=1), - ], - axis=1, - ) - candidates = (prefix_num + lower_num[:, -1:] - lower_num) / ( - prefix_den + lower_den[:, -1:] - lower_den + candidates = np.empty( + (values.shape[0], values.shape[1] + 1), dtype=np.longdouble ) + for split in range(values.shape[1] + 1): + vertex_scores = np.concatenate( + [ordered_upper[:, :split], ordered_lower[:, split:]], axis=1 + ) + vertex_shift = np.max(vertex_scores, axis=1, keepdims=True) + vertex_exp = np.exp(vertex_scores - vertex_shift) + candidates[:, split] = np.sum( + ordered_objective * vertex_exp, axis=1 + ) / np.sum(vertex_exp, axis=1) vertex = np.asarray(np.min(candidates, axis=1), dtype=np.float64) vertex -= 32.0 * np.finfo(np.float64).eps * (1.0 + np.abs(vertex)) if minimize: @@ -1548,6 +1526,10 @@ def weighted_mass(capacity, target): inequalities = np.empty((signs.shape[0], width + 1), dtype=np.float64) inequalities[:, :width] = -signed_weights[row] inequalities[:, -1] = -1.0 + assert np.all(np.isfinite(extrema[row])), ( + "HybridZ softmax invariant violated: non-finite weighted extrema " + f"at row {row}" + ) result = linprog( objective, A_ub=inequalities, @@ -2019,7 +2001,7 @@ def _hz_bounds_unconstrained(hz: HZono) -> Bounds: return Bounds(lb=(hz.c - rad).reshape(1, -1), ub=(hz.c + rad).reshape(1, -1)) -def _hz_compute_bounds_scipy(hz: HZono) -> Bounds: +def _hz_compute_bounds_scipy(hz: HZono) -> Optional[Bounds]: model = _lower_hz_milp(hz) if model.n_var == 0: LB = UB = model.value_center.copy() @@ -2049,7 +2031,7 @@ def _hz_compute_bounds_scipy(hz: HZono) -> Bounds: options=options, ) if not res_min.success or not res_max.success: - raise RuntimeError(f"HybridZ bound MILP failed at output {i}") + return None LB[i] = model.value_center[i] + res_min.fun UB[i] = model.value_center[i] - res_max.fun @@ -2080,11 +2062,9 @@ def hz_compute_bounds(hz: HZono, *, exact: bool = False) -> Bounds: if not exact: return _hz_bounds_unconstrained(hz) if _HAS_SCIPY: - try: - return _hz_compute_bounds_scipy(hz) - except Exception as e: - # Intentional: scipy linprog failures fall back to the unconstrained bounds estimate. - logger.debug("suppressed: %s", e) + exact_bounds = _hz_compute_bounds_scipy(hz) + if exact_bounds is not None: + return exact_bounds return _hz_bounds_unconstrained(hz) @@ -2310,21 +2290,17 @@ def _solve_hz_feasibility( constraints = ( LinearConstraint(A, row_lb, row_ub) if A.shape[0] else None ) - try: - result = milp( - c=np.zeros(model.n_var, dtype=np.float64), - integrality=model.integrality, - bounds=SciPyBounds(model.var_lb, model.var_ub), - constraints=constraints, - options={ - "presolve": True, - "time_limit": max(1e-3, remaining), - "mip_rel_gap": 0.0, - }, - ) - except Exception as exc: - logger.debug("HybridZ MILP failed: %s", exc) - return _MILPResult("unknown", None) + result = milp( + c=np.zeros(model.n_var, dtype=np.float64), + integrality=model.integrality, + bounds=SciPyBounds(model.var_lb, model.var_ub), + constraints=constraints, + options={ + "presolve": True, + "time_limit": max(1e-3, remaining), + "mip_rel_gap": 0.0, + }, + ) nodes = int(getattr(result, "mip_node_count", 0) or 0) x = getattr(result, "x", None) if x is not None and _valid_milp_point( @@ -2614,10 +2590,7 @@ def evaluate_spec( return self._unknown_results(B, "missing_hz_state") if not _HAS_SCIPY: return self._unknown_results(B, "scipy_unavailable") - try: - model = _lower_hz_milp(output_hz) - except Exception as exc: - return self._unknown_results(B, f"lowering_failed:{type(exc).__name__}") + model = _lower_hz_milp(output_hz) if model.value_center.size != B * int(n_out): return self._unknown_results(B, "output_shape_mismatch") @@ -2642,8 +2615,9 @@ def evaluate_spec( if not is_unsafe_linear and highspy is not None: try: warm_lp = _HighsLPRelaxation(model) - except Exception as exc: - logger.debug("HybridZ warm LP setup failed: %s", exc) + except RuntimeError: + # HiGHS setup can reject a model; MILP remains a sound fallback. + warm_lp = None def solve(extra_A=None, extra_lb=None, extra_ub=None) -> _MILPResult: nonlocal solves, nodes diff --git a/act/back_end/transfer_functions.py b/act/back_end/transfer_functions.py index fb1294efb..6de1a6df2 100644 --- a/act/back_end/transfer_functions.py +++ b/act/back_end/transfer_functions.py @@ -130,15 +130,12 @@ def ensure_active_tf(default_mode: str = "interval") -> TransferFunction: """Return the active TF, self-healing to ``default_mode`` if unset. Consolidates the self-heal pattern shared by analyze.py / cli.py / - verifier.py. Catches only ``RuntimeError`` (the documented - "not initialised" case); any other exception from a broken TF subclass - bubbles up so the bug surfaces rather than silently demoting to interval. + verifier.py. Only the "not initialised" state is healed; nothing is + caught, so a broken TF subclass surfaces rather than demoting to interval. """ - try: - return get_transfer_function() - except RuntimeError: + if _current_tf is None: set_transfer_function_mode(default_mode) - return get_transfer_function() + return get_transfer_function() def set_transfer_function_mode(mode: str = "interval", tf_config: Any = None) -> None: @@ -176,10 +173,6 @@ def set_solver_mode(mode: Optional[str]) -> None: _current_solver_mode = mode -def get_solver_mode() -> Optional[str]: - return _current_solver_mode - - def is_dual_solver_active() -> bool: return _current_solver_mode == "dual" diff --git a/act/back_end/verifier.py b/act/back_end/verifier.py index 6756f8e88..0ce7db5d6 100644 --- a/act/back_end/verifier.py +++ b/act/back_end/verifier.py @@ -404,9 +404,9 @@ def verify_lp_batched( ) violations = torch.zeros(batch_size, device=x_candidates.device, dtype=torch.bool) if bool(sat_mask.any().item()): - bab_module = importlib.import_module("act.back_end.bab.bab") + violation_module = importlib.import_module("act.back_end.bab.violation") sat_idx = torch.where(sat_mask)[0] - checked_sat = bab_module.check_violations_batched( + checked_sat = violation_module.check_violations_batched( net, x_candidates.index_select(0, sat_idx), assert_layer, ) if checked_sat.shape != (int(sat_idx.numel()),): diff --git a/act/config/backend.yaml b/act/config/backend.yaml index 28ff189ef..4a8599915 100644 --- a/act/config/backend.yaml +++ b/act/config/backend.yaml @@ -4,12 +4,70 @@ # Canonical source for ALL back-end settings (runtime, verification, generation). # CI (.github/workflows/act-bab.yml) relies on these defaults. # -# Precedence: CLI flag > environment variable > this file > code default +# Precedence: CLI flag > env var > BaB preset > this file > code default # Env vars: ACT_SOLVER / ACT_DEVICE / ACT_DTYPE # -# Layout (all under the single top-level `backend:` mapping): -# runtime selectors → verification cascade → torchlp → tf → hybridz -# → bab (branch-and-bound) +# Layout: +# bab_presets – named sparse BaB configurations +# backend – runtime selectors → verification cascade → torchlp → gurobi +# → hybridz → dual → bab (branch-and-bound) + +# ───────────────────────────────────────────────────────────────────────────── +# Branch-and-Bound presets +# ───────────────────────────────────────────────────────────────────────────── +# A preset is a named, sparse set of backend.bab overrides. Select one with +# backend.bab_preset or --bab-preset NAME; --solver dual applies "dual" when no +# preset is named. Preset values replace the backend.bab values below, except +# for BaB fields set by explicit CLI flags, which keep the flag's value. +bab_presets: + dual: + # Optimize the relaxation slopes (alpha) and the split multipliers (eta). + solver_tier: "dual_alpha_eta" + # Split the neuron with the best measured child-bound gain (babsr fallback). + branching_method: "gain" + # Children reuse the root's forward bounds, hardened by the split-derived + # interval refresh, instead of running a forward pass per node. + root_bounds_reuse: "split_refresh" + # Tighten the root's pre-activation bounds on every unstable layer. + intermediate_refine: "all" + # Split 4 neurons jointly: each lane gets 2^4 = 16 children. + multi_split_levels: 4 + + # CLIMB: query-local certificate replay, core learning and phase propagation. + # When the dual bound certifies a subproblem, its alpha/eta certificate is + # replayed to measure what each split decision (ReLU phase literal) it relied + # on contributes, and literals are dropped while the budget + # theta * slack - delta covers them (the first recheck_k lanes get one extra + # replay). The surviving literals form a core, kept in a per-query library of + # at most max_cores cores (subsumption, resolution, eviction). Pending + # subproblems whose split signs contain a core are discharged before + # bounding, and unresolved ones before splitting. + # validate_climb_config raises ConfigError unless solver_tier is + # dual_alpha_eta, bounding is a TopKBounding pool, branching splits neurons, + # and root_bounds_reuse, intermediate_refine and per_subproblem_refine are + # "none". + climb: + # Replay needs the alpha/eta dual certificate. + solver_tier: "dual_alpha_eta" + # Cores are ReLU phase literals; an input-axis split would break the + # shared-box premise of core reuse. + branching_method: "babsr" + # CLIMB prunes the pending pool, which must be a TopKBounding pool. + bounding: "depth_bound_blend" + # Prerequisite: children re-run the forward pass (no root-bound reuse). + root_bounds_reuse: "none" + # Prerequisite: no intermediate-bound refinement at the root. + intermediate_refine: "none" + # Prerequisite: no per-subproblem intermediate-bound refinement. + per_subproblem_refine: "none" + # Turn CLIMB on (off in the backend.bab defaults). + climb_enabled: true + # Spend up to half of each row's replay slack on dropping literals. + climb_theta: 0.5 + # The first 2 lanes get one extra replay to try dropping one more literal. + climb_recheck_k: 2 + # One split per step, so small nets still learn cores. + multi_split_levels: 1 backend: @@ -46,10 +104,16 @@ backend: # Verification cascade (single-shot → LP-batched → BaB) # ───────────────────────────────────────────────────────────────────────── - # Wall-clock budget in seconds for both single execution and BaB verification. + # Wall-clock budget in seconds. The LP-batched tier and each BaB run get the + # full budget; with solver hybridz it also bounds the single-shot solve unless + # hybridz.timeout is set. timeout: 60.0 - # Enable the batched LP cascade tier (tier 2) before BaB. + # Named bab_presets entry; null auto-selects "dual" when solver = "dual". + bab_preset: null + + # Enable the batched LP cascade tier (tier 2) before BaB. It runs on lanes still + # UNKNOWN after the single-shot tier and is skipped for solver dual / hybridz. # Must be false when solver = "gurobi" (Gurobi solve_batch is N=1 only). lp_enabled: true @@ -57,49 +121,62 @@ backend: # Must be 1 when solver = "gurobi" (same N=1 restriction as lp_enabled). bab_max_batch_size: 8 + # Text-verification settings. When `method` is set, these overwrite the + # corresponding bab.* values (bab.method/p/perturbed_words/eps/max_eps/ + # num_verify_iters/k/alpha_opt_steps) and bab.solver_tier. + method: null # SST/Yelp method: planar, rule, alpha, ibp or discrete + p: 2.0 # text embedding perturbation norm metadata (2 or inf) + perturbed_words: 1 # token positions perturbed together (1 or 2) + eps: 0.00001 # initial verification radius (>= 0) + max_eps: 0.01 # maximum radius for the certified-radius search (>= 0) + num_verify_iters: 5 # binary-search iterations for the certified radius (>= 0) + k: 1 # rule threshold for the rule-slope attention alpha + alpha_opt_steps: 1000 # optimization steps for optimized-alpha refinement + # ───────────────────────────────────────────────────────────────────────── # TorchLP solver tuning (torchlp) # ───────────────────────────────────────────────────────────────────────── torchlp: - # -- Penalty weights -- - rho_eq: 10.0 # equality-constraint penalty weight - rho_ineq: 10.0 # inequality-constraint penalty weight - # -- Adam optimizer -- - max_iter: 2000 # Adam iteration cap + # -- Penalty weights (objective adds rho * sum of squared violations) -- + rho_eq: 10.0 # weight on sum((A_eq x - b_eq)^2) + rho_ineq: 10.0 # weight on sum(relu(A_le x - b_le)^2) + # -- Adam optimizer (torch.optim.Adam over the LP variables x) -- + max_iter: 2000 # max Adam iterations per solve lr: 0.01 # Adam learning rate - beta1: 0.9 # Adam beta1 - beta2: 0.999 # Adam beta2 - weight_decay: 0.0 # Adam weight decay - tol_feas: 0.0001 # feasibility tolerance for SAT reporting - # -- Large-N overrides (kick in above large_n_threshold variables) -- - large_n_threshold: 20000 # variable-count threshold for large-N mode - large_n_max_iter: 800 # large-N iteration cap - large_n_tol: 0.001 # large-N feasibility tolerance + beta1: 0.9 # Adam beta1 (first-moment decay) + beta2: 0.999 # Adam beta2 (second-moment decay) + weight_decay: 0.0 # Adam weight decay (L2 penalty on x) + tol_feas: 0.0001 # a lane is SAT if its max constraint violation <= this + # -- Large-N mode (when batch size x variables per LP > large_n_threshold) -- + large_n_threshold: 20000 # batch size x variables per LP that triggers large-N mode + large_n_max_iter: 800 # large-N: max_iter is capped at this + large_n_tol: 0.001 # large-N: tol_feas is raised to at least this # -- Stagnation early-stop -- - stagnation_patience: 300 # stop after this many stagnant feasibility checks - stagnation_tol: 0.00001 # min feasibility improvement to reset stagnation - feas_check_stride: 5 # iteration stride between feasibility checks + stagnation_patience: 300 # stop after this many iterations without improvement + stagnation_tol: 0.00001 # min drop in max violation that counts as improvement + feas_check_stride: 5 # iterations between max-violation checks # ───────────────────────────────────────────────────────────────────────── # Gurobi solver tuning (gurobi) # ───────────────────────────────────────────────────────────────────────── gurobi: - time_limit: null # per-solve wall-clock cap; null = caller timeout - mip_gap: 0.0001 # relative MIP optimality gap - threads: 0 # 0 = Gurobi automatic thread selection - output_flag: 0 # 0 = silent Gurobi output + time_limit: null # Gurobi TimeLimit in seconds; null = the caller's time limit + mip_gap: 0.0001 # Gurobi MIPGap (relative MIP optimality gap) + threads: 0 # Gurobi Threads; 0 = Gurobi chooses automatically + output_flag: 0 # Gurobi OutputFlag; 0 = silent, 1 = print the solver log # ───────────────────────────────────────────────────────────────────────── # HybridZ (hybridz) # ───────────────────────────────────────────────────────────────────────── hybridz: - # Per-call wall-clock budget in seconds; null = no HybridZ-specific limit. + # Wall-clock seconds per HybridZ spec evaluation; null or 0 = backend.timeout. timeout: null # MILP feasibility/spec tolerance for HybridZ verdict decisions. tolerance: 0.0000001 - # HybridZ generator/input dimension fallback limit (used by the HybridZ TF). + # A layer whose output width or generator count (continuous + binary) exceeds + # this keeps its interval bounds instead of a hybrid zonotope. max_input_dim: 1024 - # Segments per side of the sigmoid inflection point. + # Sigmoid only: linear segments on each side of the inflection point (>= 1). sigmoid_segments: 2 # Project eligible Sigmoid -> Dense chains without materializing Sigmoid HZ rows. fuse_sigmoid_affine: false @@ -108,10 +185,15 @@ backend: # Dual # ───────────────────────────────────────────────────────────────────────── dual: - n_iters: 50 - lr_alpha: 0.1 - per_class_alpha: true - incremental_start_enabled: true + n_iters: 50 # Adam iterations for alpha/eta optimization (dual tiers) + lr_alpha: 0.1 # Adam learning rate for alpha (relaxation slopes) + lr_beta: 0.1 # Adam learning rate for eta (split-constraint multipliers) + lr_decay: 0.98 # multiplicative learning-rate decay per Adam iteration + per_class_alpha: true # separate alpha tensors per output class + incremental_start_enabled: true # children start from the parent's alpha/eta + # Linear forward bounds are kept up to this many perturbed input dims; + # above it the forward pass is interval-only. + forward_lin_max_perturbed: 100 # ═════════════════════════════════════════════════════════════════════════ # Branch-and-Bound refinement (bab) @@ -119,10 +201,19 @@ backend: bab: # -- General BaB -------------------------------------------------------- + # Run BaB (tier 3) on lanes still UNKNOWN after the earlier tiers; skipped + # for solver hybridz. enabled: true + # Max tree depth. Unresolved subproblems at this depth are dropped instead + # of split, so the verdict can then no longer be CERTIFIED. max_depth: 3 + # BaB stops after bounding this many subproblems in total; any left + # unproven make the verdict UNKNOWN. max_nodes: 10 + # Max pending subproblems; above it the lowest-priority ones are evicted, + # which also rules out CERTIFIED. 0 = no cap. frontier_cap: 0 + # Pieces per input-axis split (values below 2 are raised to 2). input_split_fanout: 2 # -- Branching: which over-approximated activation (or input axis) to split # @@ -162,6 +253,10 @@ backend: # greedy and annealed are the Oliva order-leading strategies (ECOOP 2025, # arXiv:2507.17453); mcts is specified in docs/design/mcts_bab.md. bounding: "random" + # depth_bound_blend priority (higher pops first) = depth_weight * d + bound_weight * u, + # d = depth min-max normalized over the pool, u = (max lb - lb) / (max lb - min lb). + bounding_depth_weight: 0.5 # weight on the normalized depth d + bounding_bound_weight: 0.5 # weight on the bound urgency u # Cap on subproblems popped per wave, independent of the batch size. # 0 = unbounded. Honoured by the four TopKBounding-family strategies above; @@ -170,13 +265,30 @@ backend: # pooled, so a smaller top_k only re-sorts priorities more often. top_k: 0 - # Cooling schedule for bounding: "annealed". + # annealed: priority = u / T + Gumbel noise with T = max(sa_cooling_rate**wave, 1e-6), + # so selection becomes greedier as waves pass. sa_cooling_rate: 0.99 - mcts_exploration: 1.0 - mcts_lambda: 0.5 - mcts_virtual_loss: 1.0 - intermediate_refine: "none" - reuse_root_bounds: false + + # -- MCTS pool (bounding: mcts) ------------------------------------------- + mcts_exploration: 1.0 # UCB1 exploration constant (reserved: pop is top-k for now) + mcts_lambda: 0.5 # reward = lambda*depth/n_unstable + (1-lambda)*lb rank + mcts_virtual_loss: 1.0 # virtual loss for batched selection (reserved) + + # -- Intermediate-bound refinement and root-bound reuse (dual tiers) ----- + intermediate_refine: "none" # root pre-activation refinement: none/auto/all + intermediate_refine_ratio: 10.0 # auto mode: refine layers wider than ratio x median + root_bounds_reuse: "none" # children reuse root bounds: none/plain/split_refresh + eta_only_children: false # children keep the parent's alpha, optimize only eta + presplit_levels: 0 # pre-split k root neurons into 2^k sign combinations + per_subproblem_refine: "none" # per-subproblem refinement: none/tail/all (needs reuse) + per_subproblem_refine_iters: 0 # Adam iterations (0 = one fixed-slope backward pass) + per_subproblem_refine_rows_cap: 64 # max refined neurons per layer per batch (widest first) + + # Automatic BaB batch sizing. Used only when verify_bab_batched is called with + # max_batch_size="auto"; the CLI passes the integer bab_max_batch_size. + auto_batch_safety: 0.55 # fraction of total GPU memory the sizer may target + auto_batch_cap: 2048 # upper bound on the auto batch size (also the CPU fallback) + auto_batch_floor: 8 # lower bound on the auto batch size # Neurons split jointly per branching step; each lane fans out into all 2^k # sign combinations (verdict-boundary joint splitting). 1 = single split. @@ -192,18 +304,53 @@ backend: provenance_enabled: false # -- Dual-tier selector (solver_tier != lp) ---------------------------- + # How BaB bounds each wave: lp = the backend.solver LP; dual = dual bounds + # with un-optimized slopes; dual_alpha = Adam-optimized slopes (alpha); + # dual_alpha_eta = optimized alpha and split multipliers (eta). See dual:. solver_tier: "lp" + # -- CLIMB certificate reuse (ReLU neuron branching only) --------------- + climb_enabled: false # enable CLIMB (see bab_presets.climb above) + climb_theta: 0.0 # fraction of replay slack spent on dropping literals + # Deletion-budget safety margin: delta = climb_delta_abs + climb_delta_rel * |slack|. + climb_delta_abs: 0.000000001 # absolute part of delta + climb_delta_rel: 0.0000001 # relative part of delta (scaled by |slack|) + climb_recheck_k: 0 # lanes that get one extra replay after coarsening + climb_max_cores: 1024 # max cores kept (after subsumption, then eviction) + # Cores per propagate matmul chunk (memory only; fixpoint is chunk-independent). + climb_propagate_core_chunk: 4096 + # -- LLM-probe controller ---------------------------------------------- - llm_probe_enabled: false - llm_probe_backend: "mock" - llm_probe_cadence: 1 + llm_probe_enabled: false # enable the LLM-probe BaB controller + llm_probe_backend: "mock" # mock, openrouter, openai, glm, minimax or claude_cli + llm_probe_model: "" # model name; claude_cli defaults to "sonnet" if empty + llm_probe_base_url: "" # OpenAI-compatible URL; "" = the provider's preset URL + llm_probe_api_key_env: "" # env var holding the API key; "" = provider preset + llm_probe_temperature: 0.0 # sampling temperature (OpenAI-compatible backends) + llm_probe_timeout: 30.0 # per-call timeout (s); a slow wave falls back to baseline + llm_probe_max_candidates: 8 # max unstable-neuron candidates per domain + llm_probe_max_candidates_total: 1024 # with neuron_topk = 0: above this many, fall back to FSB + llm_probe_neuron_topk: 512 # >0: send only the top-K scored candidates, never FSB + llm_probe_cadence: 1 # consult the LLM every N waves + llm_probe_history: 8 # past wave outcomes included in the prompt + llm_probe_max_failures: 3 # consecutive failures before the controller is disabled + # Decisions the LLM may steer: split (joint split depth), frontier (wave + # width), refine (per-subproblem refinement), neuron (joint neuron-group + # selection), input_split (axis and fanout; input-splitting BaB only). llm_probe_decisions: "split,frontier,refine" - llm_probe_log: false + llm_probe_log: false # append per-wave JSONL to act/pipeline/log/llm_probe/ + + # Reserved: BaB does not read this yet (use backend.verbose for diagnostics). + verbose: false # -- BERT/text method --------------------------------------------------- - method: null - p: 2.0 - eps: 0.00001 - max_eps: 0.01 - k: 1 + method: null # SST/Yelp method: planar, rule, alpha, ibp or discrete + baf: true # BaF flag; true for planar/rule/alpha, false for ibp/discrete + alpha_mode: "fixed" # attention alpha: fixed (planar), rule, optimized (alpha), none + p: 2.0 # text embedding perturbation norm metadata (2 or inf) + perturbed_words: 1 # token positions perturbed together (1 or 2) + eps: 0.00001 # initial verification radius (>= 0) + max_eps: 0.01 # maximum radius for the certified-radius search (>= 0) + num_verify_iters: 5 # binary-search iterations for the certified radius (>= 0) + k: 1 # rule threshold for the rule-slope attention alpha + alpha_opt_steps: 1000 # optimization steps for optimized-alpha refinement diff --git a/act/config/backend_cli.py b/act/config/backend_cli.py index 8f11f4305..db608cb6c 100644 --- a/act/config/backend_cli.py +++ b/act/config/backend_cli.py @@ -8,12 +8,17 @@ - Network serialization (save/load ACT Net structures) - Analysis and constraint inspection +Exit status (--verify): + 0 every lane received a verdict (CERTIFIED / FALSIFIED / UNKNOWN / TIMEOUT) + 1 error, no verdict, or a network the selected solver cannot handle + 2 invalid configuration (ConfigError), reported as "❌ : " + Copyright (C) 2025 SVF-tools/ACT License: AGPLv3+ """ import argparse -from dataclasses import fields +from dataclasses import fields, replace import datetime import glob import json @@ -25,7 +30,7 @@ from pathlib import Path from typing import Any, Dict, List, NamedTuple, Optional, Union, cast, get_args, get_origin, get_type_hints -from act.config.config import DualConfig, GurobiConfig, TorchLPConfig, VALID_BERT_METHODS, VALID_BOUNDINGS, VALID_SOLVER_TIERS, _VALID_SOLVERS +from act.config.config import ConfigError, GurobiConfig, TorchLPConfig, VALID_BERT_METHODS, VALID_BOUNDINGS, VALID_ROOT_BOUNDS_REUSE, VALID_SOLVER_TIERS, _VALID_SOLVERS from act.back_end.layer_schema import LayerKind from act.front_end.specs import OutKind from act.util.cli_utils import add_device_args, initialize_from_args @@ -196,45 +201,53 @@ def _make_solver( return GurobiSolver(config=gurobi_config) if solver_name == "torchlp": return TorchLPSolver(config=torchlp_config) - # "auto": try Gurobi, fall back to TorchLP - try: - from act.back_end.solver.solver_gurobi import GurobiSolver + # "auto": Gurobi when gurobipy is importable, otherwise TorchLP + from act.back_end.solver.solver_gurobi import GurobiSolver, is_gurobi_available + if is_gurobi_available(): return GurobiSolver(config=gurobi_config) - except Exception: - return TorchLPSolver(config=torchlp_config) + return TorchLPSolver(config=torchlp_config) -_DUAL_BAB_PRESET: dict[str, Any] = { - "solver_tier": "dual_alpha_eta", - "branching_method": "gain", - "reuse_root_bounds": True, - "intermediate_refine": "all", - "multi_split_levels": 4, -} - -_DUAL_PRESET_ARG_DESTS: dict[str, str] = { - "solver_tier": "bab_solver_tier", - "branching_method": "bab_branching", - "reuse_root_bounds": "bab_reuse_root_bounds", - "intermediate_refine": "bab_intermediate_refine", - "multi_split_levels": "bab_multi_split_levels", -} - - -def explicit_dual_preset_fields(args: argparse.Namespace) -> set[str]: - """BaB preset fields the user pinned on the command line. +def explicit_bab_fields(args: argparse.Namespace) -> set[str]: + """BaB fields the user pinned on the command line. Every BaB override flag parses with ``default=None``, so a non-None value is - proof the user typed it. The dual preset must yield on exactly those fields. + proof the user typed it. A named preset must yield on exactly those fields. """ + from act.config.config import BaBConfig + + bab_fields = {fld.name for fld in fields(BaBConfig)} + destinations = { + key[4:]: attr + for key, attr, *_ in _BACKEND_OVERRIDE_SPEC + if key.startswith("bab_") and key[4:] in bab_fields + } return { - field - for field, dest in _DUAL_PRESET_ARG_DESTS.items() - if getattr(args, dest, None) is not None + field_name + for field_name, destination in destinations.items() + if getattr(args, destination, None) is not None } +def _resolve_bab_config( + backend_cfg: Any, + pinned_fields: set[str], + *, + is_dual: bool, +) -> Any: + """Apply the selected BaB preset below explicit CLI overrides.""" + preset_name = backend_cfg.bab_preset or ("dual" if is_dual else None) + if preset_name is None: + return backend_cfg.bab + preset_values = { + field_name: value + for field_name, value in backend_cfg.bab_preset_values(preset_name).items() + if field_name not in pinned_fields + } + return replace(backend_cfg.bab, **preset_values) if preset_values else backend_cfg.bab + + def _verify_one_net( net_path: str, backend_cfg, @@ -242,13 +255,19 @@ def _verify_one_net( ) -> tuple[list[Any], Optional[Union[_SkipUnsupported, str]], Optional[int]]: """[BATCHED-API] Verify *net_path* via 3-tier cascade. + Named BaB presets yield to fields pinned by CLI flags. + Returns ``(results, err, n_layers)`` where ``err`` is one of: * ``None`` on success * ``_SkipUnsupported(tf_name, kinds)`` when the active TF cannot handle the net (unsupported layer kinds and/or unsupported ASSERT spec). - Treated as a clean skip by callers, NOT as a verifier bug. + Reported distinctly from verifier bugs, but still a nonzero outcome + because no lane received a verdict. * ``str`` for any other exception (genuine error). + Raises ``ConfigError`` when the resolved configuration is invalid (BaB + preset application, CLIMB prerequisites) so the caller can exit with 2. + Tier 1 — interval (verify_once): always runs; certifies or falsifies via pure-tensor bounds propagation. Tier 2 — LP-batched (verify_lp_batched): runs on UNKNOWN lanes when @@ -258,7 +277,7 @@ def _verify_one_net( backend_cfg.bab_enabled is True AND active TF propagates LP constraints. bab_max_batch_size=1 disables K-batching. """ - from act.back_end.bab.bab import clear_violation_check_module_cache + from act.back_end.bab.violation import clear_violation_check_module_cache from act.back_end.serialization.serialization import load_net_from_file from act.back_end.transfer_functions import ( ensure_active_tf, @@ -360,22 +379,12 @@ def _verify_one_net( from act.back_end.bab.bab import verify_bab_batched as _vbb from act.back_end.verifier import slice_net_to_sample - bab_cfg = backend_cfg.bab - if is_dual: - # Convenience preset, NOT a mandate: --solver dual fills in - # optimized alpha/eta bounds with gain-tested joint multi-neuron - # (verdict-boundary) branching for users who pass no BaB flags, - # and yields on every field the user pinned explicitly. - import dataclasses - - pinned = explicit_bab_fields or set() - preset = { - field: value - for field, value in _DUAL_BAB_PRESET.items() - if field not in pinned - } - if preset: - bab_cfg = dataclasses.replace(bab_cfg, **preset) + pinned = explicit_bab_fields or set() + bab_cfg = _resolve_bab_config( + backend_cfg, + pinned, + is_dual=is_dual, + ) try: results = [ @@ -403,26 +412,45 @@ def _verify_one_net( raise return results, None, n_layers + except ConfigError: + raise except Exception as e: # noqa: BLE001 — surface per-net error, keep iterating return [], str(e), None +def _report_config_error(label: str, err: ConfigError) -> int: + """Print *err* in the CLI error format and return exit status 2.""" + print(f"❌ {label}: {err}") + return 2 + + def run_verification(args, backend_cfg): - """Run verification on a network using *backend_cfg*.""" + """Run verification on a network using *backend_cfg*. + + Returns 0 when every lane received a verdict, 1 on error / no verdict / + unsupported network, and 2 on invalid configuration. + """ from act.util.stats import VerifyStatus - results, err, n_layers = _verify_one_net( - args.network, backend_cfg, explicit_dual_preset_fields(args) - ) + pinned_bab_fields = explicit_bab_fields(args) + try: + results, err, n_layers = _verify_one_net( + args.network, backend_cfg, pinned_bab_fields + ) + except ConfigError as e: + return _report_config_error(args.network, e) if err is not None: if isinstance(err, _SkipUnsupported): print( f"⏭️ {args.network}: {err.tf_name} cannot handle: " f"{','.join(err.kinds)}" ) - return 0 + return 1 print(f"❌ {args.network}: {err}") return 1 + if not results: + print(f"❌ {args.network}: no verdict") + return 1 print(f"Loaded {n_layers}-layer net; solver={backend_cfg.solver}") valid_outcomes = ( @@ -930,6 +958,11 @@ def main(): # Use GPU with float64 python -m act.back_end --verify --network cifar.json --device cuda --dtype float64 + +Exit status (--verify): + 0 every lane received a verdict (CERTIFIED / FALSIFIED / UNKNOWN / TIMEOUT) + 1 error, no verdict, or a network the selected solver cannot handle + 2 invalid configuration (ConfigError) """, ) @@ -1173,6 +1206,21 @@ def main(): dest="bab_solver_tier", help="Solver tier for BaB bound computation (default: from config.yaml)", ) + verify_group.add_argument( + "--bab-root-bounds-reuse", + type=str, + default=None, + choices=VALID_ROOT_BOUNDS_REUSE, + dest="bab_root_bounds_reuse", + help=( + "Descendant reuse of the root box's forward bounds (dual tiers): " + "none = every node re-propagates; " + "plain = children get the root dict untouched; " + "split_refresh = root dict hardened by the split-derived interval " + "refresh and --bab-per-subproblem-refine " + "(default: from config.yaml)" + ), + ) verify_group.add_argument( "--bab-branching", type=str, @@ -1222,7 +1270,7 @@ def main(): verify_group.add_argument( "--bab-llm-probe-enabled", action="store_true", - default=False, + default=None, dest="bab_llm_probe_enabled", help="Enable the LLM-probe BaB controller (default: from config.yaml)", ) @@ -1323,7 +1371,7 @@ def main(): verify_group.add_argument( "--bab-llm-probe-log", action="store_true", - default=False, + default=None, dest="bab_llm_probe_log", help="Enable per-wave LLM-probe decision logging (default: from config.yaml)", ) @@ -1398,14 +1446,17 @@ def main(): return run_diff_nets(args) # ── Build BackendConfig ────────────────────────────────────────────── - # Load YAML as baseline, then overlay env vars and CLI flags on top. - # Precedence: CLI flag > env var > config.yaml > dataclass default + # Load YAML, then apply env vars and CLI flags. BaB presets are resolved at + # verification time so explicit BaB flags remain the highest precedence. from act.config.config import BackendConfig - backend_cfg = BackendConfig.from_yaml( - config_path=args.backend_config, - **_collect_backend_overrides(args, _user_set), - ) + try: + backend_cfg = BackendConfig.from_yaml( + config_path=args.backend_config, + **_collect_backend_overrides(args, _user_set), + ) + except ConfigError as e: + return _report_config_error(args.network or "Error", e) # Initialize device manager from the resolved config import argparse as _ap @@ -1458,7 +1509,7 @@ def main(): return 130 except Exception as e: print(f"\n❌ Error: {e}") - if args.verbose: + if backend_cfg.verbose: import traceback traceback.print_exc() @@ -1471,6 +1522,7 @@ def main(): ("solver", "solver", "ACT_SOLVER", None, "not_none"), ("device", "device", "ACT_DEVICE", None, "user_set"), ("dtype", "dtype", "ACT_DTYPE", None, "user_set"), + ("verbose", "verbose", None, None, "user_set"), ("timeout", "timeout", None, None, "not_none"), ("method", "method", None, None, "not_none"), ("p", "p", None, None, "not_none"), @@ -1525,7 +1577,10 @@ def _build_backend_override_spec() -> list[tuple[str, str, Optional[str], Any, s def _collect_backend_overrides(args: Any, _user_set: Any) -> dict[str, Any]: - """Build overrides dict from CLI flags + env vars (precedence: CLI > env > yaml).""" + """Build overrides dict from CLI flags + env vars. + + Precedence before BaB preset resolution: CLI > env > backend.yaml. + """ overrides: dict[str, Any] = {} for key, attr, env, env_cast, check in _BACKEND_OVERRIDE_SPEC: cli_val = getattr(args, attr, None) @@ -1539,9 +1594,6 @@ def _collect_backend_overrides(args: Any, _user_set: Any) -> dict[str, Any]: elif env is not None and os.environ.get(env): overrides[key] = env_cast(os.environ[env]) if env_cast else os.environ[env] - if args.verbose: - overrides["verbose"] = True - return overrides diff --git a/act/config/check_parity.py b/act/config/check_parity.py index f00eaa787..87bd8af03 100644 --- a/act/config/check_parity.py +++ b/act/config/check_parity.py @@ -4,14 +4,12 @@ (``python -m act.config.check_parity``). Exits non-zero on any drift, so an option added to one surface can never be silently missed by the others. -The config.py dataclasses are the single source of truth. Two documented -asymmetries are allowed: +The config.py dataclasses are the single source of truth. One documented +asymmetry is allowed: * pipeline ``verification.bab`` is a SPARSE override of ``backend.yaml`` (its YAML lists only non-default keys), so there the YAML keys must be a subset of the CLI surface rather than equal to it. -* fields declared ``metadata={"in_yaml": False}`` (BackendConfig's text-verify - scalars mirrored into ``bab.*``) are CLI/dataclass-only and absent from YAML. """ from __future__ import annotations @@ -52,7 +50,7 @@ def require(self, label: str, passed: bool, detail: str = "") -> None: def check_backend(report: ParityReport) -> None: - from act.config.config import _BACKEND_YAML, _NETGEN_YAML, BaBConfig, BackendConfig, DualConfig + from act.config.config import _BACKEND_YAML, _NETGEN_YAML, BaBConfig from act.config.backend_cli import ( _BACKEND_OVERRIDE_SPEC, _BACKEND_SUBCONFIG_PREFIX, @@ -65,7 +63,8 @@ def check_backend(report: ParityReport) -> None: cli_options = {key for key, *_ in _BACKEND_OVERRIDE_SPEC} yaml_keys: set[str] = set() - backend_yaml = _load_yaml(_BACKEND_YAML).get("backend", {}) + backend_document = _load_yaml(_BACKEND_YAML) + backend_yaml = backend_document.get("backend", {}) for key, value in backend_yaml.items(): prefix = _BACKEND_SUBCONFIG_PREFIX.get(key) if prefix is None: @@ -82,25 +81,6 @@ def check_backend(report: ParityReport) -> None: for key in _load_yaml(_NETGEN_YAML): yaml_keys.add(f"gen_{key}") - # A backend option may be CLI-only (settable but absent from the YAML) only if - # its dataclass field is declared metadata={"in_yaml": False}. No hand-coded - # list here -- the field declares it, and a normal new field defaults to - # requiring a YAML entry. - cli_only_allowed = { - f.name for f in fields(BackendConfig) if not f.metadata.get("in_yaml", True) - } - bab_cli_only_allowed = { - f"bab_{f.name}" - for f in fields(BaBConfig) - if not f.metadata.get("in_yaml", True) - } - dual_cli_only_allowed = { - f"dual_{f.name}" - for f in fields(DualConfig) - if not f.metadata.get("in_yaml", True) - } - cli_only_allowed |= bab_cli_only_allowed | dual_cli_only_allowed - print("[backend]") report.require("CLI options are backed by a dataclass field", cli_options <= dataclass_fields, _fmt(cli_options - dataclass_fields)) @@ -108,11 +88,32 @@ def check_backend(report: ParityReport) -> None: dataclass_fields <= cli_options, _fmt(dataclass_fields - cli_options)) report.require("YAML keys are backed by a dataclass field", yaml_keys <= dataclass_fields, _fmt(yaml_keys - dataclass_fields)) - report.require("YAML keys are all CLI-settable", - yaml_keys <= cli_options, _fmt(yaml_keys - cli_options)) - report.require("CLI options absent from YAML are in_yaml=False fields", - (cli_options - yaml_keys) <= cli_only_allowed, - _fmt((cli_options - yaml_keys) - cli_only_allowed)) + report.require("CLI and YAML are 1-to-1", + cli_options == yaml_keys, + f"cli-only={_fmt(cli_options - yaml_keys)} yaml-only={_fmt(yaml_keys - cli_options)}") + + print("[backend.bab_presets]") + bab_yaml_keys = set((backend_yaml.get("bab") or {}).keys()) - {"enabled"} + for preset_name, preset in (backend_document.get("bab_presets") or {}).items(): + preset_keys = set(preset or {}) + report.require( + f"preset '{preset_name}' only uses backend.bab keys", + preset_keys <= bab_yaml_keys, + _fmt(preset_keys - bab_yaml_keys), + ) + try: + BaBConfig(**(preset or {})) + except (TypeError, ValueError) as exc: + report.require( + f"preset '{preset_name}' constructs a valid BaBConfig", + False, + str(exc), + ) + else: + report.require( + f"preset '{preset_name}' constructs a valid BaBConfig", + True, + ) def check_pipeline(report: ParityReport) -> None: diff --git a/act/config/config.py b/act/config/config.py index e56e6deea..9cbf8364b 100644 --- a/act/config/config.py +++ b/act/config/config.py @@ -1,8 +1,9 @@ from __future__ import annotations from copy import deepcopy -from dataclasses import asdict, dataclass, field, fields +from dataclasses import dataclass, field, fields from importlib import import_module +import math from pathlib import Path from typing import Any, Final, List, Optional, Union @@ -13,12 +14,28 @@ _PIPELINE_YAML = Path(__file__).parent / "pipeline.yaml" _FRONTEND_YAML = Path(__file__).parent / "frontend.yaml" + +class ConfigError(ValueError): + """Invalid configuration; CLIs exit with status 2.""" + + _VALID_SOLVERS = {"auto", "gurobi", "torchlp", "dual", "hybridz"} _VALID_DEVICES = {"cpu", "cuda", "gpu"} _VALID_DTYPES = {"float32", "float64"} _VALID_REGISTRY_MODES = {"intersection", "union"} _VALID_COVERAGE_MODES = {"basic", "full"} VALID_SOLVER_TIERS: Final[tuple[str, ...]] = ("lp", "dual", "dual_alpha", "dual_alpha_eta") +CLIMB_SOLVER_TIER: Final[str] = VALID_SOLVER_TIERS[-1] +NO_REFINEMENT_MODE: Final[str] = "none" +# Descendant reuse of the root box's forward bounds (``bab.root_bounds_reuse``). +# +# value | root dict passed to children | split-derived refresh +# ---------------|------------------------------|---------------------------- +# none | no (children re-propagate) | n/a +# plain | yes, untouched | no +# split_refresh | yes | _interval_refresh_bounds + +# | | per_subproblem_refine +VALID_ROOT_BOUNDS_REUSE: Final[tuple[str, ...]] = ("none", "plain", "split_refresh") # BaB subproblem-pool selection strategies (``--bab-bounding``). # # value | pool class | order function | --bab-top-k | reference @@ -42,6 +59,12 @@ "mcts", ) TOP_K_BOUNDINGS: Final[tuple[str, ...]] = VALID_BOUNDINGS[:4] +NEURON_BRANCHING_METHODS: Final[tuple[str, ...]] = ( + "babsr", + "fsb", + "gain", + "witness_residual", +) TOP_K_INCOMPATIBLE_BOUNDINGS: Final[tuple[str, ...]] = VALID_BOUNDINGS[4:] VALID_BERT_METHODS: Final[tuple[str, ...]] = ( "planar", @@ -87,7 +110,7 @@ def normalize_bert_method(method: str) -> str: key = method.strip().lower().replace("-", "_") if key not in _BERT_METHOD_SELECTIONS: valid = ", ".join(name.replace("_", "-") for name in VALID_BERT_METHODS) - raise ValueError(f"Invalid bert method {method!r}; expected one of: {valid}") + raise ConfigError(f"Invalid bert method {method!r}; expected one of: {valid}") return key @@ -127,8 +150,8 @@ class BaBConfig: branching_method: str = "random" bounding: str = "random" - bounding_depth_weight: float = field(default=0.5, metadata={"in_yaml": False}) - bounding_bound_weight: float = field(default=0.5, metadata={"in_yaml": False}) + bounding_depth_weight: float = 0.5 + bounding_bound_weight: float = 0.5 sa_cooling_rate: float = 0.99 mcts_exploration: float = 1.0 mcts_lambda: float = 0.5 @@ -137,69 +160,93 @@ class BaBConfig: # Dual-tier solver knobs — support solver_tier="dual_alpha_eta" with # Iterative slope + Lagrange-multiplier optimization for the dual backward pass. solver_tier: str = "lp" - f"""Solver tier for BaB bound computation. Valid: {VALID_SOLVER_TIERS}.""" + """Solver tier for BaB bound computation. + Valid: 'lp', 'dual', 'dual_alpha', 'dual_alpha_eta'.""" + + climb_enabled: bool = False + """Enable query-local CLIMB certificate replay and core propagation.""" + climb_theta: float = 0.0 + """Fraction of replay slack available to the vector deletion budget.""" + climb_delta_abs: float = 1e-9 + """Absolute deletion-budget margin: delta = climb_delta_abs + climb_delta_rel * |slack|.""" + climb_delta_rel: float = 1e-7 + """Relative deletion-budget margin: delta = climb_delta_abs + climb_delta_rel * |slack|.""" + climb_recheck_k: int = 0 + """Maximum lanes given one extra direct replay after budget coarsening.""" + climb_max_cores: int = 1024 + """Maximum retained cores after subsumption and deterministic eviction.""" + climb_propagate_core_chunk: int = 4096 + """Cores per propagate matmul chunk (memory only; the fixpoint is chunk-independent).""" provenance_enabled: bool = False """Track logical BaB node ids and parent ids in TopKBounding.""" - eta_only_children: bool = field(default=False, metadata={"in_yaml": False}) + eta_only_children: bool = False """Freeze alpha in child subproblems (depth > 0): children inherit the parent's optimized alpha and refine only the split multipliers (eta). - Cuts the per-node Adam graph and, combined with reuse_root_bounds, + Cuts the per-node Adam graph and, combined with root_bounds_reuse, removes the per-iteration forward pass entirely.""" - presplit_levels: int = field(default=0, metadata={"in_yaml": False}) + presplit_levels: int = 0 """Pre-split the root's top-k scored unstable neurons into all 2^k sign combinations before the main loop (LEAPS-style leap: descendants are materialized directly, intermediate tree levels are never bounded). The combinations exactly partition the root region, so soundness is unaffected. Requires a dual tier with neuron branching state.""" - intermediate_refine: str = "none" + intermediate_refine: str = NO_REFINEMENT_MODE """Backward refinement of intermediate pre-activation bounds at the root: 'none' (off), 'auto' (refine activation layers whose mean width exceeds intermediate_refine_ratio x the median - targets wide fan-in concretization loss), 'all' (every unstable activation layer).""" - intermediate_refine_ratio: float = field(default=10.0, metadata={"in_yaml": False}) + intermediate_refine_ratio: float = 10.0 """Width-blowup threshold multiplier for intermediate_refine='auto'.""" - reuse_root_bounds: bool = False - """Reuse the root box's forward bounds for every descendant (dual tiers). + root_bounds_reuse: str = "none" + """Reuse of the root box's forward bounds by descendants (dual tiers). + Valid: 'none', 'plain', 'split_refresh'. Sound by monotonicity: a child box is contained in the root box, so the root's per-layer bounds remain valid over-approximations. Children only override the INPUT/INPUT_SPEC bounds with their own sub-box; intermediate ReLU relaxations stay at root tightness, with branching gains recovered by the input-term concretization and the eta split multipliers. Eliminates - the per-node forward pass (the dominant time and memory cost).""" + the per-node forward pass (the dominant time and memory cost). - per_subproblem_refine: str = field(default="none", metadata={"in_yaml": False}) + 'none' re-propagates per node. Both reuse modes override the INPUT and + INPUT_SPEC entries of the reused dict with each lane's own sub-box; + 'plain' stops there, while 'split_refresh' additionally hardens the dict + with the split-derived interval refresh plus per_subproblem_refine. So + 'plain' is the split-independent reference: it skips only the + split-derived tightening, not the per-lane input override.""" + + per_subproblem_refine: str = NO_REFINEMENT_MODE """Per-subproblem sparse backward refinement of intermediate bounds in the - BaB loop (requires reuse_root_bounds): 'none' (off), 'tail' (last two + BaB loop (requires root_bounds_reuse != 'none'): 'none' (off), 'tail' (last two unstable activation layers), 'all' (every unstable activation layer). For each child batch, the split-hardened bounds are re-tightened by a K-lane backward pass over the unstable-neuron union only (stable phases are exact, so refining them gains nothing), so splits propagate relationally downstream instead of only through the interval refresh.""" - per_subproblem_refine_iters: int = field(default=0, metadata={"in_yaml": False}) + per_subproblem_refine_iters: int = 0 """Adam iterations for per-subproblem refine rows (0 = single fixed-slope backward, cheapest).""" - per_subproblem_refine_rows_cap: int = field(default=64, metadata={"in_yaml": False}) + per_subproblem_refine_rows_cap: int = 64 """Max refined neurons per layer per batch (top-cap by interval width); bounds the K x 2*cap backward cost.""" - auto_batch_safety: float = field(default=0.55, metadata={"in_yaml": False}) + auto_batch_safety: float = 0.55 """Fraction of GPU memory the auto batch sizer (max_batch_size='auto') may target; lowered on a shared GPU. The sizer also never exceeds 90% of the currently-reclaimable memory (free + this process's reserved cache).""" - auto_batch_cap: int = field(default=2048, metadata={"in_yaml": False}) + auto_batch_cap: int = 2048 """Hard upper bound on the auto-sized batch (also the CPU fallback).""" - auto_batch_floor: int = field(default=8, metadata={"in_yaml": False}) + auto_batch_floor: int = 8 """Lower bound on the auto-sized batch.""" multi_split_levels: int = 1 @@ -212,17 +259,17 @@ class BaBConfig: llm_probe_enabled: bool = False llm_probe_backend: str = "mock" - llm_probe_model: str = field(default="", metadata={"in_yaml": False}) - llm_probe_base_url: str = field(default="", metadata={"in_yaml": False}) - llm_probe_api_key_env: str = field(default="", metadata={"in_yaml": False}) - llm_probe_temperature: float = field(default=0.0, metadata={"in_yaml": False}) - llm_probe_timeout: float = field(default=30.0, metadata={"in_yaml": False}) - llm_probe_max_candidates: int = field(default=8, metadata={"in_yaml": False}) - llm_probe_max_candidates_total: int = field(default=1024, metadata={"in_yaml": False}) - llm_probe_neuron_topk: int = field(default=512, metadata={"in_yaml": False}) + llm_probe_model: str = "" + llm_probe_base_url: str = "" + llm_probe_api_key_env: str = "" + llm_probe_temperature: float = 0.0 + llm_probe_timeout: float = 30.0 + llm_probe_max_candidates: int = 8 + llm_probe_max_candidates_total: int = 1024 + llm_probe_neuron_topk: int = 512 llm_probe_cadence: int = 1 - llm_probe_history: int = field(default=8, metadata={"in_yaml": False}) - llm_probe_max_failures: int = field(default=3, metadata={"in_yaml": False}) + llm_probe_history: int = 8 + llm_probe_max_failures: int = 3 llm_probe_decisions: str = "split,frontier,refine" """Comma-separated decision types the LLM may steer: 'split' (joint neuron split depth), 'frontier' (wave width), 'refine' (per-subproblem refinement), @@ -230,48 +277,63 @@ class BaBConfig: dimension to bisect and its fanout, input-domain-splitting BaB only).""" llm_probe_log: bool = False - verbose: bool = field(default=False, metadata={"in_yaml": False}) + verbose: bool = False method: Optional[str] = None - baf: bool = field(default=True, metadata={"in_yaml": False}) - alpha_mode: str = field(default="fixed", metadata={"in_yaml": False}) + baf: bool = True + alpha_mode: str = "fixed" p: float = 2.0 - perturbed_words: int = field(default=1, metadata={"in_yaml": False}) + perturbed_words: int = 1 eps: float = 1e-5 max_eps: float = 0.01 - num_verify_iters: int = field(default=5, metadata={"in_yaml": False}) + num_verify_iters: int = 5 k: int = 1 - alpha_opt_steps: int = field(default=1000, metadata={"in_yaml": False}) + alpha_opt_steps: int = 1000 def __post_init__(self) -> None: if self.solver_tier not in VALID_SOLVER_TIERS: - raise ValueError( + raise ConfigError( f"Invalid solver_tier {self.solver_tier!r}; expected {VALID_SOLVER_TIERS}" ) if self.bounding not in VALID_BOUNDINGS: - raise ValueError( + raise ConfigError( f"Invalid bounding {self.bounding!r}; expected {VALID_BOUNDINGS}" ) + if self.root_bounds_reuse not in VALID_ROOT_BOUNDS_REUSE: + raise ConfigError( + f"Invalid root_bounds_reuse {self.root_bounds_reuse!r}; " + f"expected {VALID_ROOT_BOUNDS_REUSE}" + ) if self.top_k < 0: - raise ValueError(f"top_k must be non-negative, got {self.top_k}") + raise ConfigError(f"top_k must be non-negative, got {self.top_k}") + if not math.isfinite(self.climb_theta) or not 0.0 <= self.climb_theta <= 1.0: + raise ConfigError("climb_theta must be finite and in [0, 1]") + if not math.isfinite(self.climb_delta_abs) or self.climb_delta_abs < 0.0: + raise ConfigError("climb_delta_abs must be finite and non-negative") + if not math.isfinite(self.climb_delta_rel) or self.climb_delta_rel < 0.0: + raise ConfigError("climb_delta_rel must be finite and non-negative") + if self.climb_recheck_k < 0: + raise ConfigError("climb_recheck_k must be non-negative") + if self.climb_max_cores < 1: + raise ConfigError("climb_max_cores must be positive") + if self.climb_propagate_core_chunk < 1: + raise ConfigError("climb_propagate_core_chunk must be positive") if self.top_k > 0 and self.bounding in TOP_K_INCOMPATIBLE_BOUNDINGS: - raise ValueError( + raise ConfigError( f"top_k={self.top_k} is not supported by bounding={self.bounding!r}; " f"it applies only to the order-ranked pools {TOP_K_BOUNDINGS}" ) if self.method is not None: selection = select_bert_method(self.method) self.method = selection.method - self.baf = selection.baf - self.alpha_mode = selection.alpha_mode if self.solver_tier == "lp": self.solver_tier = selection.solver_tier if self.perturbed_words not in (1, 2): - raise ValueError("perturbed_words must be 1 or 2") + raise ConfigError("perturbed_words must be 1 or 2") if self.num_verify_iters < 0: - raise ValueError("num_verify_iters must be non-negative") + raise ConfigError("num_verify_iters must be non-negative") if self.max_eps < 0 or self.eps < 0: - raise ValueError("eps and max_eps must be non-negative") + raise ConfigError("eps and max_eps must be non-negative") @classmethod def from_yaml( @@ -304,17 +366,64 @@ def from_yaml( return cls(**merged) - def to_yaml(self, path: Union[str, Path]) -> Path: - """Write BaB settings to a standalone YAML file (top-level ``bab`` key).""" - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - with open(path, "w") as f: - yaml.dump( - {"bab": asdict(self)}, f, default_flow_style=False, sort_keys=False +def _validate_bab_presets( + raw_presets: Any, + config_path: Path, +) -> dict[str, dict[str, Any]]: + """Validate and copy named sparse ``BaBConfig`` presets.""" + if raw_presets is None: + return {} + if not isinstance(raw_presets, dict): + raise ConfigError(f"bab_presets in {config_path} must be a YAML mapping") + + valid_keys = {fld.name for fld in fields(BaBConfig)} + presets: dict[str, dict[str, Any]] = {} + for name, values in raw_presets.items(): + if not isinstance(name, str) or not isinstance(values, dict): + raise ConfigError( + f"BaB preset {name!r} in {config_path} must be a YAML mapping" + ) + unknown = set(values) - valid_keys + if unknown: + names = ", ".join(sorted(unknown)) + raise ConfigError( + f"Unknown BaB preset key(s) in {config_path} preset {name!r}: {names}" ) + copied = deepcopy(values) + BaBConfig(**copied) + presets[name] = copied + return presets - return path + +def load_bab_presets( + config_path: Optional[Union[str, Path]] = None, +) -> dict[str, dict[str, Any]]: + """Load and validate named BaB presets from a backend YAML file.""" + path = Path(config_path) if config_path else _BACKEND_YAML + if not path.exists(): + raise FileNotFoundError(f"Backend config not found: {path}") + return _validate_bab_presets(_load_yaml(path).get("bab_presets"), path) + + +def _select_bab_preset( + name: str, + presets: dict[str, dict[str, Any]], +) -> dict[str, Any]: + if name not in presets: + available = ", ".join(sorted(presets)) or "(none)" + raise ConfigError( + f"Unknown BaB preset {name!r}; available presets: {available}" + ) + return deepcopy(presets[name]) + + +def load_bab_preset( + name: str, + config_path: Optional[Union[str, Path]] = None, +) -> dict[str, Any]: + """Load one named BaB preset, listing available names on failure.""" + return _select_bab_preset(name, load_bab_presets(config_path)) # --------------------------------------------------------------------------- @@ -344,12 +453,12 @@ class GenerationConfig: def __post_init__(self) -> None: if self.registry_mode not in _VALID_REGISTRY_MODES: - raise ValueError( + raise ConfigError( f"Invalid registry_mode {self.registry_mode!r}; " f"expected one of {_VALID_REGISTRY_MODES}" ) if self.coverage_mode not in _VALID_COVERAGE_MODES: - raise ValueError( + raise ConfigError( f"Invalid coverage_mode {self.coverage_mode!r}; " f"expected one of {_VALID_COVERAGE_MODES}" ) @@ -379,10 +488,10 @@ class DualConfig: lr_alpha: float = 0.1 """Adam learning rate for α (slope) variables.""" - lr_beta: float = field(default=0.1, metadata={"in_yaml": False}) + lr_beta: float = 0.1 """Adam learning rate for η (split-constraint KKT multipliers).""" - lr_decay: float = field(default=0.98, metadata={"in_yaml": False}) + lr_decay: float = 0.98 """Multiplicative learning-rate decay applied each Adam iteration.""" per_class_alpha: bool = True @@ -391,6 +500,17 @@ class DualConfig: incremental_start_enabled: bool = True """Reuse α/η tensors from the parent subproblem as the child initialization.""" + forward_lin_max_perturbed: int = 100 + """Cap on the number of perturbed input dims for which the forward pass + keeps an explicit linear bound; above it, interval-only.""" + + def __post_init__(self) -> None: + if self.forward_lin_max_perturbed < 0: + raise ConfigError( + "forward_lin_max_perturbed must be non-negative, got " + f"{self.forward_lin_max_perturbed}" + ) + @dataclass class TorchLPConfig: @@ -436,6 +556,8 @@ class BackendConfig: verbose: bool = False timeout: float = 300.0 + bab_preset: Optional[str] = None + bab_enabled: bool = False bab: BaBConfig = field(default_factory=BaBConfig) @@ -462,36 +584,34 @@ class BackendConfig: torchlp: TorchLPConfig = field(default_factory=TorchLPConfig) dual: DualConfig = field(default_factory=DualConfig) - method: Optional[str] = field(default=None, metadata={"in_yaml": False}) - p: float = field(default=2.0, metadata={"in_yaml": False}) - perturbed_words: int = field(default=1, metadata={"in_yaml": False}) - eps: float = field(default=1e-5, metadata={"in_yaml": False}) - max_eps: float = field(default=0.01, metadata={"in_yaml": False}) - num_verify_iters: int = field(default=5, metadata={"in_yaml": False}) - k: int = field(default=1, metadata={"in_yaml": False}) - alpha_opt_steps: int = field(default=1000, metadata={"in_yaml": False}) + method: Optional[str] = None + p: float = 2.0 + perturbed_words: int = 1 + eps: float = 1e-5 + max_eps: float = 0.01 + num_verify_iters: int = 5 + k: int = 1 + alpha_opt_steps: int = 1000 # -- validation --------------------------------------------------------- def __post_init__(self) -> None: if self.solver not in _VALID_SOLVERS: - raise ValueError( + raise ConfigError( f"Invalid solver {self.solver!r}; expected one of {_VALID_SOLVERS}" ) if self.device not in _VALID_DEVICES: - raise ValueError( + raise ConfigError( f"Invalid device {self.device!r}; expected one of {_VALID_DEVICES}" ) if self.dtype not in _VALID_DTYPES: - raise ValueError( + raise ConfigError( f"Invalid dtype {self.dtype!r}; expected one of {_VALID_DTYPES}" ) if self.method is not None: selection = select_bert_method(self.method) self.method = selection.method self.bab.method = selection.method - self.bab.baf = selection.baf - self.bab.alpha_mode = selection.alpha_mode self.bab.solver_tier = selection.solver_tier self.bab.p = float(self.p) self.bab.perturbed_words = int(self.perturbed_words) @@ -504,7 +624,7 @@ def __post_init__(self) -> None: # Fail loud at config-load time rather than at the first batched call. if self.solver == "gurobi": if self.lp_enabled: - raise ValueError( + raise ConfigError( "BackendConfig: solver='gurobi' is incompatible with " "lp_enabled=True. GurobiSolver.solve_batch raises for N>1 " "(Gurobi does not expose a truly parallel multi-LP API for " @@ -512,7 +632,7 @@ def __post_init__(self) -> None: "Either set lp_enabled=False or switch to solver='torchlp'." ) if self.bab_max_batch_size > 1: - raise ValueError( + raise ConfigError( f"BackendConfig: solver='gurobi' is incompatible with " f"bab_max_batch_size={self.bab_max_batch_size} > 1. " f"GurobiSolver.solve_batch raises for N>1. " @@ -527,7 +647,7 @@ def from_yaml( config_path: Optional[Union[str, Path]] = None, **overrides, ) -> BackendConfig: - """Load config from YAML with optional keyword overrides. + """Load backend YAML and apply keyword overrides. YAML layout:: @@ -553,6 +673,7 @@ def from_yaml( raise FileNotFoundError(f"Backend config not found: {path}") raw = _load_yaml(path) + bab_presets = load_bab_presets() backend_raw: dict[str, Any] = raw.get("backend", {}) bab_raw: dict[str, Any] = backend_raw.pop("bab", {}) @@ -596,12 +717,7 @@ def from_yaml( top_overrides[k] = v # Build BaBConfig - bab_in_yaml = { - fld.name for fld in fields(BaBConfig) if fld.metadata.get("in_yaml", True) - } - bab_merged = { - k: v for k, v in bab_raw.items() if k in bab_fields and k in bab_in_yaml - } + bab_merged = {k: v for k, v in bab_raw.items() if k in bab_fields} bab_merged.update(bab_overrides) bab_config = BaBConfig(**bab_merged) @@ -622,12 +738,7 @@ def from_yaml( torchlp_merged.update(torchlp_overrides) torchlp_config = TorchLPConfig(**torchlp_merged) - dual_in_yaml = { - fld.name for fld in fields(DualConfig) if fld.metadata.get("in_yaml", True) - } - dual_merged = { - k: v for k, v in dual_raw.items() if k in dual_fields and k in dual_in_yaml - } + dual_merged = {k: v for k, v in dual_raw.items() if k in dual_fields} dual_merged.update(dual_overrides) dual_config = DualConfig(**dual_merged) @@ -650,7 +761,7 @@ def from_yaml( top_merged.update({k: v for k, v in top_overrides.items() if k in top_fields}) - return cls( + config = cls( bab=bab_config, generation=gen_config, hybridz=hz_config, @@ -659,38 +770,13 @@ def from_yaml( dual=dual_config, **top_merged, ) + if config.bab_preset is not None: + _select_bab_preset(config.bab_preset, bab_presets) + return config - def to_yaml(self, path: Union[str, Path]) -> Path: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - - d = asdict(self) - bab_d = d.pop("bab") - d.pop("generation") - hz_d = d.pop("hybridz") - gurobi_d = d.pop("gurobi") - torchlp_d = d.pop("torchlp") - dual_d = d.pop("dual") - bab_enabled = d.pop("bab_enabled") - bab_d["enabled"] = bab_enabled - - with open(path, "w") as f: - yaml.dump( - { - "backend": { - **d, - "bab": bab_d, - "hybridz": hz_d, - "gurobi": gurobi_d, - "torchlp": torchlp_d, - "dual": dual_d, - } - }, - f, - default_flow_style=False, - sort_keys=False, - ) - return path + def bab_preset_values(self, name: str) -> dict[str, Any]: + """Return a validated named preset from the canonical backend YAML.""" + return load_bab_preset(name) if __name__ == "__main__": @@ -772,7 +858,7 @@ def build_vnncomp_bab_config( frontier_cap=25000, max_depth=max_depth, max_nodes=max_nodes, - reuse_root_bounds=True, + root_bounds_reuse="split_refresh", intermediate_refine="all", presplit_levels=0, eta_only_children=False, diff --git a/act/config/frontend_cli.py b/act/config/frontend_cli.py index c22ba98c0..e51cc95ee 100644 --- a/act/config/frontend_cli.py +++ b/act/config/frontend_cli.py @@ -221,6 +221,7 @@ def print_unified_info(name: str, explicit_creator: Optional[str] = None): except ValueError as e: print(f"Error: {e}") print(f"\nTip: Use --search to find available names or --creator to override detection") + raise def handle_unified_download(name: str, explicit_creator: Optional[str] = None): @@ -243,23 +244,16 @@ def handle_unified_download(name: str, explicit_creator: Optional[str] = None): info = tv_mapping.get_dataset_info(normalized_name) models = info.get('models', []) - if not models: - print(f"⚠️ No models available for {normalized_name}") - return + assert models, f"No models available for {normalized_name}" print(f"Downloading dataset '{normalized_name}' with {len(models)} models...\n") success_count = 0 for model in models: - try: - result = tv_loader.download_dataset_model_pair(normalized_name, model) - if result['status'] == 'success': - print(f"✓ {normalized_name} + {model}") - success_count += 1 - else: - print(f"✗ {normalized_name} + {model} - {result['message']}") - except Exception as e: - print(f"✗ {normalized_name} + {model} - Error: {e}") + result = tv_loader.download_dataset_model_pair(normalized_name, model) + assert result['status'] == 'success', f"{normalized_name} + {model}: {result['message']}" + print(f"✓ {normalized_name} + {model}") + success_count += 1 print(f"\n{rule(100)}") print(f"Downloaded {success_count}/{len(models)} model pairs") @@ -275,36 +269,29 @@ def handle_unified_download(name: str, explicit_creator: Optional[str] = None): print(f" • VNNLIB property files") print(f" • instances.csv mapping\n") - try: - result = vnnlib_loader.download_vnnlib_category(normalized_name) - - if result['status'] == 'success': - print(f"\n{rule(100)}") - print(f"✓ Successfully downloaded category: {normalized_name}") - print(f" Location: {result['category_path']}") - print(f" Instances: {result['num_instances']}") - print(f"{rule(100)}\n") - else: - print(f"\n{rule(100)}") - print(f"✗ Download failed: {result['message']}") - print(f"\nNote: VNNLIB benchmarks must be downloaded manually from VNN-COMP.") - print(f"Expected location: data/vnnlib/{normalized_name}/") - print(f"\nManual download steps:") - print(f" 1. Visit: https://github.com/ChristopherBrix/vnncomp_benchmarks") - print(f" 2. Download the '{normalized_name}' benchmark") - print(f" 3. Extract to: data/vnnlib/{normalized_name}/") - print(f" 4. Ensure structure includes:") - print(f" - onnx/ (ONNX model files)") - print(f" - vnnlib/ (VNNLIB property files)") - print(f" - instances.csv (benchmark instances)") - print(f"{rule(100)}\n") - - except Exception as e: + result = vnnlib_loader.download_vnnlib_category(normalized_name) + + if result['status'] != 'success': print(f"\n{rule(100)}") - print(f"✗ Download error: {e}") + print(f"✗ Download failed: {result['message']}") print(f"\nNote: VNNLIB benchmarks must be downloaded manually from VNN-COMP.") print(f"Expected location: data/vnnlib/{normalized_name}/") + print(f"\nManual download steps:") + print(f" 1. Visit: https://github.com/ChristopherBrix/vnncomp_benchmarks") + print(f" 2. Download the '{normalized_name}' benchmark") + print(f" 3. Extract to: data/vnnlib/{normalized_name}/") + print(f" 4. Ensure structure includes:") + print(f" - onnx/ (ONNX model files)") + print(f" - vnnlib/ (VNNLIB property files)") + print(f" - instances.csv (benchmark instances)") print(f"{rule(100)}\n") + assert result['status'] == 'success', f"Download failed: {result['message']}" + + print(f"\n{rule(100)}") + print(f"✓ Successfully downloaded category: {normalized_name}") + print(f" Location: {result['category_path']}") + print(f" Instances: {result['num_instances']}") + print(f"{rule(100)}\n") elif creator_name == 'bert': print("BERT datasets are file-based.") print(f"Expected raw files under data/{normalized_name}/") @@ -312,6 +299,7 @@ def handle_unified_download(name: str, explicit_creator: Optional[str] = None): except ValueError as e: print(f"Error: {e}") + raise def print_list_downloads(creator: Optional[str] = None): @@ -408,7 +396,7 @@ def print_creator_info(name: str) -> None: except ValueError as e: print(f"Error: {e}") print(f"Available creators: {list_creators()}") - return + raise print(f"\n{rule(100)}") print(f"CREATOR: {name.upper()}") @@ -754,28 +742,25 @@ def main(): print(f"MODEL SYNTHESIS - {creator_name.upper()}") print(f"{rule(100)}\n") - try: - from act.front_end.model_synthesis import model_synthesis - from act.util.model_inference import model_inference - - wrapped_models = model_synthesis( - creator=creator_name, - spec_overrides=_build_spec_overrides(args), - text_verification_overrides=_build_text_verification_overrides(args), - ) - print(f"\n✓ Successfully synthesized {len(wrapped_models)} models") - - # Automatically run inference after synthesis - print(f"\n{rule(100)}") - print(f"MODEL INFERENCE - {creator_name.upper()}") - print(f"{rule(100)}\n") - - # model_inference extracts input from InputLayer - successful_models = model_inference(cast(Any, wrapped_models)) - print(f"\n✓ Successfully ran inference on {len(successful_models)}/{len(wrapped_models)} models") - print(f" Models are ready for verification") - except Exception as e: - print(f"\n✗ Synthesis/Inference failed: {e}") + from act.front_end.model_synthesis import model_synthesis + from act.util.model_inference import model_inference + + wrapped_models = model_synthesis( + creator=creator_name, + spec_overrides=_build_spec_overrides(args), + text_verification_overrides=_build_text_verification_overrides(args), + ) + print(f"\n✓ Successfully synthesized {len(wrapped_models)} models") + + # Automatically run inference after synthesis + print(f"\n{rule(100)}") + print(f"MODEL INFERENCE - {creator_name.upper()}") + print(f"{rule(100)}\n") + + # model_inference extracts input from InputLayer + successful_models = model_inference(cast(Any, wrapped_models)) + print(f"\n✓ Successfully ran inference on {len(successful_models)}/{len(wrapped_models)} models") + print(f" Models are ready for verification") elif args.inference: creator_name = args.creator if args.creator else 'torchvision' @@ -783,86 +768,80 @@ def main(): print(f"MODEL INFERENCE - {creator_name.upper()}") print(f"{rule(100)}\n") - try: - # Get downloaded pairs for the creator - downloaded_pairs = [] - if creator_name == 'torchvision': - downloaded_pairs = tv_loader.list_downloaded_pairs() - elif creator_name == 'vnnlib': - from act.front_end.vnnlib_loader import data_model_loader as vnnlib_loader - downloaded_pairs = vnnlib_loader.list_downloaded_pairs() - - if not downloaded_pairs: - print(f"⚠️ No downloaded models found for creator '{creator_name}'") - print(f" Use --download to download datasets/categories first") - print(f" Example: python -m act.front_end --download MNIST --creator {creator_name}") - else: - print(f"Found {len(downloaded_pairs)} downloaded pair(s) for {creator_name}") - print(f"Running inference on downloaded models...\n") - - results = [] - - # Run inference on each downloaded pair - if creator_name == 'torchvision': - for pair_info in downloaded_pairs: - dataset_name = pair_info['dataset'] - model_name = pair_info['model'] - - # Test the pair - result = tv_loader.model_inference_with_dataset( - dataset_name=dataset_name, - model_name=model_name, - split='test', - verbose=False - ) - results.append(result) - - # Print result - if result['status'] == 'success': - print(f"✓ {dataset_name:20s} + {model_name:20s} - Output: {result['output_shape']}") - else: - print(f"✗ {dataset_name:20s} + {model_name:20s} - Error: {result.get('error', 'Unknown')[:50]}") + # Get downloaded pairs for the creator + downloaded_pairs = [] + if creator_name == 'torchvision': + downloaded_pairs = tv_loader.list_downloaded_pairs() + elif creator_name == 'vnnlib': + from act.front_end.vnnlib_loader import data_model_loader as vnnlib_loader + downloaded_pairs = vnnlib_loader.list_downloaded_pairs() + + if not downloaded_pairs: + print(f"⚠️ No downloaded models found for creator '{creator_name}'") + print(f" Use --download to download datasets/categories first") + print(f" Example: python -m act.front_end --download MNIST --creator {creator_name}") + assert downloaded_pairs, f"No downloaded models found for creator '{creator_name}'" + print(f"Found {len(downloaded_pairs)} downloaded pair(s) for {creator_name}") + print(f"Running inference on downloaded models...\n") + + results = [] + + # Run inference on each downloaded pair + if creator_name == 'torchvision': + for pair_info in downloaded_pairs: + dataset_name = pair_info['dataset'] + model_name = pair_info['model'] - elif creator_name == 'vnnlib': - from act.front_end.vnnlib_loader import data_model_loader as vnnlib_loader - - for pair_info in downloaded_pairs: - category = pair_info['category'] - onnx_model = pair_info['onnx_model'] - vnnlib_spec = pair_info['vnnlib_spec'] - - # Test the pair - result = vnnlib_loader.model_inference_with_vnnlib( - category=category, - onnx_model=onnx_model, - vnnlib_spec=vnnlib_spec, - verbose=False - ) - results.append(result) - - # Print result - if result['status'] == 'success': - print(f"✓ {category:20s} + {result['model']:30s} - Output: {result['output_shape']}") - else: - print(f"✗ {category:20s} + {result['model']:30s} - Error: {result.get('error', 'Unknown')[:50]}") + # Test the pair + result = tv_loader.model_inference_with_dataset( + dataset_name=dataset_name, + model_name=model_name, + split='test', + verbose=False + ) + results.append(result) - # Summary - successful = sum(1 for r in results if r['status'] == 'success') - failed = len(results) - successful + # Print result + if result['status'] == 'success': + print(f"✓ {dataset_name:20s} + {model_name:20s} - Output: {result['output_shape']}") + else: + print(f"✗ {dataset_name:20s} + {model_name:20s} - Error: {result.get('error', 'Unknown')[:50]}") + + elif creator_name == 'vnnlib': + from act.front_end.vnnlib_loader import data_model_loader as vnnlib_loader + + for pair_info in downloaded_pairs: + category = pair_info['category'] + onnx_model = pair_info['onnx_model'] + vnnlib_spec = pair_info['vnnlib_spec'] - print(f"\n{rule(100)}") - print(f"INFERENCE SUMMARY") - print(f"{rule(100)}") - print(f"✓ Successful: {successful}/{len(results)}") - if failed > 0: - print(f"✗ Failed: {failed}/{len(results)}") - print(f"{rule(100)}\n") + # Test the pair + result = vnnlib_loader.model_inference_with_vnnlib( + category=category, + onnx_model=onnx_model, + vnnlib_spec=vnnlib_spec, + verbose=False + ) + results.append(result) - except Exception as e: - import traceback - print(f"\n✗ Inference failed: {e}") - print(f"\nTraceback:") - traceback.print_exc() + # Print result + if result['status'] == 'success': + print(f"✓ {category:20s} + {result['model']:30s} - Output: {result['output_shape']}") + else: + print(f"✗ {category:20s} + {result['model']:30s} - Error: {result.get('error', 'Unknown')[:50]}") + + # Summary + successful = sum(1 for r in results if r['status'] == 'success') + failed = len(results) - successful + + print(f"\n{rule(100)}") + print(f"INFERENCE SUMMARY") + print(f"{rule(100)}") + print(f"✓ Successful: {successful}/{len(results)}") + if failed > 0: + print(f"✗ Failed: {failed}/{len(results)}") + print(f"{rule(100)}\n") + assert failed == 0, f"{failed}/{len(results)} inference run(s) failed" else: parser.print_help() diff --git a/act/config/pipeline_cli.py b/act/config/pipeline_cli.py index 8a2dbb484..77695d9ce 100644 --- a/act/config/pipeline_cli.py +++ b/act/config/pipeline_cli.py @@ -2,6 +2,11 @@ """ ACT Pipeline Command-Line Interface. +Exit statuses: +- 0: command completed successfully +- 1: runtime, verification, or unexpected error +- 3: unsupported model/specification + Provides fuzzing capabilities for neural network verification with support for: - VNNLib verification benchmarks (default) - TorchVision datasets (alternative) @@ -26,15 +31,18 @@ logger = logging.getLogger(__name__) from act.front_end.specs import OutputSpec from act.front_end.vnnlib_loader.create_specs import VNNLibSpecCreator +from act.front_end.vnnlib_loader.vnnlib_parser import UnsupportedSpecError from act.front_end.vnnlib_loader import data_model_loader as vnnlib_loader -from act.front_end.vnnlib_loader import category_mapping as vnnlib_mapping from act.front_end.torchvision_loader.create_specs import TorchVisionSpecCreator from act.front_end.torchvision_loader import data_model_loader as tv_loader -from act.front_end.torchvision_loader import data_model_mapping as tv_mapping from act.front_end.model_synthesis import synthesize_models_and_seeds_from_specs +from act.config.frontend_cli import ( + print_unified_list, print_unified_search, print_unified_info, + handle_unified_download, print_list_downloads, +) from act.pipeline.fuzzing.actfuzzer import ACTFuzzer, FuzzingConfig from act.pipeline.verification.per_neuron_bounds import PerNeuronCheckConfig -from act.config.config import PipelineConfig +from act.config.config import BackendConfig, PipelineConfig _FUZZ_MUTATION_WEIGHT_KEYS = frozenset( @@ -271,7 +279,6 @@ def _apply_pipeline_config_defaults(args: Any) -> PipelineConfig: return config - def print_header(): """Print simple header.""" print(f"\n{rule()}") @@ -280,257 +287,12 @@ def print_header(): print(f"{rule()}\n") -# ============================================================================ -# Data-Model Pair Management Commands -# ============================================================================ - - -def cmd_list_available(creator: str): - """List available datasets/categories.""" - print(f"\n{rule()}") - print(f"AVAILABLE DATA-MODEL PAIRS ({creator.upper()})") - print(f"{rule()}\n") - - if creator == "vnnlib": - categories = vnnlib_mapping.list_categories() - print(f"VNNLIB Categories ({len(categories)}):") - print(rule(80, "-")) - for cat_name in sorted(categories): - info = vnnlib_mapping.get_category_info(cat_name) - print(f" {cat_name:30s} ({info['type']}) - {info['description']}") - print(f" └─ Models: {info['models']}, Properties: {info['properties']}") - - elif creator == "torchvision": - datasets = sorted(tv_mapping.DATASET_MODEL_MAPPING.keys()) - print(f"TorchVision Datasets ({len(datasets)}):") - print(rule(80, "-")) - for ds_name in datasets: - info = tv_mapping.DATASET_MODEL_MAPPING[ds_name] - models = info.get("models", []) - print(f" {ds_name:30s} [{info.get('category', 'N/A')}]") - if models: - print( - f" └─ Models: {', '.join(models[:5])}{'...' if len(models) > 5 else ''}" - ) - - print(f"\n{rule()}\n") - - -def cmd_search(query: str, creator: str): - """Search for datasets/categories.""" - print(f"\n{rule()}") - print(f"SEARCH RESULTS: '{query}' ({creator.upper()})") - print(f"{rule()}\n") - - if creator == "vnnlib": - matches = vnnlib_mapping.search_categories(query) - if matches: - print(f"Found {len(matches)} VNNLIB categories:") - print(rule(80, "-")) - for cat_name in sorted(matches): - info = vnnlib_mapping.get_category_info(cat_name) - print(f" {cat_name:30s} ({info['type']}) - {info['description']}") - else: - print(f"No VNNLIB categories found for '{query}'") - - elif creator == "torchvision": - matches = tv_mapping.search_datasets(query) - if matches: - print(f"Found {len(matches)} TorchVision datasets:") - print(rule(80, "-")) - for ds_name in sorted(matches): - info = tv_mapping.DATASET_MODEL_MAPPING[ds_name] - print(f" {ds_name:30s} [{info.get('category', 'N/A')}]") - else: - print(f"No TorchVision datasets found for '{query}'") - - print(f"\n{rule()}\n") - - -def cmd_info(name: str, creator: str): - """Show detailed information about dataset/category.""" - print(f"\n{rule()}") - print(f"INFO: {name} ({creator.upper()})") - print(f"{rule()}\n") - - if creator == "vnnlib": - try: - info = vnnlib_mapping.get_category_info(name) - print(f"Category: {name}") - print(f"Type: {info['type']}") - print(f"Year: {info['year']}") - print(f"Description: {info['description']}") - print(f"\nModel Information:") - print(f" • Models: {info['models']}") - print(f" • Properties: {info['properties']}") - print(f" • Input Dim: {info['input_dim']}") - print(f" • Output Dim: {info['output_dim']}") - - # Check if downloaded - downloaded = vnnlib_loader.list_downloaded_pairs() - matching = [p for p in downloaded if p["category"] == name] - if matching: - print(f"\n✓ Downloaded: {len(matching)} instances") - else: - print(f"\n⚠ Not downloaded (use --download {name})") - except ValueError as e: - print(f"Error: {e}") - - elif creator == "torchvision": - try: - info = tv_mapping.get_dataset_info(name) - print(f"Dataset: {name}") - print(f"Category: {info.get('category', 'N/A')}") - print(f"Input Size: {info.get('input_size', 'N/A')}") - print(f"Classes: {info.get('num_classes', 'N/A')}") - - models = info.get("models", []) - if models: - print(f"\nRecommended Models ({len(models)}):") - for model in models: - print(f" • {model}") - - # Check if downloaded - downloaded = tv_loader.list_downloaded_pairs() - matching = [p for p in downloaded if p["dataset"] == name] - if matching: - print(f"\n✓ Downloaded: {len(matching)} model pairs") - else: - print( - f"\n⚠ Not downloaded (use --download {name} --creator torchvision)" - ) - except ValueError as e: - print(f"Error: {e}") - - print(f"\n{rule()}\n") - - -def cmd_download(name: str, creator: str): - """Download dataset/category.""" - print(f"\n{rule()}") - print(f"DOWNLOADING: {name} ({creator.upper()})") - print(f"{rule()}\n") - - if creator == "vnnlib": - try: - result = vnnlib_loader.download_vnnlib_category(name) - - if result["status"] == "success": - print(f"✓ Successfully downloaded: {name}") - print(f" Location: {result['category_path']}") - print(f" Instances: {result['num_instances']}") - else: - print(f"✗ Download failed: {result['message']}") - print( - f"\nNote: VNNLIB benchmarks must be downloaded manually from VNN-COMP." - ) - print(f"Expected location: data/vnnlib/{name}/") - print(f"\nManual steps:") - print( - f" 1. Visit: https://github.com/ChristopherBrix/vnncomp_benchmarks" - ) - print(f" 2. Download '{name}' benchmark") - print(f" 3. Extract to: data/vnnlib/{name}/") - print(f" 4. Ensure structure:") - print(f" - onnx/ (ONNX model files)") - print(f" - vnnlib/ (VNNLIB property files)") - print(f" - instances.csv (benchmark instances)") - except Exception as e: - print(f"✗ Download error: {e}") - - elif creator == "torchvision": - try: - info = tv_mapping.get_dataset_info(name) - models = info.get("models", []) - - if not models: - print(f"⚠ No models available for {name}") - return - - print(f"Downloading {name} with {len(models)} models...\n") - - success_count = 0 - for model in models: - result = tv_loader.download_dataset_model_pair(name, model) - if result["status"] == "success": - print(f"✓ {name} + {model}") - success_count += 1 - else: - print(f"✗ {name} + {model} - {result['message']}") - - print(f"\n{rule()}") - print(f"Downloaded {success_count}/{len(models)} model pairs") - print(f"{rule()}") - except Exception as e: - print(f"✗ Download error: {e}") - - print() - - -def cmd_list_downloaded(creator: str): - """List downloaded data-model pairs.""" - print(f"\n{rule()}") - print(f"DOWNLOADED DATA-MODEL PAIRS ({creator.upper()})") - print(f"{rule()}\n") - - if creator == "vnnlib": - downloaded = vnnlib_loader.list_downloaded_pairs() - if downloaded: - # Group by category - categories = {} - for item in downloaded: - cat = item["category"] - if cat not in categories: - categories[cat] = [] - categories[cat].append(item) - - print(f"VNNLIB Downloads ({len(downloaded)} instances):") - print(rule(80, "-")) - for cat in sorted(categories.keys()): - instances = categories[cat] - print(f" {cat:30s} ({len(instances)} instances)") - if len(instances) <= 5: - for inst in instances: - print( - f" └─ {inst['instance_id']}: {inst['onnx_model']} + {inst['vnnlib_spec']}" - ) - else: - print("No VNNLIB downloads found") - print("Use --download to download benchmarks") - - elif creator == "torchvision": - downloaded = tv_loader.list_downloaded_pairs() - if downloaded: - # Group by dataset - datasets = {} - for item in downloaded: - ds = item["dataset"] - if ds not in datasets: - datasets[ds] = [] - datasets[ds].append(item["model"]) - - print(f"TorchVision Downloads ({len(downloaded)} pairs):") - print(rule(80, "-")) - for ds in sorted(datasets.keys()): - models = datasets[ds] - print(f" {ds:30s} ({len(models)} models)") - for model in sorted(models): - print(f" └─ {model}") - else: - print("No TorchVision downloads found") - print( - "Use --download --creator torchvision to download data-model pairs" - ) - - print(f"\n{rule()}\n") - - # ============================================================================ # Fuzzing Commands # ============================================================================ -def cmd_fuzz(args): +def cmd_fuzz(args) -> int: """Run ACTFuzzer.""" print_header() @@ -566,7 +328,7 @@ def cmd_fuzz(args): if not downloaded: print("❌ No VNNLIB categories downloaded!") print("Use: python -m act.pipeline --download ") - return + return 1 categories = list(set(p["category"] for p in downloaded)) print(f"Loading {len(categories)} VNNLIB category(ies):") @@ -592,7 +354,7 @@ def cmd_fuzz(args): print( "Use: python -m act.pipeline --download --creator torchvision" ) - return + return 1 datasets = list(set(p["dataset"] for p in downloaded)) print(f"Loading {len(datasets)} TorchVision dataset(s):") @@ -615,7 +377,7 @@ def cmd_fuzz(args): if not model_names: print("❌ No models found for selected datasets!") - return + return 1 spec_results = spec_creator.create_specs_for_data_model_pairs( dataset_names=datasets, @@ -644,11 +406,11 @@ def cmd_fuzz(args): import traceback traceback.print_exc() - return + return 1 if not spec_results: print("❌ No spec results generated!") - return + return 1 print(f"✓ Generated {len(spec_results)} spec result(s)\n") @@ -671,11 +433,11 @@ def cmd_fuzz(args): import traceback traceback.print_exc() - return + return 1 if not wrapped_models: print("❌ No models synthesized!") - return + return 1 print(f"✓ Synthesized {len(wrapped_models)} wrapped model(s)\n") @@ -689,7 +451,7 @@ def cmd_fuzz(args): if not initial_seeds: print("❌ No initial seeds extracted!") - return + return 1 print(f"✓ Extracted {len(initial_seeds)} initial seeds\n") @@ -722,15 +484,12 @@ def cmd_fuzz(args): if report.counterexamples and config.save_counterexamples: import os - import torch as _torch from act.front_end.vnnlib_loader.vnnlib_parser import write_vnncomp_result os.makedirs(config.output_dir, exist_ok=True) ce0 = report.counterexamples[0] - x = cast(Any, getattr(ce0, "input", ce0)) - with _torch.no_grad(): - y_raw = wrapped_model(x) - y = cast(Any, y_raw["output"] if isinstance(y_raw, dict) else y_raw) + x = ce0.input.unsqueeze(0) + y = ce0.output.unsqueeze(0) fname = "_".join(map(str, model_id)) if isinstance(model_id, tuple) else str(model_id) write_vnncomp_result( os.path.join(config.output_dir, f"{fname}_result.txt"), @@ -745,7 +504,9 @@ def cmd_fuzz(args): import traceback traceback.print_exc() - return + return 1 + + return 0 # ============================================================================ @@ -1023,6 +784,7 @@ def _run_bab_on_net(net, args, bab_first_sample_only: bool = False): pipeline_config = PipelineConfig.from_yaml(**_collect_pipeline_config_overrides(args)) config = pipeline_config.bab dual_config = pipeline_config.dual + max_batch_size = BackendConfig.from_yaml().bab_max_batch_size budget = float(getattr(args, "timeout", 60.0) or 60.0) spec_layers = gather_input_spec_layers(net) @@ -1034,7 +796,7 @@ def _run_bab_on_net(net, args, bab_first_sample_only: bool = False): net=net, solver_factory=TorchLPSolver, config=config, - max_batch_size=None, + max_batch_size=max_batch_size, time_budget_s=budget, dual_config=dual_config, ) @@ -1049,7 +811,7 @@ def _run_bab_on_net(net, args, bab_first_sample_only: bool = False): net=sliced_net, solver_factory=TorchLPSolver, config=config, - max_batch_size=None, + max_batch_size=max_batch_size, time_budget_s=budget, dual_config=dual_config, ) @@ -1285,6 +1047,8 @@ def cmd_verify(target: str, args): try: soundness_failed = _run_vnnlib_verify(args) results[test_name] = "FAILED" if soundness_failed else "PASSED" + except UnsupportedSpecError: + raise except Exception as e: print(f"\n❌ Test failed: {e}") import traceback @@ -1391,6 +1155,11 @@ def main(): python -m act.pipeline --verify netfactory --solvers torchlp --tf-modes interval --validate-soundness python -m act.pipeline --verify vnnlib --category acasxu_2023 --max-instances 3 --validate-soundness python -m act.pipeline --verify torchvision --dataset MNIST --model simple_cnn --num-samples 2 --validate-soundness + +Exit statuses: + 0 Command completed successfully + 1 Runtime, verification, or unexpected error + 3 Unsupported model/specification """, ) @@ -1784,22 +1553,25 @@ def main(): # Execute command try: if args.list: - cmd_list_available(args.creator) + print_unified_list(args.creator) elif args.search: - cmd_search(args.search, args.creator) + print_unified_search(args.search, args.creator) elif args.info: - cmd_info(args.info, args.creator) + print_unified_info(args.info, explicit_creator=args.creator) elif args.download: - cmd_download(args.download, args.creator) + handle_unified_download(args.download, explicit_creator=args.creator) elif args.list_downloaded: - cmd_list_downloaded(args.creator) + print_list_downloads(args.creator) elif args.fuzz: - cmd_fuzz(args) + sys.exit(cmd_fuzz(args)) elif args.verify: cmd_verify(args.verify, args) except KeyboardInterrupt: print("\n\n⚠️ Interrupted by user") sys.exit(1) + except UnsupportedSpecError as e: + print(f"\n❌ Unsupported model/specification: {e}") + sys.exit(3) except Exception as e: print(f"\n❌ Error: {e}") import traceback diff --git a/act/front_end/bert_loader/data_loader.py b/act/front_end/bert_loader/data_loader.py index ff22e42cd..ce74b9d5a 100644 --- a/act/front_end/bert_loader/data_loader.py +++ b/act/front_end/bert_loader/data_loader.py @@ -362,7 +362,7 @@ def tokenize(sentence: str) -> list[str]: import nltk # type: ignore[import-untyped] return nltk.word_tokenize(sentence) - except Exception as exc: + except (ImportError, LookupError) as exc: # nltk or its punkt data not installed logger.debug("Falling back to regex tokenization: %s", exc) return _TOKEN_RE.findall(sentence) diff --git a/act/front_end/spec_creator_base.py b/act/front_end/spec_creator_base.py index fc713cc0c..3f61178d4 100644 --- a/act/front_end/spec_creator_base.py +++ b/act/front_end/spec_creator_base.py @@ -9,7 +9,7 @@ from __future__ import annotations from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import List, Tuple, Dict, Any, Optional, Callable, Union +from typing import List, Tuple, Dict, Any, Optional, Union import logging import torch @@ -244,20 +244,7 @@ def __init__( def _load_config(self, config_name: Optional[str] = None) -> Dict[str, Any]: """Load a named spec configuration from the front-end config.""" - try: - return FrontEndConfig.from_yaml().spec_config(config_name) - except Exception as e: - print(f"⚠️ Could not load config: {e}, using defaults") - return self._get_default_config() - - def _get_default_config(self) -> Dict[str, Any]: - """Fallback default configuration""" - return { - 'input_spec_types': ['BOX', 'LINF_BALL'], - 'output_spec_types': ['MARGIN_ROBUST', 'TOP1_ROBUST'], - 'perturbation': {'epsilon_values': [0.01, 0.03], 'norm': 'inf'}, - 'combination_strategy': 'balanced' - } + return FrontEndConfig.from_yaml().spec_config(config_name) # ==================== SHAPE VALIDATION ==================== # @@ -355,16 +342,12 @@ def validate_spec_pair_with_model( errors = [] # Get model I/O shapes - try: - input_shape, output_shape = self._get_model_io_shapes(model, sample_input) - # Assume last dimension is classes for classification - if len(output_shape) > 1: - num_classes = output_shape[-1] - else: - num_classes = output_shape[0] - except Exception as e: - errors.append(f"Failed to infer model shapes: {e}") - return False, errors + input_shape, output_shape = self._get_model_io_shapes(model, sample_input) + # Assume last dimension is classes for classification + if len(output_shape) > 1: + num_classes = output_shape[-1] + else: + num_classes = output_shape[0] # Validate input spec input_valid, input_err = self._validate_input_spec_shape(input_spec, input_shape) @@ -387,19 +370,16 @@ def _validate_and_filter_specs( """Validate spec pairs against model and filter invalid ones.""" valid_pairs = [] for input_spec, output_spec in spec_pairs: - try: - is_valid, errors = self.validate_spec_pair_with_model( - input_spec, output_spec, pytorch_model, sample_input + is_valid, errors = self.validate_spec_pair_with_model( + input_spec, output_spec, pytorch_model, sample_input + ) + if is_valid: + valid_pairs.append((input_spec, output_spec)) + else: + logger.debug( + "Spec pair validation failed: %s, %s, with errors:\n%s", + input_spec.kind, output_spec.kind, "\n".join(errors), ) - if is_valid: - valid_pairs.append((input_spec, output_spec)) - else: - logger.debug( - "Spec pair validation failed: %s, %s, with errors:\n%s", - input_spec.kind, output_spec.kind, "\n".join(errors), - ) - except Exception as e: - logger.debug("Spec validation error: %s", e) return valid_pairs # ==================== SPEC COMBINATION ==================== # @@ -439,17 +419,13 @@ def _create_spec_combinations( @abstractmethod def create_specs_for_data_model_pairs( self, - max_samples: Optional[int] = None, - filter_fn: Optional[Callable] = None, - validate_shapes: bool = True - ) -> List[Tuple[str, str, torch.nn.Module, List[torch.Tensor], List[Tuple[InputSpec, OutputSpec]]]]: + *args: Any, + **kwargs: Any, + ) -> Any: """ Create specs for data-model pairs (must be implemented by subclasses). - - Args: - max_samples: Maximum number of samples/instances to process - filter_fn: Optional filter function (source, model) -> bool - validate_shapes: Whether to validate spec shapes with model + + Concrete creators expose source-specific selection and limit options. Returns: List of (data_source, model_name, pytorch_model, input_tensors, spec_pairs) tuples diff --git a/act/front_end/torchvision_loader/cli.py b/act/front_end/torchvision_loader/cli.py index e2d48fc83..f7b1b03ac 100644 --- a/act/front_end/torchvision_loader/cli.py +++ b/act/front_end/torchvision_loader/cli.py @@ -313,12 +313,14 @@ def main(): print(f" • {name}") except ValueError as e: print(f"Error: {e}") + raise elif args.dataset: try: print_dataset_detail(args.dataset) except ValueError as e: print(f"Error: {e}") + raise elif args.summary: print_mapping_summary() @@ -335,10 +337,7 @@ def main(): split=args.split ) - if result['status'] == 'error': - print(f"\n❌ Error: {result['message']}") - if 'traceback' in result: - print(f"\nDetails:\n{result['traceback']}") + assert result['status'] == 'success', f"Download failed: {result['message']}" # Success message already printed by download function elif args.list_downloads: @@ -432,12 +431,10 @@ def main(): print(f"✗ Error: {e}") print(f"\nThe dataset-model pair has not been downloaded yet.") print(f"Use: python -m act.front_end.torchvision_loader.cli --download {dataset_name} {model_name}") + raise except ValueError as e: print(f"✗ Error: {e}") - except Exception as e: - print(f"✗ Unexpected error: {e}") - import traceback - print(f"\nDetails:\n{traceback.format_exc()}") + raise elif args.models_for: # Show all recommended models for a dataset (case-insensitive) @@ -497,6 +494,7 @@ def main(): except ValueError as e: print(f"Error: {e}") print(f"Use --list to see all available datasets") + raise elif args.datasets_for: # Show all datasets that can run inference with a model (case-insensitive) @@ -603,18 +601,17 @@ def main(): print(f" Dataset samples: {result['num_samples']}") print(f"{rule()}") + assert result['status'] == 'success', f"Inference failed: {result.get('error')}" except FileNotFoundError as e: print(f"\n✗ Error: Dataset-model pair not found locally") print(f" {e}") print(f"\nDownload first with:") print(f" python -m act.front_end.torchvision_loader.cli --download {dataset_name} {model_name}") + raise except ValueError as e: print(f"\n✗ Error: {e}") - except Exception as e: - print(f"\n✗ Unexpected error: {e}") - import traceback - print(f"\nDetails:\n{traceback.format_exc()}") + raise elif args.validate: user_dataset, user_model = args.validate @@ -643,6 +640,7 @@ def main(): print(f"{rule(100)}") except ValueError as e: print(f"Error: {e}") + raise elif args.show_preprocessing: try: @@ -677,6 +675,7 @@ def main(): print(f"{rule(100)}") except ValueError as e: print(f"Error: {e}") + raise else: # Default: show quick examples diff --git a/act/front_end/torchvision_loader/create_specs.py b/act/front_end/torchvision_loader/create_specs.py index 7a75506ca..a3a511f4f 100644 --- a/act/front_end/torchvision_loader/create_specs.py +++ b/act/front_end/torchvision_loader/create_specs.py @@ -135,6 +135,7 @@ def create_specs_for_data_model_pairs( logger.info(f"Processing {len(all_pairs)} dataset-model pairs") results = [] + failures: List[str] = [] for pair_info in all_pairs: dataset_name = pair_info['dataset'] @@ -166,8 +167,7 @@ def create_specs_for_data_model_pairs( validate_shapes=validate_shapes ) - if result is not None: - results.append(result) + results.append(result) # Memory optimization: Free dataset/dataloader after extracting input_tensors # pair_data contains the dataset (476 MB for MNIST) which is no longer needed @@ -175,11 +175,17 @@ def create_specs_for_data_model_pairs( del pair_data, dataloader gc.collect() - except Exception as e: + except Exception as e: # noqa: BLE001 — collect every failure, raise below logger.error( f"Failed to create specs for {dataset_name} + {model_name}: {e}" ) + failures.append(f"{dataset_name} + {model_name}: {e}") + if failures: + raise RuntimeError( + f"Spec creation failed for {len(failures)}/{len(all_pairs)} dataset-model pair(s):\n " + + "\n ".join(failures) + ) logger.info(f"Successfully created specs for {len(results)} pairs") return results @@ -192,7 +198,7 @@ def _create_specs_for_single_instance( num_samples: int, start_index: int, validate_shapes: bool - ) -> Optional[Tuple[str, str, nn.Module, List[LabeledInputTensor], List[Tuple[InputSpec, OutputSpec]]]]: + ) -> Tuple[str, str, nn.Module, List[LabeledInputTensor], List[Tuple[InputSpec, OutputSpec]]]: """ Create specs for a single dataset-model pair. @@ -218,8 +224,7 @@ def _create_specs_for_single_instance( labeled_tensors.append(LabeledInputTensor(tensor=tensor, label=label)) if not labeled_tensors: - logger.warning(f"No samples collected for {data_source}") - return None + raise RuntimeError(f"No samples collected for {data_source} + {model_name}") logger.info(f"Collected {len(labeled_tensors)} samples") @@ -261,8 +266,7 @@ def _create_specs_for_single_instance( spec_pairs = validated_pairs if not spec_pairs: - logger.warning(f"No valid specs generated for {data_source} + {model_name}") - return None + raise RuntimeError(f"No valid specs generated for {data_source} + {model_name}") return (data_source, model_name, pytorch_model, labeled_tensors, spec_pairs) diff --git a/act/front_end/torchvision_loader/data_model_loader.py b/act/front_end/torchvision_loader/data_model_loader.py index 22e7132aa..e05cbad9c 100644 --- a/act/front_end/torchvision_loader/data_model_loader.py +++ b/act/front_end/torchvision_loader/data_model_loader.py @@ -131,19 +131,14 @@ def _download_dataset( # Archive datasets ship train+test in one archive: download # once, then both splits are available regardless of `split`. print(f" • Downloading archive from {dataset_info['download']['url']} ...") - try: - for split_name in ('test', 'train'): - ds = _resolve_external_dataset( - dataset_info, raw_dir, - train=(split_name == 'train'), - download=True, - ) - downloaded_splits.append(split_name) - print(f" ✓ {split_name.capitalize()} split: {len(ds)} samples") - except Exception as e: - error = f"External archive download failed: {e}" - errors.append(error) - print(" ⚠ " + error) + for split_name in ('test', 'train'): + ds = _resolve_external_dataset( + dataset_info, raw_dir, + train=(split_name == 'train'), + download=True, + ) + downloaded_splits.append(split_name) + print(f" ✓ {split_name.capitalize()} split: {len(ds)} samples") else: import torchvision.datasets dataset_class = getattr(torchvision.datasets, dataset_name, None) @@ -159,19 +154,14 @@ def _download_dataset( if split not in (split_name, 'both'): continue print(f" • Downloading {split_name} split...") - try: - # Split-selection kwarg differs per dataset (train= / split= / image_set= / version= / background= / none), so unpack the per-dataset set declared in split_kwarg. - ds = dataset_class( - root=str(raw_dir), - **dataset_info["split_kwarg"][split_name], - download=True - ) - downloaded_splits.append(split_name) - print(f" ✓ {split_name.capitalize()} split: {len(ds)} samples") - except Exception as e: - error = f"{split_name.capitalize()} split failed: {e}" - errors.append(error) - print(" ⚠ " + error) + # Split-selection kwarg differs per dataset (train= / split= / image_set= / version= / background= / none), so unpack the per-dataset set declared in split_kwarg. + ds = dataset_class( + root=str(raw_dir), + **dataset_info["split_kwarg"][split_name], + download=True + ) + downloaded_splits.append(split_name) + print(f" ✓ {split_name.capitalize()} split: {len(ds)} samples") return downloaded_splits, errors @@ -211,98 +201,93 @@ def download_dataset_model_pair( >>> print(result['dataset_path']) '/path/to/ACT/data/torchvision/MNIST/raw/' """ - import traceback - # Use path from config if not specified if root_dir is None: root_dir = get_torchvision_data_root() - try: - # Normalize names (case-insensitive) - dataset_name = find_dataset_name(dataset_name) - model_name = find_model_name(model_name) - - # Validate compatibility - validation = validate_dataset_model_compatibility(dataset_name, model_name) - if not validation['compatible']: + # Normalize names (case-insensitive) + dataset_name = find_dataset_name(dataset_name) + model_name = find_model_name(model_name) + + # Validate compatibility + validation = validate_dataset_model_compatibility(dataset_name, model_name) + if not validation['compatible']: + return { + 'status': 'error', + 'message': f"Incompatible pair: {', '.join(validation['issues'])}" + } + + # Get dataset info + dataset_info = get_dataset_info(dataset_name) + + # Create directory structure + dataset_dir = Path(root_dir) / dataset_name + raw_dir = dataset_dir / "raw" + models_dir = dataset_dir / "models" + + # Check if dataset already exists (for incremental download) + dataset_exists = dataset_dir.exists() and raw_dir.exists() + model_file_path = models_dir / f"{model_name}.py" + model_already_exists = model_file_path.exists() + + if dataset_exists: + if model_already_exists: + print(f"\n⚠️ Model '{model_name}' already exists for dataset '{dataset_name}'") + print(f" Skipping download (pair already complete)") return { - 'status': 'error', - 'message': f"Incompatible pair: {', '.join(validation['issues'])}" + 'status': 'success', + 'message': f"Model {model_name} already exists for {dataset_name}", + 'path': str(dataset_dir) } - - # Get dataset info - dataset_info = get_dataset_info(dataset_name) - - # Create directory structure - dataset_dir = Path(root_dir) / dataset_name - raw_dir = dataset_dir / "raw" - models_dir = dataset_dir / "models" - - # Check if dataset already exists (for incremental download) - dataset_exists = dataset_dir.exists() and raw_dir.exists() - model_file_path = models_dir / f"{model_name}.py" - model_already_exists = model_file_path.exists() - - if dataset_exists: - if model_already_exists: - print(f"\n⚠️ Model '{model_name}' already exists for dataset '{dataset_name}'") - print(f" Skipping download (pair already complete)") - return { - 'status': 'success', - 'message': f"Model {model_name} already exists for {dataset_name}", - 'path': str(dataset_dir) - } - else: - print(f"\n✓ Dataset '{dataset_name}' already downloaded - adding model '{model_name}'") - print(f" Location: {raw_dir}") - - # Load existing metadata to check splits - info_path = dataset_dir / "info.json" - if info_path.exists(): - with open(info_path, 'r') as f: - existing_metadata = json.load(f) - downloaded_splits = existing_metadata.get('splits_downloaded', []) - print(f" Existing splits: {', '.join(downloaded_splits)}") - else: - downloaded_splits = [] - - # Create models directory if it doesn't exist - models_dir.mkdir(parents=True, exist_ok=True) else: - # Fresh download - create directories - raw_dir.mkdir(parents=True, exist_ok=True) + print(f"\n✓ Dataset '{dataset_name}' already downloaded - adding model '{model_name}'") + print(f" Location: {raw_dir}") + + # Load existing metadata to check splits + info_path = dataset_dir / "info.json" + if info_path.exists(): + with open(info_path, 'r') as f: + existing_metadata = json.load(f) + downloaded_splits = existing_metadata.get('splits_downloaded', []) + print(f" Existing splits: {', '.join(downloaded_splits)}") + else: + downloaded_splits = [] + + # Create models directory if it doesn't exist models_dir.mkdir(parents=True, exist_ok=True) - downloaded_splits = [] - - print(f"\n{rule()}") - print(f"{'ADDING MODEL' if dataset_exists else 'DOWNLOADING'}: {dataset_name} + {model_name}") - print(f"{rule()}") - print(f"Target: {dataset_dir}") - if not dataset_exists: - print(f"Split: {split}") - - # Download dataset (only if not already present) - if not dataset_exists: - print(f"\n[1/3] Downloading dataset...") + else: + # Fresh download - create directories + raw_dir.mkdir(parents=True, exist_ok=True) + models_dir.mkdir(parents=True, exist_ok=True) + downloaded_splits = [] + + print(f"\n{rule()}") + print(f"{'ADDING MODEL' if dataset_exists else 'DOWNLOADING'}: {dataset_name} + {model_name}") + print(f"{rule()}") + print(f"Target: {dataset_dir}") + if not dataset_exists: + print(f"Split: {split}") + + # Download dataset (only if not already present) + if not dataset_exists: + print(f"\n[1/3] Downloading dataset...") - downloaded_splits, errors = _download_dataset(root_dir, dataset_name, dataset_info, split) + downloaded_splits, errors = _download_dataset(root_dir, dataset_name, dataset_info, split) - if not downloaded_splits: - return { - 'status': 'error', - 'message': f"Failed to download any splits for {dataset_name}, with errors: {"\n".join(errors)}" - } - else: - print(f"\n[1/3] Dataset already present, skipping download...") - - # Save model architecture - print(f"\n[2/3] Saving model architecture...") - model_path = models_dir / f"{model_name}.py" - - import torchvision.models - if hasattr(torchvision.models, model_name): - # Standard TorchVision model - model_code = f""" + assert downloaded_splits and not errors, ( + f"Failed to download {dataset_name}: {'; '.join(errors)}" + ) + else: + print(f"\n[1/3] Dataset already present, skipping download...") + + # Save model architecture + print(f"\n[2/3] Saving model architecture...") + model_path = models_dir / f"{model_name}.py" + + import torchvision.models + if hasattr(torchvision.models, model_name): + # Standard TorchVision model + model_code = f""" # Standard TorchVision Model: {model_name} # Can be loaded with: torchvision.models.{model_name}() @@ -321,106 +306,92 @@ def download_dataset_model_pair( print(f"Model: {model_name}") print(f"Parameters: {{sum(p.numel() for p in model.parameters()):,}}") """ - else: - # Custom model - get definition - model_code = _get_custom_model_definition(model_name, dataset_info['num_classes']) - - with open(model_path, 'w') as f: - f.write(model_code) - print(f" ✓ Model architecture saved: {model_path}") - - # Save/update metadata - print(f"\n[3/3] Saving metadata...") - info_path = dataset_dir / "info.json" + else: + # Custom model - get definition + model_code = _get_custom_model_definition(model_name, dataset_info['num_classes']) + + with open(model_path, 'w') as f: + f.write(model_code) + print(f" ✓ Model architecture saved: {model_path}") + + # Save/update metadata + print(f"\n[3/3] Saving metadata...") + info_path = dataset_dir / "info.json" + + # Load existing metadata if it exists + if info_path.exists(): + with open(info_path, 'r') as f: + existing_metadata = json.load(f) - # Load existing metadata if it exists - if info_path.exists(): - try: - with open(info_path, 'r') as f: - existing_metadata = json.load(f) - - # Check if this is old format (single model) or new format (list of models) - if 'models' not in existing_metadata: - # Old format - convert to new format - old_model = existing_metadata.get('model') - models_list = [old_model] if old_model else [] - else: - models_list = existing_metadata.get('models', []) - - # Add new model if not already in list - if model_name not in models_list: - models_list.append(model_name) - - # Preserve existing splits if incremental - if dataset_exists: - downloaded_splits = existing_metadata.get('splits_downloaded', downloaded_splits) - except Exception as e: - print(f" ⚠ Could not read existing metadata: {e}") - models_list = [model_name] + # Check if this is old format (single model) or new format (list of models) + if 'models' not in existing_metadata: + # Old format - convert to new format + old_model = existing_metadata.get('model') + models_list = [old_model] if old_model else [] else: - models_list = [model_name] - - # Create updated metadata with list of models - metadata = { - 'dataset': dataset_name, - 'models': models_list, # Changed from single 'model' to list 'models' - 'category': dataset_info['category'], - 'input_size': dataset_info['input_size'], - 'num_classes': dataset_info['num_classes'], - 'splits_downloaded': downloaded_splits, - 'preprocessing_required': validation['preprocessing_required'], - 'preprocessing_steps': validation['preprocessing_steps'], - 'notes': dataset_info['notes'], - 'paths': { - 'raw_data': str(raw_dir), - 'models_directory': str(models_dir), - 'metadata': str(info_path) - } - } + models_list = existing_metadata.get('models', []) - with open(info_path, 'w') as f: - json.dump(metadata, f, indent=2) - print(f" ✓ Metadata saved: {info_path}") - print(f" ✓ Models registered: {', '.join(models_list)}") + # Add new model if not already in list + if model_name not in models_list: + models_list.append(model_name) - # Calculate directory size - total_size_bytes = _get_directory_size(dataset_dir) - total_size_formatted = _format_size(total_size_bytes) - - # Print summary - print(f"\n{rule()}") - print(f"✓ DOWNLOAD COMPLETE") - print(f"{rule()}") - print(f"Dataset: {raw_dir}") - print(f"Model: {model_path}") - print(f"Info: {info_path}") - print(f"Size: {total_size_formatted} ({total_size_bytes:,} bytes)") - - if validation['preprocessing_required']: - print(f"\n⚠ Preprocessing Required:") - for step in validation['preprocessing_steps']: - print(f" • {step}") - - return { - 'status': 'success', - 'dataset_path': str(raw_dir), - 'model_path': str(model_path), - 'info_path': str(info_path), - 'message': 'Download completed successfully', - 'metadata': metadata, - 'size_bytes': total_size_bytes, - 'size_formatted': total_size_formatted - } - - except AssertionError: - # Re-raise AssertionError from validation without catching - raise - except Exception as e: - return { - 'status': 'error', - 'message': f"Download failed: {str(e)}", - 'traceback': traceback.format_exc() + # Preserve existing splits if incremental + if dataset_exists: + downloaded_splits = existing_metadata.get('splits_downloaded', downloaded_splits) + else: + models_list = [model_name] + + # Create updated metadata with list of models + metadata = { + 'dataset': dataset_name, + 'models': models_list, # Changed from single 'model' to list 'models' + 'category': dataset_info['category'], + 'input_size': dataset_info['input_size'], + 'num_classes': dataset_info['num_classes'], + 'splits_downloaded': downloaded_splits, + 'preprocessing_required': validation['preprocessing_required'], + 'preprocessing_steps': validation['preprocessing_steps'], + 'notes': dataset_info['notes'], + 'paths': { + 'raw_data': str(raw_dir), + 'models_directory': str(models_dir), + 'metadata': str(info_path) } + } + + with open(info_path, 'w') as f: + json.dump(metadata, f, indent=2) + print(f" ✓ Metadata saved: {info_path}") + print(f" ✓ Models registered: {', '.join(models_list)}") + + # Calculate directory size + total_size_bytes = _get_directory_size(dataset_dir) + total_size_formatted = _format_size(total_size_bytes) + + # Print summary + print(f"\n{rule()}") + print(f"✓ DOWNLOAD COMPLETE") + print(f"{rule()}") + print(f"Dataset: {raw_dir}") + print(f"Model: {model_path}") + print(f"Info: {info_path}") + print(f"Size: {total_size_formatted} ({total_size_bytes:,} bytes)") + + if validation['preprocessing_required']: + print(f"\n⚠ Preprocessing Required:") + for step in validation['preprocessing_steps']: + print(f" • {step}") + + return { + 'status': 'success', + 'dataset_path': str(raw_dir), + 'model_path': str(model_path), + 'info_path': str(info_path), + 'message': 'Download completed successfully', + 'metadata': metadata, + 'size_bytes': total_size_bytes, + 'size_formatted': total_size_formatted + } def list_downloaded_pairs(root_dir: Optional[str] = None) -> List[dict]: @@ -455,46 +426,42 @@ def list_downloaded_pairs(root_dir: Optional[str] = None) -> List[dict]: if not info_path.exists(): continue - try: - with open(info_path, 'r') as f: - metadata = json.load(f) - - # Calculate directory size (once per dataset) - size_bytes = _get_directory_size(dataset_dir) - size_formatted = _format_size(size_bytes) - - # info.json may store either a list under 'models' or a single - # 'model' string; normalise both into models_list. - if 'models' in metadata: - models_list = metadata['models'] - else: - model = metadata.get('model') - models_list = [model] if model else [] - - # Create one entry per model - for model_name in models_list: - pair_metadata = { - 'dataset': metadata['dataset'], - 'model': model_name, - 'category': metadata.get('category', 'unknown'), - 'input_size': metadata.get('input_size', []), - 'num_classes': metadata.get('num_classes', 0), - 'splits_downloaded': metadata.get('splits_downloaded', []), - 'preprocessing_required': metadata.get('preprocessing_required', True), - 'preprocessing_steps': metadata.get('preprocessing_steps', []), - 'notes': metadata.get('notes', ''), - 'paths': { - 'raw_data': metadata['paths'].get('raw_data', ''), - 'model_architecture': str(dataset_dir / "models" / f"{model_name}.py"), - 'metadata': str(info_path) - }, - 'size_bytes': size_bytes, - 'size_formatted': size_formatted - } - downloaded.append(pair_metadata) - - except Exception as e: - print(f"Warning: Could not read {info_path}: {e}") + with open(info_path, 'r') as f: + metadata = json.load(f) + + # Calculate directory size (once per dataset) + size_bytes = _get_directory_size(dataset_dir) + size_formatted = _format_size(size_bytes) + + # info.json may store either a list under 'models' or a single + # 'model' string; normalise both into models_list. + if 'models' in metadata: + models_list = metadata['models'] + else: + model = metadata.get('model') + models_list = [model] if model else [] + + # Create one entry per model + for model_name in models_list: + pair_metadata = { + 'dataset': metadata['dataset'], + 'model': model_name, + 'category': metadata.get('category', 'unknown'), + 'input_size': metadata.get('input_size', []), + 'num_classes': metadata.get('num_classes', 0), + 'splits_downloaded': metadata.get('splits_downloaded', []), + 'preprocessing_required': metadata.get('preprocessing_required', True), + 'preprocessing_steps': metadata.get('preprocessing_steps', []), + 'notes': metadata.get('notes', ''), + 'paths': { + 'raw_data': metadata['paths'].get('raw_data', ''), + 'model_architecture': str(dataset_dir / "models" / f"{model_name}.py"), + 'metadata': str(info_path) + }, + 'size_bytes': size_bytes, + 'size_formatted': size_formatted + } + downloaded.append(pair_metadata) return downloaded @@ -761,114 +728,88 @@ def model_inference_with_dataset( combo_id = f"m:{model_name}|x:{dataset_name}|split:{split}" - try: - # Load the pair - if verbose: - print(f"Testing: {dataset_name} + {model_name} ({split} split)") - - result = load_dataset_model_pair( - dataset_name=dataset_name, - model_name=model_name, - split=split, - batch_size=1, - shuffle=False, - auto_download=True, - root_dir=root_dir - ) - - model = result['model'] - dataloader = result['dataloader'] - num_samples = result['num_samples'] - - # Get first sample to test - input_tensor = None - try: - for batch in dataloader: - if isinstance(batch, (tuple, list)) and len(batch) >= 2: - # Standard (image, label) tuple - input_tensor, labels = batch[0], batch[1] - elif isinstance(batch, (tuple, list)) and len(batch) == 1: - input_tensor = batch[0] - else: - input_tensor = batch - break # Only need first batch - except Exception as e: - error_msg = f'Batch loading error: {str(e)[:50]}' - if verbose: - print(f" ✗ FAILED - {error_msg}") - return { - 'dataset': dataset_name, - 'model': model_name, - 'split': split, - 'status': 'failed', - 'error': error_msg, - 'num_samples': num_samples - } - - if input_tensor is None: - error_msg = 'Empty dataloader' - if verbose: - print(f" ✗ FAILED - {error_msg}") - return { - 'dataset': dataset_name, - 'model': model_name, - 'split': split, - 'status': 'failed', - 'error': error_msg, - 'num_samples': num_samples - } - - # Verify we have a tensor - if not isinstance(input_tensor, torch.Tensor): - error_msg = f'Input type {type(input_tensor).__name__}, expected torch.Tensor' - if verbose: - print(f" ✗ FAILED - {error_msg}") - return { - 'dataset': dataset_name, - 'model': model_name, - 'split': split, - 'status': 'failed', - 'error': error_msg, - 'num_samples': num_samples - } - - # Run inference - success, output, error_msg = infer_single_model(combo_id, model, input_tensor) - - if success: - output_shape = tuple(output.shape) if output is not None else None - if verbose: - print(f" ✓ SUCCESS - Output shape: {output_shape}") - return { - 'dataset': dataset_name, - 'model': model_name, - 'split': split, - 'status': 'success', - 'output_shape': output_shape, - 'num_samples': num_samples - } + # Load the pair + if verbose: + print(f"Testing: {dataset_name} + {model_name} ({split} split)") + + result = load_dataset_model_pair( + dataset_name=dataset_name, + model_name=model_name, + split=split, + batch_size=1, + shuffle=False, + auto_download=True, + root_dir=root_dir + ) + + model = result['model'] + dataloader = result['dataloader'] + num_samples = result['num_samples'] + + # Get first sample to test + input_tensor = None + for batch in dataloader: + if isinstance(batch, (tuple, list)) and len(batch) >= 2: + # Standard (image, label) tuple + input_tensor, labels = batch[0], batch[1] + elif isinstance(batch, (tuple, list)) and len(batch) == 1: + input_tensor = batch[0] else: - if verbose: - print(f" ✗ FAILED - {error_msg}") - return { - 'dataset': dataset_name, - 'model': model_name, - 'split': split, - 'status': 'failed', - 'error': error_msg, - 'num_samples': num_samples - } - - except Exception as e: - error_msg = str(e)[:100] + input_tensor = batch + break # Only need first batch + + if input_tensor is None: + error_msg = 'Empty dataloader' if verbose: - print(f" ✗ ERROR - {error_msg}") + print(f" ✗ FAILED - {error_msg}") return { 'dataset': dataset_name, 'model': model_name, 'split': split, - 'status': 'error', - 'error': error_msg + 'status': 'failed', + 'error': error_msg, + 'num_samples': num_samples + } + + # Verify we have a tensor + if not isinstance(input_tensor, torch.Tensor): + error_msg = f'Input type {type(input_tensor).__name__}, expected torch.Tensor' + if verbose: + print(f" ✗ FAILED - {error_msg}") + return { + 'dataset': dataset_name, + 'model': model_name, + 'split': split, + 'status': 'failed', + 'error': error_msg, + 'num_samples': num_samples + } + + # Run inference + success, output, error_msg = infer_single_model(combo_id, model, input_tensor) + + if success: + output_shape = tuple(output.shape) if output is not None else None + if verbose: + print(f" ✓ SUCCESS - Output shape: {output_shape}") + return { + 'dataset': dataset_name, + 'model': model_name, + 'split': split, + 'status': 'success', + 'output_shape': output_shape, + 'num_samples': num_samples + } + else: + if verbose: + print(f" ✗ FAILED - {error_msg}") + return { + 'dataset': dataset_name, + 'model': model_name, + 'split': split, + 'status': 'failed', + 'error': error_msg, + 'num_samples': num_samples } @@ -957,6 +898,7 @@ def main(): print(f" • {pair_key}: {reason}") print(rule()) + assert failed_pairs == 0, f"{failed_pairs}/{total_pairs} dataset-model pair(s) failed inference" if __name__ == "__main__": diff --git a/act/front_end/vnnlib_loader/cli.py b/act/front_end/vnnlib_loader/cli.py index 6781275de..4efd17700 100644 --- a/act/front_end/vnnlib_loader/cli.py +++ b/act/front_end/vnnlib_loader/cli.py @@ -102,6 +102,7 @@ def print_category_detail(category_name: str): except ValueError as e: print(f"Error: {e}") + raise def print_summary_statistics(): @@ -206,6 +207,7 @@ def main(): except ValueError as e: print(f"Error: {e}") print(f"\nAvailable types: {', '.join(get_all_types())}") + raise elif args.info: print_category_detail(args.info) @@ -261,24 +263,25 @@ def main(): force_redownload=args.force ) - if result['status'] == 'success': - print(f"\n{rule(100)}") - print(f"✓ DOWNLOAD SUCCESSFUL") - print(f"{rule(100)}") - print(f"Category: {actual_name}") - print(f"Path: {result['category_path']}") - print(f"Instances: {result['num_instances']}") - if 'num_onnx_models' in result: - print(f"ONNX models: {result['num_onnx_models']}") - if 'num_vnnlib_specs' in result: - print(f"VNNLIB specs: {result['num_vnnlib_specs']}") - print(f"{rule(100)}\n") - else: + if result['status'] != 'success': print(f"\n{rule(100)}") print(f"⚠️ DOWNLOAD FAILED") print(f"{rule(100)}") print(f"Error: {result.get('message', 'Unknown error')}") print(f"{rule(100)}\n") + assert result['status'] == 'success', f"Download failed: {result.get('message')}" + + print(f"\n{rule(100)}") + print(f"✓ DOWNLOAD SUCCESSFUL") + print(f"{rule(100)}") + print(f"Category: {actual_name}") + print(f"Path: {result['category_path']}") + print(f"Instances: {result['num_instances']}") + if 'num_onnx_models' in result: + print(f"ONNX models: {result['num_onnx_models']}") + if 'num_vnnlib_specs' in result: + print(f"VNNLIB specs: {result['num_vnnlib_specs']}") + print(f"{rule(100)}\n") except ValueError as e: print(f"\n{rule(100)}") @@ -287,51 +290,42 @@ def main(): print(f"Error: {e}") print(f"\nUse --list to see available categories") print(f"{rule(100)}\n") - except Exception as e: - print(f"\n{rule(100)}") - print(f"⚠️ DOWNLOAD ERROR") - print(f"{rule(100)}") - print(f"Error: {e}") - print(f"{rule(100)}\n") + raise elif args.list_downloads: # List downloaded categories - try: - downloaded = list_downloaded_pairs() - - if not downloaded: - print(f"\nNo downloaded VNNLIB categories found.") - print(f"Use --download to download a benchmark.") - return - - # Group by category - by_category = {} - for pair in downloaded: - cat = pair['category'] - if cat not in by_category: - by_category[cat] = [] - by_category[cat].append(pair) - - print(f"\n{rule(100)}") - print(f"DOWNLOADED VNNLIB CATEGORIES ({len(by_category)})") - print(f"{rule(100)}") - - for category in sorted(by_category.keys()): - instances = by_category[category] - info = CATEGORY_MAPPING.get(category, {}) - - print(f"\n{category} ({len(instances)} instances)") - if info: - print(f" Type: {info.get('type', 'N/A')}") - print(f" Description: {info.get('description', 'N/A')}") - print(f" Path: {instances[0]['category_path']}") - - print(f"\n{rule(100)}") - print(f"Total instances: {len(downloaded)}") - print(f"{rule(100)}\n") + downloaded = list_downloaded_pairs() + + if not downloaded: + print(f"\nNo downloaded VNNLIB categories found.") + print(f"Use --download to download a benchmark.") + return + + # Group by category + by_category = {} + for pair in downloaded: + cat = pair['category'] + if cat not in by_category: + by_category[cat] = [] + by_category[cat].append(pair) + + print(f"\n{rule(100)}") + print(f"DOWNLOADED VNNLIB CATEGORIES ({len(by_category)})") + print(f"{rule(100)}") + + for category in sorted(by_category.keys()): + instances = by_category[category] + info = CATEGORY_MAPPING.get(category, {}) - except Exception as e: - print(f"Error listing downloads: {e}") + print(f"\n{category} ({len(instances)} instances)") + if info: + print(f" Type: {info.get('type', 'N/A')}") + print(f" Description: {info.get('description', 'N/A')}") + print(f" Path: {instances[0]['category_path']}") + + print(f"\n{rule(100)}") + print(f"Total instances: {len(downloaded)}") + print(f"{rule(100)}\n") else: parser.print_help() diff --git a/act/front_end/vnnlib_loader/create_specs.py b/act/front_end/vnnlib_loader/create_specs.py index 4abdf8bfe..9c7cd3b12 100644 --- a/act/front_end/vnnlib_loader/create_specs.py +++ b/act/front_end/vnnlib_loader/create_specs.py @@ -14,7 +14,7 @@ from __future__ import annotations from pathlib import Path -from typing import Any, List, Tuple, Dict, Optional +from typing import Any, List, Tuple, Dict, Optional, override import logging import torch.nn as nn @@ -57,7 +57,7 @@ class VNNLibSpecCreator(BaseSpecCreator): def __init__( self, config_name: Optional[str] = "vnnlib_default", - config_dict: Optional[Dict] = None + config_dict: Optional[Dict[str, Any]] = None ): """ Initialize VNNLIB spec creator. @@ -68,6 +68,7 @@ def __init__( """ super().__init__(config_name, config_dict) + @override def create_specs_for_data_model_pairs( self, categories: Optional[List[str]] = None, @@ -159,6 +160,7 @@ def create_specs_for_data_model_pairs( # different run. _onnx_cache: Dict[Tuple[str, str, Optional[str]], Dict[str, Any]] = {} + failures: List[str] = [] for instance_info in all_instances: category = instance_info['category'] onnx_model = instance_info['onnx_model'] @@ -193,14 +195,21 @@ def create_specs_for_data_model_pairs( validate_shapes=validate_shapes ) - if result is not None: - results.append(result) + results.append(result) - except Exception as e: + except UnsupportedSpecError: + raise + except Exception as e: # noqa: BLE001 — collect every failure, raise below logger.error( f"Failed to create specs for {category}/{instance_id}: {e}" ) + failures.append(f"{category}/{instance_id}: {e}") + if failures: + raise RuntimeError( + f"Spec creation failed for {len(failures)}/{len(all_instances)} VNNLIB instance(s):\n " + + "\n ".join(failures) + ) logger.info(f"Successfully created specs for {len(results)} instances") return results @@ -208,15 +217,14 @@ def _create_specs_for_single_instance( self, category: str, instance_id: str, - instance_data: Dict, + instance_data: Dict[str, Any], validate_shapes: bool - ) -> Optional[Tuple[str, str, nn.Module, List[LabeledInputTensor], List[Tuple[InputSpec, OutputSpec]]]]: + ) -> Tuple[str, str, nn.Module, List[LabeledInputTensor], List[Tuple[InputSpec, OutputSpec]]]: """ Create specs for a single VNNLIB instance. Returns: Tuple of (category, instance_id, pytorch_model, labeled_tensors, spec_pairs) - or None if failed """ logger.info(f"Generating specs for {category}/{instance_id}") @@ -227,23 +235,16 @@ def _create_specs_for_single_instance( # Parse VNNLIB to create specs # Pass input_shape to ensure specs match tensor shape (not flattened) # Pass true_label to promote RANGE to TOP1_ROBUST for classification - try: - queries = parse_vnnlib_queries( - vnnlib_path, - labeled_tensor=labeled_tensor, - ) - if not queries: - raise RuntimeError("parse_vnnlib_queries returned no queries") - logger.info( - f"Parsed VNNLIB specs: {len(queries)} queries, " - f"first kind=({queries[0][0].kind}, {queries[0][1].kind})" - ) - except UnsupportedSpecError as e: - logger.warning("UNSUPPORTED spec %s: %s", instance_id, e) - return None - except Exception as e: - logger.error(f"Failed to parse VNNLIB specs: {e}") - return None + queries = parse_vnnlib_queries( + vnnlib_path, + labeled_tensor=labeled_tensor, + ) + if not queries: + raise RuntimeError(f"parse_vnnlib_queries returned no queries for {vnnlib_path}") + logger.info( + f"Parsed VNNLIB specs: {len(queries)} queries, " + f"first kind=({queries[0][0].kind}, {queries[0][1].kind})" + ) spec_pairs = list(queries) @@ -256,8 +257,9 @@ def _create_specs_for_single_instance( ) if not validated_pairs: - logger.warning(f"Spec validation failed for {category}/{instance_id}") - return None + raise RuntimeError( + f"No spec pair of {category}/{instance_id} matches the model's I/O shapes" + ) spec_pairs = validated_pairs @@ -318,6 +320,4 @@ def create_specs_from_paths(onnx_path, vnnlib_path, category: str = "custom"): } sr = VNNLibSpecCreator()._create_specs_for_single_instance( category, f"{onnx_p.stem}__{vnnlib_p.stem}", instance_data, validate_shapes=True) - if sr is None: - raise SystemExit(f"Spec creation failed (unsupported/invalid spec) for {vnnlib_p.name}") return sr diff --git a/act/front_end/vnnlib_loader/data_model_loader.py b/act/front_end/vnnlib_loader/data_model_loader.py index 411356cfa..02008003b 100644 --- a/act/front_end/vnnlib_loader/data_model_loader.py +++ b/act/front_end/vnnlib_loader/data_model_loader.py @@ -181,121 +181,112 @@ def download_vnnlib_category( 'message': f"Failed to download instances.csv from all sources. Tried {len(VNNCOMP_REPO_URLS)} repositories." } - try: - # Parse instances.csv to get ONNX and VNNLIB files - instances: List[Dict[str, Any]] = [] - onnx_files: set[str] = set() - vnnlib_files: set[str] = set() - with open(instances_file, 'r') as f: - reader = csv.reader(f) - rows = list(reader) - if rows and len(rows[0]) >= 2 and rows[0][0].strip().lower() == 'onnx': - rows = rows[1:] - - for row in rows: - if len(row) >= 3: - onnx_field, vnnlib_field, timeout = row[0], row[1], row[2] - onnx_entries = _parse_onnx_field(onnx_field) - vnnlib_file = _normalize_rel_path(vnnlib_field) - for _, onnx_file in onnx_entries: - onnx_files.add(onnx_file) - vnnlib_files.add(vnnlib_file) - instances.append({ - 'onnx': onnx_entries, - 'vnnlib': vnnlib_file, - 'timeout': timeout - }) + # Parse instances.csv to get ONNX and VNNLIB files + instances: List[Dict[str, Any]] = [] + onnx_files: set[str] = set() + vnnlib_files: set[str] = set() + with open(instances_file, 'r') as f: + reader = csv.reader(f) + rows = list(reader) + if rows and len(rows[0]) >= 2 and rows[0][0].strip().lower() == 'onnx': + rows = rows[1:] - logger.info(f"Found {len(instances)} instances in category '{category}'") + for row in rows: + if len(row) >= 3: + onnx_field, vnnlib_field, timeout = row[0], row[1], row[2] + onnx_entries = _parse_onnx_field(onnx_field) + vnnlib_file = _normalize_rel_path(vnnlib_field) + for _, onnx_file in onnx_entries: + onnx_files.add(onnx_file) + vnnlib_files.add(vnnlib_file) + instances.append({ + 'onnx': onnx_entries, + 'vnnlib': vnnlib_file, + 'timeout': timeout + }) + + logger.info(f"Found {len(instances)} instances in category '{category}'") + + # Download ONNX and VNNLIB files using the successful base URL + downloaded_onnx: set[str] = set() + downloaded_vnnlib: set[str] = set() + all_files: List[Tuple[str, set[str]]] = [(rel_path, downloaded_onnx) for rel_path in sorted(onnx_files)] + all_files.extend((rel_path, downloaded_vnnlib) for rel_path in sorted(vnnlib_files)) + + for idx, (rel_path, downloaded_set) in enumerate(all_files, 1): + target_path = category_dir / rel_path + target_path.parent.mkdir(parents=True, exist_ok=True) - # Download ONNX and VNNLIB files using the successful base URL - downloaded_onnx: set[str] = set() - downloaded_vnnlib: set[str] = set() - all_files: List[Tuple[str, set[str]]] = [(rel_path, downloaded_onnx) for rel_path in sorted(onnx_files)] - all_files.extend((rel_path, downloaded_vnnlib) for rel_path in sorted(vnnlib_files)) + # Try compressed first, then uncompressed. + tried_urls = [] + success = False - for idx, (rel_path, downloaded_set) in enumerate(all_files, 1): - target_path = category_dir / rel_path - target_path.parent.mkdir(parents=True, exist_ok=True) - - # Try compressed first, then uncompressed. - tried_urls = [] - success = False + for extension in ['.gz', '']: + file_url = f"{successful_base_url}/{repo_prefix}/{rel_path}{extension}" + tried_urls.append(file_url) - for extension in ['.gz', '']: - file_url = f"{successful_base_url}/{repo_prefix}/{rel_path}{extension}" - tried_urls.append(file_url) + try: + logger.debug(f"[{idx}/{len(all_files)}] Trying {rel_path}{extension}") - try: - logger.debug(f"[{idx}/{len(all_files)}] Trying {rel_path}{extension}") + if extension == '.gz': + # Download compressed file + gz_path = target_path.parent / f"{target_path.name}.gz" + urllib.request.urlretrieve(file_url, gz_path) - if extension == '.gz': - # Download compressed file - gz_path = target_path.parent / f"{target_path.name}.gz" - urllib.request.urlretrieve(file_url, gz_path) - - # Decompress - with gzip.open(gz_path, 'rb') as f_in: - with open(target_path, 'wb') as f_out: - shutil.copyfileobj(f_in, f_out) - - # Remove .gz file - gz_path.unlink() - logger.debug(f" ✓ Decompressed {rel_path}") - else: - # Download uncompressed file - urllib.request.urlretrieve(file_url, target_path) + # Decompress + with gzip.open(gz_path, 'rb') as f_in: + with open(target_path, 'wb') as f_out: + shutil.copyfileobj(f_in, f_out) - downloaded_set.add(rel_path) - success = True - break - except Exception as e: - logger.debug(f" Failed with {extension}: {e}") - continue - - if not success: - logger.warning(f"Failed to download {rel_path} from any URL") - - # Create metadata - metadata = { - 'category': category, - 'repo_name': repo_name, - 'version': version, - 'num_instances': len(instances), - 'num_onnx_models': len(downloaded_onnx), - 'num_vnnlib_specs': len(downloaded_vnnlib), - 'source': successful_base_url, - 'paths': { - 'onnx': str(onnx_dir), - 'vnnlib': str(vnnlib_dir), - 'instances': str(instances_file) - } - } - - # Save metadata - metadata_path = category_dir / "info.json" - with open(metadata_path, 'w') as f: - json.dump(metadata, f, indent=2) - - logger.info( - f"Successfully downloaded category '{category}': " - f"{len(instances)} instances, {len(downloaded_onnx)} ONNX models, " - f"{len(downloaded_vnnlib)} VNNLIB specs" - ) - - return { - 'status': 'success', - 'message': f"Downloaded {len(instances)} instances", - 'category_path': str(category_dir), - 'num_instances': len(instances) - } + # Remove .gz file + gz_path.unlink() + logger.debug(f" ✓ Decompressed {rel_path}") + else: + # Download uncompressed file + urllib.request.urlretrieve(file_url, target_path) + + downloaded_set.add(rel_path) + success = True + break + except Exception as e: + logger.debug(f" Failed with {extension}: {e}") + continue - except Exception as e: - logger.error(f"Failed to download category '{category}': {e}") - return { - 'status': 'error', - 'message': str(e) + assert success, f"Failed to download {rel_path} from any of {tried_urls}" + + # Create metadata + metadata = { + 'category': category, + 'repo_name': repo_name, + 'version': version, + 'num_instances': len(instances), + 'num_onnx_models': len(downloaded_onnx), + 'num_vnnlib_specs': len(downloaded_vnnlib), + 'source': successful_base_url, + 'paths': { + 'onnx': str(onnx_dir), + 'vnnlib': str(vnnlib_dir), + 'instances': str(instances_file) } + } + + # Save metadata + metadata_path = category_dir / "info.json" + with open(metadata_path, 'w') as f: + json.dump(metadata, f, indent=2) + + logger.info( + f"Successfully downloaded category '{category}': " + f"{len(instances)} instances, {len(downloaded_onnx)} ONNX models, " + f"{len(downloaded_vnnlib)} VNNLIB specs" + ) + + return { + 'status': 'success', + 'message': f"Downloaded {len(instances)} instances", + 'category_path': str(category_dir), + 'num_instances': len(instances) + } def list_downloaded_pairs(root_dir: Optional[str] = None) -> List[Dict[str, Any]]: @@ -338,44 +329,40 @@ def list_downloaded_pairs(root_dir: Optional[str] = None) -> List[Dict[str, Any] if not instances_file.exists(): continue - try: - with open(instances_file, 'r') as f: - reader = csv.reader(f) - rows = list(reader) - if rows and len(rows[0]) >= 2 and rows[0][0].strip().lower() == 'onnx': - rows = rows[1:] - - for row in rows: - if len(row) >= 3: - onnx_field, vnnlib_file, timeout = row[0], row[1], row[2] - onnx_entries = _parse_onnx_field(onnx_field) - vnnlib_rel = _normalize_rel_path(vnnlib_file) - - onnx_paths = [category_dir / onnx_rel for _, onnx_rel in onnx_entries] - vnnlib_path = category_dir / vnnlib_rel + with open(instances_file, 'r') as f: + reader = csv.reader(f) + rows = list(reader) + if rows and len(rows[0]) >= 2 and rows[0][0].strip().lower() == 'onnx': + rows = rows[1:] + + for row in rows: + if len(row) >= 3: + onnx_field, vnnlib_file, timeout = row[0], row[1], row[2] + onnx_entries = _parse_onnx_field(onnx_field) + vnnlib_rel = _normalize_rel_path(vnnlib_file) + + onnx_paths = [category_dir / onnx_rel for _, onnx_rel in onnx_entries] + vnnlib_path = category_dir / vnnlib_rel + + # Only include if all referenced files exist + if all(onnx_path.exists() for onnx_path in onnx_paths) and vnnlib_path.exists(): + first_onnx_rel = onnx_entries[0][1] + paths: Dict[str, Any] = { + 'onnx': str(onnx_paths[0]), + 'vnnlib': str(vnnlib_path) + } + if len(onnx_entries) > 1: + paths['onnx_models'] = [str(onnx_path) for onnx_path in onnx_paths] - # Only include if all referenced files exist - if all(onnx_path.exists() for onnx_path in onnx_paths) and vnnlib_path.exists(): - first_onnx_rel = onnx_entries[0][1] - paths: Dict[str, Any] = { - 'onnx': str(onnx_paths[0]), - 'vnnlib': str(vnnlib_path) - } - if len(onnx_entries) > 1: - paths['onnx_models'] = [str(onnx_path) for onnx_path in onnx_paths] - - all_instances.append({ - 'category': category_dir.name, - 'onnx_model': first_onnx_rel, - 'onnx_models': onnx_entries, - 'vnnlib_spec': vnnlib_rel, - 'timeout': float(timeout) if timeout else None, - 'is_dual_model': len(onnx_entries) > 1, - 'paths': paths - }) - - except Exception as e: - logger.warning(f"Failed to read instances from {category_dir.name}: {e}") + all_instances.append({ + 'category': category_dir.name, + 'onnx_model': first_onnx_rel, + 'onnx_models': onnx_entries, + 'vnnlib_spec': vnnlib_rel, + 'timeout': float(timeout) if timeout else None, + 'is_dual_model': len(onnx_entries) > 1, + 'paths': paths + }) return all_instances @@ -670,62 +657,49 @@ def model_inference_with_vnnlib( model_name = Path(onnx_model).stem combo_id = f"{category}+{model_name}" - try: - # Load the pair + # Load the pair + if verbose: + logger.info(f"Testing: {category} + {model_name}") + + result = load_vnnlib_pair( + category=category, + onnx_model=onnx_model, + vnnlib_spec=vnnlib_spec, + root_dir=root_dir, + auto_download=True + ) + + model = result['model'] + labeled_tensor = result['labeled_tensor'] + input_tensor = labeled_tensor.tensor # Already has batch dimension + + # Run inference + success, output, error_msg = infer_single_model( + combo_id, + model, + input_tensor # Already (1, C, H, W) + ) + + if success: + if output is None: + raise RuntimeError("Inference succeeded without an output tensor") if verbose: - logger.info(f"Testing: {category} + {model_name}") - - result = load_vnnlib_pair( - category=category, - onnx_model=onnx_model, - vnnlib_spec=vnnlib_spec, - root_dir=root_dir, - auto_download=True - ) - - model = result['model'] - labeled_tensor = result['labeled_tensor'] - input_tensor = labeled_tensor.tensor # Already has batch dimension - - # Run inference - success, output, error_msg = infer_single_model( - combo_id, - model, - input_tensor # Already (1, C, H, W) - ) - - if success: - if output is None: - raise RuntimeError("Inference succeeded without an output tensor") - if verbose: - logger.info(f" ✓ SUCCESS - Output shape: {output.shape}") - return { - 'category': category, - 'model': model_name, - 'spec': vnnlib_spec, - 'status': 'success', - 'output_shape': tuple(output.shape) - } - else: - if verbose: - logger.warning(f" ✗ FAILED - {error_msg}") - return { - 'category': category, - 'model': model_name, - 'spec': vnnlib_spec, - 'status': 'failed', - 'error': error_msg - } - - except Exception as e: - error_msg = str(e)[:100] + logger.info(f" ✓ SUCCESS - Output shape: {output.shape}") + return { + 'category': category, + 'model': model_name, + 'spec': vnnlib_spec, + 'status': 'success', + 'output_shape': tuple(output.shape) + } + else: if verbose: - logger.error(f" ✗ ERROR - {error_msg}") + logger.warning(f" ✗ FAILED - {error_msg}") return { 'category': category, 'model': model_name, 'spec': vnnlib_spec, - 'status': 'error', + 'status': 'failed', 'error': error_msg } @@ -804,9 +778,5 @@ def get_category_info(category: str, root_dir: Optional[str] = None) -> Optional if not info_path.exists(): return None - try: - with open(info_path, 'r') as f: - return json.load(f) - except Exception as e: - logger.error(f"Failed to read category info: {e}") - return None + with open(info_path, 'r') as f: + return json.load(f) diff --git a/act/front_end/vnnlib_loader/onnx_converter.py b/act/front_end/vnnlib_loader/onnx_converter.py index aba299106..2d6a530ca 100644 --- a/act/front_end/vnnlib_loader/onnx_converter.py +++ b/act/front_end/vnnlib_loader/onnx_converter.py @@ -14,7 +14,7 @@ from __future__ import annotations from pathlib import Path -from typing import Tuple, Optional +from typing import Tuple import logging import torch import torch.nn as nn @@ -458,46 +458,41 @@ def get_onnx_metadata(onnx_path: Path) -> dict: Returns: Dict with model metadata (producer, version, shapes, etc.) """ - try: - import onnx - - onnx_model = onnx.load(str(onnx_path)) - - metadata = { - 'producer_name': onnx_model.producer_name, - 'producer_version': onnx_model.producer_version, - 'ir_version': onnx_model.ir_version, - 'opset_version': None, - 'input_shapes': [], - 'output_shapes': [] - } - - # Get opset version - if onnx_model.opset_import: - metadata['opset_version'] = onnx_model.opset_import[0].version - - # Get input/output shapes - graph = onnx_model.graph - - for inp in graph.input: - shape = _extract_shape_from_tensor(inp) - metadata['input_shapes'].append({ - 'name': inp.name, - 'shape': shape - }) - - for out in graph.output: - shape = _extract_shape_from_tensor(out) - metadata['output_shapes'].append({ - 'name': out.name, - 'shape': shape - }) - - return metadata - - except Exception as e: - logger.error(f"Failed to extract ONNX metadata: {e}") - return {} + import onnx + + onnx_model = onnx.load(str(onnx_path)) + + metadata = { + 'producer_name': onnx_model.producer_name, + 'producer_version': onnx_model.producer_version, + 'ir_version': onnx_model.ir_version, + 'opset_version': None, + 'input_shapes': [], + 'output_shapes': [] + } + + # Get opset version + if onnx_model.opset_import: + metadata['opset_version'] = onnx_model.opset_import[0].version + + # Get input/output shapes + graph = onnx_model.graph + + for inp in graph.input: + shape = _extract_shape_from_tensor(inp) + metadata['input_shapes'].append({ + 'name': inp.name, + 'shape': shape + }) + + for out in graph.output: + shape = _extract_shape_from_tensor(out) + metadata['output_shapes'].append({ + 'name': out.name, + 'shape': shape + }) + + return metadata def validate_onnx_file(onnx_path: Path) -> bool: @@ -508,53 +503,18 @@ def validate_onnx_file(onnx_path: Path) -> bool: onnx_path: Path to .onnx file Returns: - True if valid, False otherwise + True if the model passes the ONNX checker, False if it is malformed. + Load errors (missing file, unreadable protobuf) are raised. """ + import onnx + + onnx_model = onnx.load(str(onnx_path)) try: - import onnx - - onnx_model = onnx.load(str(onnx_path)) onnx.checker.check_model(onnx_model) - logger.info(f"ONNX model validated: {onnx_path.name}") - return True - - except Exception as e: - logger.error(f"ONNX validation failed: {e}") + except onnx.checker.ValidationError as e: + logger.warning(f"ONNX validation failed for {onnx_path.name}: {e}") return False + logger.info(f"ONNX model validated: {onnx_path.name}") + return True -def convert_and_save_pytorch( - onnx_path: Path, - output_path: Optional[Path] = None, - simplify: bool = True -) -> Path: - """ - Convert ONNX model to PyTorch and save as .pt file. - - Args: - onnx_path: Path to .onnx file - output_path: Path for .pt file (defaults to same dir as ONNX) - simplify: Whether to simplify ONNX before conversion - - Returns: - Path to saved .pt file - - Raises: - ONNXConversionError: If conversion or saving fails - """ - try: - # Convert to PyTorch - pytorch_model = convert_onnx_to_pytorch(onnx_path, simplify=simplify) - - # Determine output path - if output_path is None: - output_path = onnx_path.with_suffix('.pt') - - # Save model - torch.save(pytorch_model.state_dict(), output_path) - logger.info(f"Saved PyTorch model to {output_path}") - - return output_path - - except Exception as e: - raise ONNXConversionError(f"Failed to convert and save: {str(e)}") diff --git a/act/front_end/vnnlib_loader/vnnlib_parser.py b/act/front_end/vnnlib_loader/vnnlib_parser.py index b7f39abc6..34f21e860 100644 --- a/act/front_end/vnnlib_loader/vnnlib_parser.py +++ b/act/front_end/vnnlib_loader/vnnlib_parser.py @@ -32,11 +32,89 @@ class VNNLibParseError(Exception): pass -class UnsupportedSpecError(Exception): +class UnsupportedSpecError(ValueError): """Exception raised for soundly unsupported VNNLIB features.""" pass +_DECLARE_NETWORK_RE = re.compile(r"\(\s*declare-network\s+([A-Za-z_]\w*)\b") +_DECLARE_INPUT_RE = re.compile( + r"\(\s*declare-input\s+([A-Za-z_]\w*)\s+\S+\s+\[([^\]]+)\]\s*\)", + re.MULTILINE, +) +_ASSERT_RE = re.compile(r"\(\s*assert\b") + + +def _declaration_prefix(vnnlib_path: Path) -> str: + """Read only the declarations preceding the first VNNLIB assertion.""" + lines: List[str] = [] + try: + with open(vnnlib_path, "r") as source: + for line in source: + assert_match = _ASSERT_RE.search(line) + lines.append(line if assert_match is None else line[:assert_match.start()]) + if assert_match is not None: + break + except OSError as exc: + raise VNNLibParseError(f"Failed to read {vnnlib_path}: {exc}") from exc + return "".join(lines) + + +def _network_declaration_blocks(prefix: str) -> List[Tuple[str, str]]: + """Return complete ``declare-network`` forms from a declaration prefix.""" + uncommented = re.sub(r";[^\n]*", "", prefix) + blocks: List[Tuple[str, str]] = [] + for match in _DECLARE_NETWORK_RE.finditer(uncommented): + depth = 0 + for index in range(match.start(), len(uncommented)): + char = uncommented[index] + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + blocks.append((match.group(1), uncommented[match.start():index + 1])) + break + return blocks + + +def reject_multi_input_vnnlib(vnnlib_path: Path) -> None: + """Reject VNNLIB 2.0 networks declaring multiple input tensors. + + The scan stops at the first assertion, so even very large properties are + rejected without tokenizing or constructing their assertion AST. Separate + one-input ``declare-network`` blocks remain supported for isomorphic specs. + Legacy VNNLIB 1.0 ``declare-const X_i`` declarations denote one flat input + tensor and therefore do not trigger this check. + """ + prefix = _declaration_prefix(vnnlib_path) + network_blocks = _network_declaration_blocks(prefix) + if not network_blocks: + network_blocks = [("", prefix)] + + for network_name, block in network_blocks: + declarations = [] + for match in _DECLARE_INPUT_RE.finditer(block): + raw_shape = match.group(2) + try: + shape = tuple(int(part.strip()) for part in raw_shape.split(",") if part.strip()) + except ValueError as exc: + raise VNNLibParseError( + f"Invalid VNNLIB 2.0 declare-input shape: {raw_shape}" + ) from exc + declarations.append((match.group(1), shape)) + + if len(declarations) > 1: + rendered = ", ".join( + f"{name} [{', '.join(str(dim) for dim in shape)}]" + for name, shape in declarations + ) + raise UnsupportedSpecError( + f"{vnnlib_path}: multiple input tensors are unsupported; " + f"network {network_name!r} declares {len(declarations)} inputs: {rendered}" + ) + + # ------------------------------------------------------------------------- # Public API # ------------------------------------------------------------------------- @@ -67,6 +145,8 @@ def parse_vnnlib_to_tensors( """ if not vnnlib_path.exists(): raise VNNLibParseError(f"VNNLIB file not found: {vnnlib_path}") + + reject_multi_input_vnnlib(vnnlib_path) try: with open(vnnlib_path, 'r') as f: @@ -96,15 +176,12 @@ def parse_vnnlib_to_tensors( # Constraints inside (or ...) branches must not be intersected here — # that would produce empty boxes for disjunctive-input properties # (e.g. ACAS Xu prop_5..10). - try: - forms = _parse_all_forms(content) - simple_bodies = [ - f[1] for f in forms - if isinstance(f, list) and len(f) >= 2 and f[0] == "assert" - and _is_simple_x_bound(f[1]) - ] - except Exception: - simple_bodies = [] + forms = _parse_all_forms(content) + simple_bodies = [ + f[1] for f in forms + if isinstance(f, list) and len(f) >= 2 and f[0] == "assert" + and _is_simple_x_bound(f[1]) + ] input_bounds = _extract_input_bounds(simple_bodies, num_inputs) # Create input tensor from bounds center @@ -238,23 +315,19 @@ def extract_label_from_vnnlib(vnnlib_path: Path) -> Optional[int]: >>> print(label) 14 """ - try: - with open(vnnlib_path, 'r') as f: - # Read first few lines (label is typically in first comment) - for _ in range(5): - line = f.readline() - if not line: - break - - # Match patterns like: ; CIFAR100 property with label: 14. - match = re.search(r'label:\s*(\d+)', line, re.IGNORECASE) - if match: - return int(match.group(1)) - - return None - except Exception as e: - logger.debug(f"Failed to extract label from {vnnlib_path}: {e}") - return None + with open(vnnlib_path, 'r') as f: + # Read first few lines (label is typically in first comment) + for _ in range(5): + line = f.readline() + if not line: + break + + # Match patterns like: ; CIFAR100 property with label: 14. + match = re.search(r'label:\s*(\d+)', line, re.IGNORECASE) + if match: + return int(match.group(1)) + + return None # ------------------------------------------------------------------------- @@ -280,6 +353,7 @@ def parse_vnnlib_2_0( """Parse VNNLIB 2.0 by ravel-rewriting bracket variables to legacy names.""" if not vnnlib_path.exists(): raise VNNLibParseError(f"VNNLIB file not found: {vnnlib_path}") + reject_multi_input_vnnlib(vnnlib_path) try: with open(vnnlib_path, 'r') as f: content = f.read() @@ -597,16 +671,12 @@ def _extract_input_bounds( continue x_tok = left if x_is_left else right lit_tok = right if x_is_left else left - try: - val = float(lit_tok) - except (TypeError, ValueError): - continue + val = float(lit_tok) # _is_simple_x_bound admitted only numeric literals x_match = _X_RE.fullmatch(x_tok) if x_match is None: continue idx = int(x_match.group(1)) - if idx >= num_inputs: - continue + assert idx < num_inputs, f"X_{idx} is bounded but the property declares {num_inputs} inputs" # (<= X val) -> X <= val (upper) # (>= X val) -> X >= val (lower) # (<= val X) -> val <= X -> X >= val (lower) diff --git a/act/pipeline/__init__.py b/act/pipeline/__init__.py index 43a81b077..27298e302 100644 --- a/act/pipeline/__init__.py +++ b/act/pipeline/__init__.py @@ -22,7 +22,6 @@ - ModelFactory: Create PyTorch models from YAML configurations - TorchToACT: Convert PyTorch models to ACT representation - VerifierValidator: Validate verifier correctness with concrete tests - - PerformanceProfiler: Profile execution time and memory usage All verification utilities are located in the verification/ submodule. @@ -42,40 +41,8 @@ from act.pipeline.verification.model_factory import ModelFactory from act.pipeline.verification.torch2act import TorchToACT -# Import utilities -try: - from act.pipeline.verification.utils import ( - PerformanceProfiler, - ParallelExecutor, - print_memory_usage, - clear_torch_cache, - setup_logging, - ProgressTracker, - ) - UTILS_AVAILABLE = True -except ImportError: - UTILS_AVAILABLE = False - PerformanceProfiler = None - ParallelExecutor = None - print_memory_usage = None - clear_torch_cache = None - setup_logging = None - ProgressTracker = None - - __all__ = [ # Core model factory and conversion 'ModelFactory', 'TorchToACT', - - # Utilities - 'PerformanceProfiler', - 'ParallelExecutor', - 'print_memory_usage', - 'clear_torch_cache', - 'setup_logging', - 'ProgressTracker', - - # Availability flags - 'UTILS_AVAILABLE', ] diff --git a/act/pipeline/fuzzing/actfuzzer.py b/act/pipeline/fuzzing/actfuzzer.py index deca71594..547660801 100644 --- a/act/pipeline/fuzzing/actfuzzer.py +++ b/act/pipeline/fuzzing/actfuzzer.py @@ -549,16 +549,10 @@ def _generate_report(self) -> FuzzingReport: total_time = time.time() - self.start_time # Neurons that were never activated across all iterations - never_activated_neurons: List[Tuple[str, int]] = [] - never_activated_count = 0 - try: - uncovered = self.coverage_tracker.get_uncovered_neurons() - never_activated_count = len(uncovered) - # Deterministic small sample for logs/report - never_activated_neurons = sorted(list(uncovered))[:20] - except Exception: - never_activated_count = 0 - never_activated_neurons = [] + uncovered = self.coverage_tracker.get_uncovered_neurons() + never_activated_count = len(uncovered) + # Deterministic small sample for logs/report + never_activated_neurons: List[Tuple[str, int]] = sorted(list(uncovered))[:20] report = FuzzingReport( total_iterations=self.iterations, diff --git a/act/pipeline/verification/model_factory.py b/act/pipeline/verification/model_factory.py index 1cf65bc33..fe4a05e9c 100644 --- a/act/pipeline/verification/model_factory.py +++ b/act/pipeline/verification/model_factory.py @@ -67,10 +67,7 @@ def _discover_net_names(nets_dir: Path, manifest_path: Optional[Path]) -> List[s break if manifest_path is not None and manifest_path.exists(): - try: - names.extend(str(n) for n in _load_manifest(manifest_path)) - except Exception as e: - logger.warning(f"Failed to read manifest {manifest_path}: {e}") + names.extend(str(n) for n in _load_manifest(manifest_path)) names.extend(p.stem for p in sorted(nets_dir.glob('*.json'))) @@ -123,26 +120,25 @@ def _load_all_nets(self) -> None: - Enables O(1) lookup via get_act_net() - Costs ~10-20MB memory for typical test suites """ + # Callers that only need VNNLIB/TorchVision models construct the factory + # before nets are generated; consumers that need nets check list_networks(). if not self.nets_dir.exists(): - logger.warning(f"Nets dir not found: {self.nets_dir}") + logger.debug(f"Nets dir not found: {self.nets_dir}") return for name in self.net_names: net_path = self.nets_dir / f"{name}.json" if not net_path.exists(): - logger.warning(f"ACT Net file not found: {net_path}. Skipping '{name}'.") - continue + raise FileNotFoundError( + f"ACT Net '{name}' is listed in the manifest but {net_path} does not exist" + ) - try: - with open(net_path, 'r') as f: - net_dict = json.load(f) - act_net = NetSerializer.deserialize_net(net_dict) - self.nets[name] = act_net - logger.debug(f"Pre-loaded ACT Net '{name}' from {net_path}") - except Exception as e: - logger.error(f"Failed to load ACT Net '{name}' from {net_path}: {e}") - continue + with open(net_path, 'r') as f: + net_dict = json.load(f) + act_net = NetSerializer.deserialize_net(net_dict) + self.nets[name] = act_net + logger.debug(f"Pre-loaded ACT Net '{name}' from {net_path}") logger.info(f"Pre-loaded {len(self.nets)} ACT Nets from {self.nets_dir}") @@ -342,11 +338,17 @@ def main(): print("PyTorch Model Factory - Spec-Free Verification Testing") print(rule()) - all_passed = True + failures: List[str] = [] total_tests = 0 passed_tests = 0 - for name in factory.list_networks(): + network_names = factory.list_networks() + if not network_names: + raise RuntimeError( + f"No ACT Nets found in {factory.nets_dir}; run `python -m act.back_end --generate` first" + ) + + for name in network_names: print(f"\n{rule()}") print(f"Network: {name}") print(rule()) @@ -395,13 +397,15 @@ def main(): # Track test success total_tests += 1 + # Constraint outcomes are properties of the (random-weight) + # network on this input, not conversion errors. if input_satisfied and output_satisfied: passed_tests += 1 - print(f" ✅ Test PASSED (both constraints satisfied)") + print(f" ✅ Both constraints satisfied") elif not input_satisfied: - print(f" ⚠️ Test UNCERTAIN (input constraint violated)") + print(f" ℹ️ Input outside spec region (property not applicable)") else: - print(f" ❌ Test FAILED (output constraint violated)") + print(f" ℹ️ Output property violated on this input (counterexample)") else: # Legacy nn.Module (no verification) @@ -415,35 +419,31 @@ def main(): print(f" ❌ Test case '{test_case}' failed: {e}") import traceback traceback.print_exc() - all_passed = False + failures.append(f"{name}/{test_case}: {e}") except Exception as e: print(f"\n❌ Failed to create/test model '{name}': {e}") import traceback traceback.print_exc() - all_passed = False + failures.append(f"{name}: {e}") # Print summary print("\n" + rule()) print(f"📊 Verification Test Summary:") print(f" Total tests: {total_tests}") - print(f" ✅ Passed: {passed_tests}") - print(f" ⚠️ Uncertain/Failed: {total_tests - passed_tests}") - if total_tests > 0: - success_rate = (passed_tests / total_tests) * 100 - print(f" Success rate: {success_rate:.1f}%") + print(f" ✅ Both constraints satisfied: {passed_tests}") + print(f" ℹ️ Input outside spec / property violated: {total_tests - passed_tests}") + if total_tests == 0: + raise RuntimeError("No ACT→PyTorch test cases were executed") print(rule()) - if all_passed: - print("✅ All models created and tested successfully") - else: - print("⚠️ Some models had issues - see details above") + if failures: + raise RuntimeError( + f"{len(failures)} ACT→PyTorch model/test-case error(s):\n " + "\n ".join(failures) + ) + print("✅ All models created and tested successfully") print(rule()) - if not all_passed: - import sys - sys.exit(1) - if __name__ == "__main__": main() diff --git a/act/pipeline/verification/torch2act.py b/act/pipeline/verification/torch2act.py index 728ed0889..67b65fa89 100644 --- a/act/pipeline/verification/torch2act.py +++ b/act/pipeline/verification/torch2act.py @@ -69,7 +69,7 @@ from act.back_end.layer_util import create_layer from act.back_end.utils import avgpool2d_output_hw from act.pipeline.verification.utils import ( - _prod, _normalize_tuple, _assert_dag, _broadcast_const_to_size, + _prod, _normalize_tuple, _assert_dag, _assert_dataflow, _broadcast_const_to_size, _normalize_axes, _reduce_output_shape, _compute_slice_output_shape, ONNX_HANDLERS, ) @@ -77,8 +77,6 @@ # Imports needed for main() test harness from act.util.model_inference import model_inference from act.front_end.model_synthesis import model_synthesis -from act.back_end.solver.solver_torchlp import TorchLPSolver -from act.back_end.solver.solver_gurobi import GurobiSolver from act.util.options import PerformanceOptions from act.util.format_utils import rule @@ -148,6 +146,12 @@ def __init__( # would mis-wire them. _build_preds_succs honours this override. self._fx_pred_override: Dict[int, List[int]] = {} + # Head layer of each node converted into a chain of layers (BatchNorm + # -> SCALE + BIAS). node_to_layer_id maps a node to its LAST layer so + # successors read its output; incoming FX edges must enter the head, + # otherwise the layers in between are bypassed. See _record_chain_head. + self._node_first_layer_id: Dict[str, int] = {} + # Compile-time constants (e.g. OnnxShape values) that must NOT enter # the runtime IR. Resolved by ``_resolve_constant_tensor`` BEFORE any # placeholder/sample_input fallback — this is what stops shape chains @@ -898,11 +902,8 @@ def _evaluate_constant_subgraph(self, node_name: str) -> Optional[torch.Tensor]: arg_vals.append(v) else: arg_vals.append(a) - try: - with torch.no_grad(): - return sub(*arg_vals) - except Exception: - return None + with torch.no_grad(): + return sub(*arg_vals) def _ensure_constant_vars(self, node_name: str) -> bool: """Emit a CONSTANT layer for an ONNX initializer (consumed by Concat / @@ -972,6 +973,7 @@ def _process_fx_graph(self) -> None: if self.fx_graph is None: return for node in self.fx_graph.nodes: + first_layer_id = len(self.layers) if node.op == 'placeholder': pass # Already pre-registered in _pre_register_nodes elif node.op == 'call_module': @@ -984,6 +986,27 @@ def _process_fx_graph(self) -> None: self._handle_get_attr(node) elif node.op == 'output': self._handle_output(node) + self._record_chain_head(node.name, first_layer_id) + + def _record_chain_head(self, node_name: str, first_layer_id: int) -> None: + """Route a node's incoming FX edges to the head of the chain it emitted. + + Only a contiguous chain of data layers ending at the node's registered + layer qualifies (BatchNorm SCALE -> BIAS, const - x, DequantizeLinear). + Other expansions keep the default wiring, e.g. sibling SLICEs of a + split (pinned via _fx_pred_override) or a CONSTANT emitted before its + consumer. + """ + last = self.node_to_layer_id.get(node_name) + if last is None or last <= first_layer_id: + return + chain = range(first_layer_id, last + 1) + if not all(self.layers[lid].in_vars for lid in chain): + return + for prev, cur in zip(chain, chain[1:]): + if not set(self.layers[cur].in_vars) <= set(self.layers[prev].out_vars): + return + self._node_first_layer_id[node_name] = first_layer_id def _handle_call_module(self, node: fx.Node) -> None: """Handle call_module node.""" @@ -1007,6 +1030,10 @@ def _handle_call_module(self, node: fx.Node) -> None: def _handle_call_function(self, node: fx.Node) -> None: """Handle call_function node.""" + if node.target is nn.functional.adaptive_avg_pool2d: + self._process_adaptive_avg_pool2d_function(node) + return + target_name = str(node.target).lower() handlers = { @@ -1088,14 +1115,17 @@ def _build_preds_succs(self) -> Tuple[Dict[int, List[int]], Dict[int, List[int]] """Build preds and succs dictionaries, validating DAG property. Uses a hybrid approach: - 1. Build edges from FX graph for mapped nodes - 2. For unmapped intermediate layers (e.g., SCALE from BatchNorm), - connect them sequentially based on layer creation order + 1. Build edges from FX graph for mapped nodes; a node converted into + a chain of layers is entered at the chain head + 2. Connect the remaining layers of a multi-layer conversion (e.g. BIAS + after BatchNorm's SCALE) sequentially by layer creation order """ n_layers = len(self.layers) - # Identify which layers are mapped to FX nodes + # Identify which layers are mapped to FX nodes. The head of a chain + # conversion receives that node's incoming FX edges. mapped_layer_ids = set(lid for lid in self.node_to_layer_id.values() if lid >= 0) + mapped_layer_ids.update(self._node_first_layer_id.values()) preds: Dict[int, List[int]] = {i: [] for i in range(n_layers)} succs: Dict[int, List[int]] = {i: [] for i in range(n_layers)} @@ -1108,6 +1138,7 @@ def _build_preds_succs(self) -> Tuple[Dict[int, List[int]], Dict[int, List[int]] layer_id = self.node_to_layer_id.get(node_name) if layer_id is None or layer_id < 0: continue + layer_id = self._node_first_layer_id.get(node_name, layer_id) for pred_name in pred_names: pred_layer_id = self.node_to_layer_id.get(pred_name) @@ -1122,8 +1153,8 @@ def _build_preds_succs(self) -> Tuple[Dict[int, List[int]], Dict[int, List[int]] if pred_layer_id != layer_id and pred_layer_id not in preds[layer_id]: preds[layer_id].append(pred_layer_id) - # Second, connect unmapped layers (SCALE, BIAS from BatchNorm, etc.) sequentially - # These layers are internal to a multi-layer conversion and should connect to i-1. + # Second, connect layers internal to a multi-layer conversion (e.g. BIAS + # after BatchNorm's SCALE) sequentially: they connect to i-1. # Source layers (``in_vars == []``, e.g. CONSTANT emitted by # ``_ensure_constant_vars``) are intentionally exempt -- they have no # data dependency on anything upstream, so wiring them to ``i-1`` would @@ -1188,6 +1219,9 @@ def _convert_module(self, mod: nn.Module) -> None: if kind == LayerKind.RELU.value: self._convert_activation(mod, LayerKind.RELU) return + if kind == LayerKind.RELU6.value: + self._convert_activation(mod, LayerKind.RELU6) + return if kind == LayerKind.CONV2D.value and isinstance(mod, nn.Conv2d): self._convert_conv2d(mod) return @@ -1904,6 +1938,14 @@ def _process_mul_operation(self, node: fx.Node) -> None: self.shape = x_shape self._register_node(node.name, layer_id) + def _process_adaptive_avg_pool2d_function(self, node: fx.Node) -> None: + """Process F.adaptive_avg_pool2d() via the nn.AdaptiveAvgPool2d converter.""" + if not self._get_predecessor_state(node): + return + output_size = node.args[1] if len(node.args) > 1 else node.kwargs["output_size"] + self._convert_adaptive_avgpool2d(nn.AdaptiveAvgPool2d(output_size)) + self._register_node(node.name) + def _process_mean_operation(self, node: fx.Node) -> None: """Process torch.mean().""" if not self._get_predecessor_state(node): @@ -2076,6 +2118,7 @@ def run(self) -> Net: # Build and validate network preds, succs = self._build_layer_graph() + _assert_dataflow(self.layers, preds) net = Net(layers=self.layers, preds=preds, succs=succs) from act.back_end.layer_util import validate_graph @@ -2223,7 +2266,7 @@ def _build_layer_graph(self) -> Tuple[Dict[int, List[int]], Dict[int, List[int]] # ----------------------------------------------------------------------------- def main(): - """Main entry point for PyTorch→ACT conversion and verification testing.""" + """Main entry point for PyTorch→ACT conversion testing.""" # Initialize debug file (GUARDED) if PerformanceOptions.debug_tf: debug_file = PerformanceOptions.debug_output_file @@ -2245,61 +2288,22 @@ def main(): print(f" {len(successful_models)} models passed inference tests") if not successful_models: - print(" No successful models to verify!") - exit(1) - - # Step 3: Convert and verify all successful models (memory-efficient) - print(f"\n Step 3: Converting and verifying all {len(successful_models)} successful models...") - print(f" Processing one at a time to avoid memory issues...") + raise RuntimeError("No synthesized model passed inference; nothing to convert") - # Import verification functions - from act.back_end.verifier import verify_once + # Step 3: Convert every successful model (one at a time to bound memory). + print(f"\n Step 3: Converting all {len(successful_models)} successful models...") import gc import torch as torch_module - conversion_results = {} - verification_results = {} - conversion_success_count = 0 - verification_success_count = 0 - - # Step 4: Initialize solvers (moved earlier to reuse for all models) - print("\n Step 4: Initializing solvers...") - gurobi_solver = None - torch_solver = None - - try: - gurobi_solver = GurobiSolver() - print(" Gurobi solver available") - except Exception as e: - print(f" Gurobi initialization failed: {e}") - - try: - torch_solver = TorchLPSolver() - print(f" TorchLP solver available (device: {torch_solver._device})") - except Exception as e: - print(f" TorchLP initialization failed: {e}") - - solvers_to_test = [] - if gurobi_solver: - solvers_to_test.append(("Gurobi", gurobi_solver)) - if torch_solver: - solvers_to_test.append(("TorchLP", torch_solver)) - - if not solvers_to_test: - print(" No solvers available!") - print(" Will only test conversions without verification") - - print(f"\n Step 5: Processing all models...") + failed_conversions = {} for idx, (model_id, wrapped_model) in enumerate(successful_models.items(), 1): print(f"\n [{idx}/{len(successful_models)}] Processing '{model_id}'...") - # === CONVERSION === try: net = TorchToACT(wrapped_model).run() - # Verify the conversion produced a valid net if not net.layers: raise ValueError("Net should have layers") if net.layers[0].kind != LayerKind.INPUT.value: @@ -2309,74 +2313,30 @@ def main(): layer_types = " -> ".join([layer.kind for layer in net.layers]) print(f" Conversion: {len(net.layers)} layers ({layer_types})") - - conversion_results[model_id] = "SUCCESS" - conversion_success_count += 1 - - except Exception as e: - conversion_results[model_id] = f"FAILED: {str(e)[:100]}..." + del net + except Exception as e: # noqa: BLE001 — collect every failure, raise below + failed_conversions[model_id] = e print(f" Conversion FAILED: {e}") - continue # Skip verification if conversion failed - - # === VERIFICATION (only if solvers available) === - if solvers_to_test: - model_verification = {} - - for solver_name, solver in solvers_to_test: - try: - # TEMPORARILY COMMENTED OUT: Testing if verify_once causes memory issue - # res = verify_once(net, solver=solver, timelimit=30.0) - # status = res.status - status = "SKIPPED" # Placeholder to test memory usage - model_verification[solver_name] = status - print(f" Verification ({solver_name}): {status} (verify_once commented out)") - - if status == "UNSAT" or status == "SAT": - verification_success_count += 1 - - except Exception as e: - model_verification[solver_name] = f"ERROR: {str(e)[:50]}" - print(f" Verification ({solver_name}): ERROR - {str(e)[:50]}") - - verification_results[model_id] = model_verification - - # === MEMORY CLEANUP === - # Free memory from this net immediately (no need to store) - del net - # Clean up memory periodically if idx % 10 == 0: gc.collect() if torch_module.cuda.is_available(): torch_module.cuda.empty_cache() - # === FINAL SUMMARY === total_count = len(successful_models) print(f"\nFinal Results:") - print(f" Conversions: {conversion_success_count}/{total_count} ({conversion_success_count/total_count*100:.1f}%)") + print(f" Conversions: {total_count - len(failed_conversions)}/{total_count}") - if solvers_to_test and verification_results: - # Count successful verifications (UNSAT or SAT results) - total_verifications = sum(len(v) for v in verification_results.values()) - successful_verifications = sum( - 1 for results in verification_results.values() - for status in results.values() - if isinstance(status, str) and status in ["UNSAT", "SAT"] - ) - print(f" Verifications: {successful_verifications}/{total_verifications} successful") - - # Print failed conversions if any - failed_conversions = {k: v for k, v in conversion_results.items() if v != "SUCCESS"} - if failed_conversions: - print(f"\n Failed conversions: {len(failed_conversions)}") - for model_id, error in list(failed_conversions.items())[:5]: # Show first 5 - print(f" - {model_id}: {error}") - - # Print debug file location (GUARDED) if PerformanceOptions.debug_tf: print(f"\nDebug log written to: {PerformanceOptions.debug_output_file}") - print("\nTorch->ACT conversion and verification completed!") + if failed_conversions: + details = "\n ".join(f"{model_id}: {err}" for model_id, err in failed_conversions.items()) + raise RuntimeError( + f"{len(failed_conversions)}/{total_count} Torch→ACT conversion(s) failed:\n {details}" + ) + + print("\nTorch->ACT conversion completed!") if __name__ == "__main__": diff --git a/act/pipeline/verification/utils.py b/act/pipeline/verification/utils.py index 2c6deed78..1053fa735 100644 --- a/act/pipeline/verification/utils.py +++ b/act/pipeline/verification/utils.py @@ -14,14 +14,8 @@ #===---------------------------------------------------------------------===# import time -import psutil import logging -import functools -import threading -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import List, Callable, Any, Dict, Optional, Tuple -from dataclasses import dataclass -from contextlib import contextmanager +from typing import List, Any, Dict, Optional, Tuple import torch logger = logging.getLogger(__name__) @@ -70,6 +64,26 @@ def _assert_dag(preds: Dict[int, List[int]], succs: Dict[int, List[int]], n_laye raise ValueError(f"Layer graph contains a cycle! Nodes: {cycle_nodes}") +def _assert_dataflow(layers: List[Any], preds: Dict[int, List[int]]) -> None: + """Assert every layer reads its in_vars from its predecessors' out_vars. + + A violation means an edge bypasses a layer (e.g. BatchNorm's SCALE), so + abstract interpretation would propagate the bounds of the wrong tensor. + """ + by_id = {layer.id: layer for layer in layers} + for layer in layers: + layer_preds = preds.get(layer.id, []) + if not layer.in_vars or not layer_preds: + continue + produced = set().union(*(by_id[p].out_vars for p in layer_preds)) + missing = set(layer.in_vars) - produced + assert not missing, ( + f"Layer {layer.id} ({layer.kind}) reads {len(missing)}/{len(layer.in_vars)} " + f"input vars not produced by its predecessors " + f"{[(p, by_id[p].kind) for p in layer_preds]}" + ) + + def _normalize_axes(axes: Any, rank: int) -> List[int]: """Sort + dedupe + normalise possibly-negative axes against ``rank``.""" return sorted({(a + rank) if a < 0 else a for a in (int(x) for x in axes)}) @@ -138,363 +152,6 @@ def _broadcast_const_to_size(const: torch.Tensor, size: int, dtype: torch.dtype) ) -@dataclass -class PerformanceMetrics: # pragma: no cover - """Performance metrics for validation operations.""" - execution_time: float - peak_memory_mb: float - cpu_usage_percent: float - gpu_memory_mb: Optional[float] = None - - -@dataclass -class ParallelResult: # pragma: no cover - """Result from parallel execution.""" - results: List[Any] - failed_tasks: List[Tuple[int, Exception]] - total_time: float - metrics: PerformanceMetrics - - -class PerformanceProfiler: # pragma: no cover - """Performance profiling utilities for validation operations.""" - - def __init__(self): - self.start_time: Optional[float] = None - self.start_memory: Optional[float] = None - self.peak_memory: float = 0 - self.cpu_usage_samples: List[float] = [] - self._monitoring_thread: Optional[threading.Thread] = None - self._stop_monitoring = threading.Event() - - def start(self) -> None: - """Start performance monitoring.""" - self.start_time = time.time() - self.start_memory = self._get_memory_usage() - self.peak_memory = self.start_memory - self.cpu_usage_samples = [] - self._stop_monitoring.clear() - - # Start monitoring thread - self._monitoring_thread = threading.Thread(target=self._monitor_resources) - self._monitoring_thread.daemon = True - self._monitoring_thread.start() - - logger.debug("Performance profiling started") - - def stop(self) -> PerformanceMetrics: - """Stop monitoring and return performance metrics.""" - if self.start_time is None: - raise RuntimeError("Profiler not started") - - # Stop monitoring thread - self._stop_monitoring.set() - if self._monitoring_thread: - self._monitoring_thread.join(timeout=1.0) - - execution_time = time.time() - self.start_time - avg_cpu_usage = sum(self.cpu_usage_samples) / len(self.cpu_usage_samples) if self.cpu_usage_samples else 0 - - metrics = PerformanceMetrics( - execution_time=execution_time, - peak_memory_mb=self.peak_memory, - cpu_usage_percent=avg_cpu_usage, - gpu_memory_mb=self._get_gpu_memory() if torch.cuda.is_available() else None - ) - - logger.debug(f"Performance metrics: {metrics}") - return metrics - - def _monitor_resources(self) -> None: - """Monitor resource usage in background thread.""" - while not self._stop_monitoring.wait(0.1): # Sample every 100ms - try: - # Monitor memory - current_memory = self._get_memory_usage() - self.peak_memory = max(self.peak_memory, current_memory) - - # Monitor CPU - cpu_usage = psutil.cpu_percent(interval=None) - self.cpu_usage_samples.append(cpu_usage) - - except Exception as e: - logger.warning(f"Resource monitoring error: {e}") - - def _get_memory_usage(self) -> float: - """Get current memory usage in MB.""" - process = psutil.Process() - return process.memory_info().rss / 1024 / 1024 - - def _get_gpu_memory(self) -> Optional[float]: - """Get GPU memory usage in MB.""" - if not torch.cuda.is_available(): - return None - try: - return torch.cuda.memory_allocated() / 1024 / 1024 - except Exception: - return None - - -@contextmanager -def profile_performance(): # pragma: no cover - """Context manager for performance profiling.""" - profiler = PerformanceProfiler() - profiler.start() - try: - yield profiler - finally: - metrics = profiler.stop() - yield metrics - - -class ParallelExecutor: # pragma: no cover - """Utilities for parallel execution of validation tasks.""" - - def __init__(self, max_workers: Optional[int] = None, timeout: Optional[float] = None): - """ - Initialize parallel executor. - - Args: - max_workers: Maximum number of worker threads - timeout: Timeout for individual tasks in seconds - """ - self.max_workers = max_workers or min(32, (psutil.cpu_count() or 1) + 4) - self.timeout = timeout - - def execute_parallel(self, - tasks: List[Callable], - task_args: Optional[List[Tuple]] = None, - task_kwargs: Optional[List[Dict]] = None) -> ParallelResult: - """ - Execute tasks in parallel. - - Args: - tasks: List of callable tasks to execute - task_args: List of argument tuples for each task - task_kwargs: List of keyword argument dicts for each task - - Returns: - ParallelResult with results and performance metrics - """ - if task_args is None: - task_args = [() for _ in tasks] - if task_kwargs is None: - task_kwargs = [{} for _ in tasks] - - if len(tasks) != len(task_args) or len(tasks) != len(task_kwargs): - raise ValueError("tasks, task_args, and task_kwargs must have same length") - - results = [] - failed_tasks = [] - - with profile_performance() as profiler: - with ThreadPoolExecutor(max_workers=self.max_workers) as executor: - # Submit all tasks - future_to_index = { - executor.submit(task, *args, **kwargs): i - for i, (task, args, kwargs) in enumerate(zip(tasks, task_args, task_kwargs)) - } - - # Collect results - for future in as_completed(future_to_index, timeout=self.timeout): - task_index = future_to_index[future] - try: - result = future.result() - results.append((task_index, result)) - except Exception as e: - failed_tasks.append((task_index, e)) - logger.error(f"Task {task_index} failed: {e}") - - # Sort results by original task index - results.sort(key=lambda x: x[0]) - sorted_results = [result for _, result in results] - - metrics = profiler.stop() - - return ParallelResult( - results=sorted_results, - failed_tasks=failed_tasks, - total_time=metrics.execution_time, - metrics=metrics - ) - - def map_parallel(self, func: Callable, items: List[Any]) -> ParallelResult: - """ - Apply function to list of items in parallel. - - Args: - func: Function to apply to each item - items: List of items to process - - Returns: - ParallelResult with mapped results - """ - tasks = [func for _ in items] - task_args = [(item,) for item in items] - - return self.execute_parallel(tasks, task_args) - - -def print_memory_usage(prefix: str = "") -> None: # pragma: no cover - """Print current memory usage information.""" - process = psutil.Process() - memory_info = process.memory_info() - - memory_mb = memory_info.rss / 1024 / 1024 - cpu_percent = process.cpu_percent() - - gpu_info = "" - if torch.cuda.is_available(): - gpu_memory_mb = torch.cuda.memory_allocated() / 1024 / 1024 - gpu_max_mb = torch.cuda.max_memory_allocated() / 1024 / 1024 - gpu_info = f", GPU: {gpu_memory_mb:.1f}MB (max: {gpu_max_mb:.1f}MB)" - - logger.info(f"{prefix}Memory: {memory_mb:.1f}MB, CPU: {cpu_percent:.1f}%{gpu_info}") - - -def clear_torch_cache() -> None: # pragma: no cover - """Clear PyTorch GPU cache if available.""" - if torch.cuda.is_available(): - torch.cuda.empty_cache() - logger.debug("Cleared PyTorch GPU cache") - - -def setup_logging(level: str = "INFO", format_str: Optional[str] = None) -> None: # pragma: no cover - """ - Setup logging configuration for the pipeline. - - Args: - level: Logging level (DEBUG, INFO, WARNING, ERROR) - format_str: Custom format string for log messages - """ - if format_str is None: - format_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" - - # Ensure log file goes to pipeline/log directory - import os - from pathlib import Path - pipeline_dir = Path(__file__).parent - log_dir = pipeline_dir / "log" - log_dir.mkdir(exist_ok=True) - log_file_path = log_dir / "pipeline_tests.log" - - logging.basicConfig( - level=getattr(logging, level.upper()), - format=format_str, - handlers=[ - logging.StreamHandler(), - logging.FileHandler(log_file_path) - ] - ) - - # Reduce noise from some libraries - logging.getLogger("matplotlib").setLevel(logging.WARNING) - logging.getLogger("PIL").setLevel(logging.WARNING) - - -def retry_on_failure(max_retries: int = 3, delay: float = 1.0, backoff: float = 2.0): # pragma: no cover - """ - Decorator to retry function on failure. - - Args: - max_retries: Maximum number of retry attempts - delay: Initial delay between retries in seconds - backoff: Backoff multiplier for delay - """ - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - current_delay = delay - last_exception = None - - for attempt in range(max_retries + 1): - try: - return func(*args, **kwargs) - except Exception as e: - last_exception = e - if attempt < max_retries: - logger.warning(f"Attempt {attempt + 1} failed for {func.__name__}: {e}") - time.sleep(current_delay) - current_delay *= backoff - else: - logger.error(f"All {max_retries + 1} attempts failed for {func.__name__}") - - raise last_exception - - return wrapper - return decorator - - -def timeout_handler(timeout_seconds: float): # pragma: no cover - """ - Decorator to add timeout to function execution. - - Args: - timeout_seconds: Timeout in seconds - """ - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - import signal - - def timeout_signal_handler(signum, frame): - raise TimeoutError(f"Function {func.__name__} timed out after {timeout_seconds} seconds") - - # Set up signal handler - old_handler = signal.signal(signal.SIGALRM, timeout_signal_handler) - signal.alarm(int(timeout_seconds)) - - try: - result = func(*args, **kwargs) - finally: - # Restore old signal handler - signal.alarm(0) - signal.signal(signal.SIGALRM, old_handler) - - return result - - return wrapper - return decorator - - -class ProgressTracker: # pragma: no cover - """Track progress of long-running operations.""" - - def __init__(self, total_items: int, description: str = "Processing"): - self.total_items = total_items - self.description = description - self.completed_items = 0 - self.start_time = time.time() - - def update(self, completed: int = 1) -> None: - """Update progress by specified number of completed items.""" - self.completed_items += completed - self._print_progress() - - def _print_progress(self) -> None: - """Print current progress.""" - if self.total_items == 0: - return - - percentage = (self.completed_items / self.total_items) * 100 - elapsed_time = time.time() - self.start_time - - if self.completed_items > 0: - eta = (elapsed_time / self.completed_items) * (self.total_items - self.completed_items) - eta_str = f", ETA: {eta:.1f}s" - else: - eta_str = "" - - logger.info(f"{self.description}: {self.completed_items}/{self.total_items} " - f"({percentage:.1f}%) - {elapsed_time:.1f}s elapsed{eta_str}") - - def finish(self) -> None: - """Mark progress as complete.""" - self.completed_items = self.total_items - elapsed_time = time.time() - self.start_time - logger.info(f"{self.description} completed in {elapsed_time:.1f}s") - - # ----------------------------------------------------------------------------- # ONNX -> ACT layer handlers (bound onto _LayerGraphBuilder via setattr in # torch2act.py). Kept here only to keep that file manageable; ``self`` is always @@ -1257,10 +914,7 @@ def _convert_OnnxMinMax(self, mod: nn.Module, node: fx.Node) -> None: raise ValueError(f"OnnxMinMax at {node.name}: '{n.name}' not registered") xv, yv = self.node_outputs[args[0].name], self.node_outputs[args[1].name] xs, ys = self.node_shapes[args[0].name], self.node_shapes[args[1].name] - try: - output_shape = tuple(int(d) for d in torch.broadcast_shapes(xs, ys)) - except RuntimeError: - output_shape = xs if _prod(xs) >= _prod(ys) else ys + output_shape = tuple(int(d) for d in torch.broadcast_shapes(xs, ys)) out_vars = self._alloc_ids(_prod(output_shape) or 1) layer_id = self._add_layer( kind.value, @@ -1335,14 +989,11 @@ def _convert_OnnxWhere(self, mod: nn.Module, node: fx.Node) -> None: cv = self.node_outputs[args[0].name] xv = self.node_outputs[args[1].name] yv = self.node_outputs[args[2].name] - try: - output_shape = tuple(int(d) for d in torch.broadcast_shapes( - self.node_shapes[args[0].name], - self.node_shapes[args[1].name], - self.node_shapes[args[2].name], - )) - except RuntimeError: - output_shape = self.node_shapes[args[1].name] + output_shape = tuple(int(d) for d in torch.broadcast_shapes( + self.node_shapes[args[0].name], + self.node_shapes[args[1].name], + self.node_shapes[args[2].name], + )) out_vars = self._alloc_ids(_prod(output_shape) or 1) layer_id = self._add_layer( LayerKind.WHERE.value, @@ -1539,10 +1190,7 @@ def _convert_OnnxCompare(self, mod: nn.Module, node: fx.Node) -> None: raise ValueError(f"OnnxCompare at {node.name}: '{n.name}' not registered") xv, yv = self.node_outputs[x_node.name], self.node_outputs[y_node.name] xs, ys = self.node_shapes[x_node.name], self.node_shapes[y_node.name] - try: - output_shape = tuple(int(d) for d in torch.broadcast_shapes(xs, ys)) - except RuntimeError: - output_shape = xs if _prod(xs) >= _prod(ys) else ys + output_shape = tuple(int(d) for d in torch.broadcast_shapes(xs, ys)) out_vars = self._alloc_ids(_prod(output_shape) or 1) layer_id = self._add_layer( LayerKind.COMPARE.value, diff --git a/act/pipeline/verification/validate_verifier.py b/act/pipeline/verification/validate_verifier.py index d1eea42c8..b280ad178 100644 --- a/act/pipeline/verification/validate_verifier.py +++ b/act/pipeline/verification/validate_verifier.py @@ -319,24 +319,17 @@ def _batchify_assert_layer(self, L: Layer, target_B: int) -> None: (v for v in high.values() if torch.is_tensor(v)), None ) ) - try: - spec = OutputSpec(kind=kind, **high) - new_params = spec.encode_linear( - B=target_B, - n_out=n_out, - device=ref.device if ref is not None else torch.device(self.device), - dtype=( - ref.dtype - if ref is not None and ref.dtype.is_floating_point - else self.dtype - ), - ) - except Exception as e: - logger.warning( - f"_batchify_assert_layer({kind}) at B={target_B}: " - f"re-encode failed: {e}" - ) - return + spec = OutputSpec(kind=kind, **high) + new_params = spec.encode_linear( + B=target_B, + n_out=n_out, + device=ref.device if ref is not None else torch.device(self.device), + dtype=( + ref.dtype + if ref is not None and ref.dtype.is_floating_point + else self.dtype + ), + ) L.params.clear() L.params.update(new_params) diff --git a/act/util/device_manager.py b/act/util/device_manager.py index 3b99ea431..e92f52b48 100644 --- a/act/util/device_manager.py +++ b/act/util/device_manager.py @@ -5,10 +5,8 @@ import torch from typing import Tuple, Optional -try: +if hasattr(torch.sparse, "check_sparse_tensor_invariants"): torch.sparse.check_sparse_tensor_invariants.disable() -except Exception: - pass # Global initialization state _INITIALIZED = False @@ -66,61 +64,50 @@ def initialize_device(device: str, dtype: str) -> None: """ global _INITIALIZED - try: - # Handle gpu/cuda aliasing - if device == 'gpu': - device = 'cuda' - print(f"🔄 Device alias: 'gpu' → 'cuda'") + # Handle gpu/cuda aliasing + if device == 'gpu': + device = 'cuda' + print(f"🔄 Device alias: 'gpu' → 'cuda'") - # Determine target device - if device == 'cpu': - target_device = torch.device("cpu") - elif device == 'cuda': - if torch.cuda.is_available(): - target_device = torch.device("cuda:0") - else: - target_device = torch.device("cpu") - print(f"⚠️ CUDA not available, falling back to CPU") - elif device == 'mps': - if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): - target_device = torch.device("mps") - else: - target_device = torch.device("cpu") - print(f"⚠️ MPS not available, falling back to CPU") - else: - # Unknown device, default to CPU - target_device = torch.device("cpu") - print(f"⚠️ Unknown device '{device}', using CPU") + # Determine target device. CLI defaults already select an available device, + # so an unavailable explicit request is a configuration error. + if device == 'cpu': + target_device = torch.device("cpu") + elif device == 'cuda': + if not torch.cuda.is_available(): + raise RuntimeError("CUDA device requested but CUDA is not available") + target_device = torch.device("cuda:0") + elif device == 'mps': + if not hasattr(torch.backends, 'mps') or not torch.backends.mps.is_available(): + raise RuntimeError("MPS device requested but MPS is not available") + target_device = torch.device("mps") + else: + raise ValueError( + f"Unknown device {device!r}; expected one of 'cpu', 'cuda', 'gpu', or 'mps'" + ) - # Determine target dtype - if dtype == 'float32': - target_dtype = torch.float32 - elif dtype == 'float64': - target_dtype = torch.float64 - else: - # Unknown dtype, default to float64 - target_dtype = torch.float64 - print(f"⚠️ Unknown dtype '{dtype}', using float64") + # Determine target dtype. + if dtype == 'float32': + target_dtype = torch.float32 + elif dtype == 'float64': + target_dtype = torch.float64 + else: + raise ValueError( + f"Unknown dtype {dtype!r}; expected 'float32' or 'float64'" + ) - # Set PyTorch global defaults - torch.set_default_dtype(target_dtype) - if hasattr(torch, 'set_default_device'): - torch.set_default_device(target_device) + # Set PyTorch global defaults + torch.set_default_dtype(target_dtype) + if hasattr(torch, 'set_default_device'): + torch.set_default_device(target_device) - precision = _apply_precision_policy() + precision = _apply_precision_policy() - print( - f"✅ Device Manager Initialized: device={target_device}, " - f"dtype={target_dtype}, {precision}" - ) - _INITIALIZED = True - - except Exception as e: - print(f"❌ Device initialization failed: {e}") - print(f" Falling back to CPU + float64") - torch.set_default_dtype(torch.float64) - _apply_precision_policy() - _INITIALIZED = True + print( + f"✅ Device Manager Initialized: device={target_device}, " + f"dtype={target_dtype}, {precision}" + ) + _INITIALIZED = True def get_default_device() -> torch.device: @@ -133,10 +120,7 @@ def get_default_device() -> torch.device: _ensure_initialized() if hasattr(torch, 'get_default_device'): - try: - return torch.get_default_device() - except: - return torch.device("cpu") + return torch.get_default_device() else: # For older PyTorch versions, check where a test tensor is created test_tensor = torch.zeros(1) @@ -186,19 +170,13 @@ def _ensure_initialized(): device = 'cuda' if torch.cuda.is_available() else 'cpu' dtype = 'float64' - # Initialize with defaults (no print statements for lazy init) - try: - target_device = torch.device("cuda:0") if device == 'cuda' else torch.device("cpu") - target_dtype = torch.float64 - - torch.set_default_dtype(target_dtype) - if hasattr(torch, 'set_default_device'): - torch.set_default_device(target_device) - - _apply_precision_policy() - _INITIALIZED = True - except Exception: - # Silent fallback to CPU + float64 - torch.set_default_dtype(torch.float64) - _apply_precision_policy() - _INITIALIZED = True + # Initialize with defaults (no print statements for lazy init). + target_device = torch.device("cuda:0") if device == 'cuda' else torch.device("cpu") + target_dtype = torch.float64 + + torch.set_default_dtype(target_dtype) + if hasattr(torch, 'set_default_device'): + torch.set_default_device(target_device) + + _apply_precision_policy() + _INITIALIZED = True diff --git a/act/util/format_utils.py b/act/util/format_utils.py index 31bd82e77..3042bb764 100644 --- a/act/util/format_utils.py +++ b/act/util/format_utils.py @@ -33,14 +33,16 @@ def format_bytes(size_bytes: Union[int, float], precision: int = 1) -> str: def dir_size(path: Union[str, Path]) -> int: - """Return total size in bytes of all files under ``path`` (0 on error).""" + """Return total size in bytes of all files under ``path``.""" + root = Path(path) + if not root.exists(): + raise FileNotFoundError(f"Directory does not exist: {root}") + if not root.is_dir(): + raise NotADirectoryError(f"Expected a directory, got: {root}") total = 0 - try: - for item in Path(path).rglob("*"): - if item.is_file(): - total += item.stat().st_size - except Exception: - pass + for item in root.rglob("*"): + if item.is_file(): + total += item.stat().st_size return total diff --git a/act/util/model_inference.py b/act/util/model_inference.py index 4753ba6b7..07c8b8bc7 100644 --- a/act/util/model_inference.py +++ b/act/util/model_inference.py @@ -42,10 +42,14 @@ def infer_single_model(combo_id: Union[str, Tuple], model: nn.Module, input_tens output = model(input_tensor) # Extract tensor if model returns dict (VerifiableModel) if isinstance(output, dict): + if 'output' not in output: + raise KeyError( + f"Model {combo_id!r} returned a dictionary without an 'output' entry" + ) output = output['output'] return True, output, None - except Exception as e: - return False, None, str(e)[:100] + except Exception as exc: + raise RuntimeError(f"Model inference failed for {combo_id!r}: {exc}") from exc # Main model inference function @@ -65,8 +69,9 @@ def model_inference(models: Dict[Union[str, Tuple], nn.Module]) -> Dict[Union[st # Handle case where no models were generated if not models: - print("⚠️ No models to test - check spec generation and synthesis configuration") - return {} + raise ValueError( + "No models to test; check spec generation and synthesis configuration" + ) success_count = 0 failure_count = 0 @@ -78,15 +83,17 @@ def model_inference(models: Dict[Union[str, Tuple], nn.Module]) -> Dict[Union[st # Extract input and label from InputLayer (named child on VerifiableModel) input_layer = model.input_layer if not hasattr(input_layer, 'input_tensor'): - print(f"⚠️ Model {combo_id} missing input_tensor in InputLayer") - failure_count += 1 - continue + raise AttributeError( + f"Model {combo_id!r} InputLayer is missing input_tensor" + ) test_input = input_layer.input_tensor test_label = input_layer.label # Move input to same device as model - model_device = next(model.parameters()).device + # Parameter-free models (e.g. pure-constant ONNX nets) keep the input's device. + first_param = next(model.parameters(), None) + model_device = first_param.device if first_param is not None else test_input.device test_input = test_input.to(model_device) # Parse combo_id: handle tuple keys (dataset, model, in_kind, out_kind) @@ -107,9 +114,12 @@ def model_inference(models: Dict[Union[str, Tuple], nn.Module]) -> Dict[Union[st success, output, error_msg = infer_single_model(combo_id, model, test_input) if success: - # Validate prediction against ground truth label tensor (N,), output is (N, num_classes) - pred_classes = output.argmax(dim=1) # (N,) - is_correct = (pred_classes == test_label).all().item() + # Validate prediction against ground truth label tensor (N,), output is (N, num_classes). + # Unlabelled or non-classification outputs (e.g. VNNLIB safety nets) have no accuracy. + is_correct = False + if test_label is not None and test_label.numel() > 0 and output.dim() >= 2: + pred_classes = output.argmax(dim=1) # (N,) + is_correct = (pred_classes == test_label).all().item() success_count += 1 if is_correct: @@ -137,4 +147,3 @@ def model_inference(models: Dict[Union[str, Tuple], nn.Module]) -> Dict[Union[st print(f" 💡 Tip: Use domain-matched combinations (mnist+mnist, cifar+cifar) for 100% success") return successful_models - diff --git a/act/util/stats.py b/act/util/stats.py index 42ffd454e..af59c7bec 100644 --- a/act/util/stats.py +++ b/act/util/stats.py @@ -23,13 +23,11 @@ # #===----------------------------------------------------------------------===# -import os import torch from enum import Enum from dataclasses import dataclass, field -from typing import Dict, Any, List, Tuple, Optional +from typing import Dict, Any, List, Optional -from act.util.format_utils import rule # ============================================================================= @@ -153,382 +151,3 @@ def to_verify_results(self, *, source: str = "dual_bound") -> List["VerifyResult }, )) return results - - -# ============================================================================= -# Statistics and Logging Utilities -# ============================================================================= - -class ACTStats: - """ - Utility class for monitoring and reporting memory usage during verification. - - Provides comprehensive memory monitoring including system memory, process memory, - and GPU memory statistics. Gracefully handles missing dependencies. - """ - - @staticmethod - def print_memory_usage(stage_name: str = "") -> float: - """ - Print comprehensive memory usage statistics for the current process and system. - - Reports process memory usage, system memory statistics, and GPU memory - information when available. Handles missing psutil dependency gracefully. - - Args: - stage_name: Optional label for the memory report stage - - Returns: - float: Process memory usage in MB, or 0 if psutil is unavailable - """ - try: - import psutil - process = psutil.Process(os.getpid()) - memory_info = process.memory_info() - memory_mb = memory_info.rss / 1024 / 1024 - - system_memory = psutil.virtual_memory() - total_mb = system_memory.total / 1024 / 1024 - available_mb = system_memory.available / 1024 / 1024 - used_percent = (memory_mb / total_mb) * 100 - - print(f"[{stage_name}] Memory Usage:") - print(f"Process: {memory_mb:.1f} MB ({used_percent:.1f}% of total)") - print(f"System: {total_mb:.1f} MB total, {available_mb:.1f} MB available") - - if torch.cuda.is_available(): - gpu_memory_mb = torch.cuda.memory_allocated() / 1024 / 1024 - gpu_cached_mb = torch.cuda.memory_reserved() / 1024 / 1024 - print(f"GPU: {gpu_memory_mb:.1f} MB allocated, {gpu_cached_mb:.1f} MB cached") - - return memory_mb - - except ImportError: - print(f"⚠️ [{stage_name}] psutil not available, cannot monitor memory") - return 0.0 - except Exception as e: - print(f"⚠️ [{stage_name}] Error monitoring memory: {e}") - return 0.0 - - @classmethod - def get_memory_usage_mb(cls) -> float: - """ - Get current process memory usage in MB without printing. - - Returns: - float: Process memory usage in MB, or 0 if psutil is unavailable - """ - try: - import psutil - process = psutil.Process(os.getpid()) - memory_info = process.memory_info() - return memory_info.rss / 1024 / 1024 - except (ImportError, Exception): - return 0.0 - - @classmethod - def get_current_memory_usage(cls) -> float: - """ - Get current memory usage in MB (GPU if available, otherwise CPU). - - This is an enhanced version that automatically detects the best memory - source (GPU vs CPU) for tracking during bounds propagation. - - Returns: - float: Current memory usage in MB - """ - try: - if torch.cuda.is_available() and torch.cuda.device_count() > 0: - # Get GPU memory usage - gpu_memory_bytes = torch.cuda.memory_allocated() - return gpu_memory_bytes / (1024 * 1024) # Convert to MB - else: - # Get CPU memory usage for current process - import psutil - process = psutil.Process(os.getpid()) - cpu_memory_bytes = process.memory_info().rss - return cpu_memory_bytes / (1024 * 1024) # Convert to MB - except (ImportError, Exception): - # Fallback if psutil unavailable or memory monitoring fails - return 0.0 - - @classmethod - def get_gpu_memory_info(cls) -> Tuple[float, float]: - """ - Get GPU memory info (available, total) in MB. - - Returns: - tuple: (available_memory_mb, total_memory_mb) - (0.0, 0.0) if no GPU - """ - try: - if torch.cuda.is_available() and torch.cuda.device_count() > 0: - total_memory = torch.cuda.get_device_properties(0).total_memory - allocated_memory = torch.cuda.memory_allocated() - available_memory = total_memory - allocated_memory - return available_memory / (1024 * 1024), total_memory / (1024 * 1024) - else: - return 0.0, 0.0 - except Exception: - return 0.0, 0.0 - - @classmethod - def get_cpu_memory_usage(cls) -> float: - """ - Get CPU memory usage for current process in MB. - - Returns: - float: CPU memory usage in MB, or 0.0 if psutil unavailable - """ - try: - import psutil - process = psutil.Process(os.getpid()) - return process.memory_info().rss / (1024 * 1024) - except (ImportError, Exception): - return 0.0 - - @staticmethod - def print_verification_stats(prediction_stats: Dict[str, Any]) -> None: - """ - Print verification statistics from the prediction stats. - - Args: - prediction_stats: Dictionary containing various verification statistics - """ - # Hardcode verification_result_stat_dict as empty for now - verification_result_stat_dict = {} - - print("\n📊 Verification Results Summary:") - - # Display known statistics - known_keys = ['total', 'safe', 'unsafe', 'timeout', 'model_infer_failure'] - for key in known_keys: - if key in prediction_stats: - print(f" {key}: {prediction_stats[key]}") - - # Display any additional statistics - additional_stats = {k: v for k, v in prediction_stats.items() if k not in known_keys} - if additional_stats: - print(" Additional stats:") - for key, value in additional_stats.items(): - print(f" {key}: {value}") - - # Display verification result statistics if available - if verification_result_stat_dict: - print("\n🔍 Verification VerifyStatus Distribution:") - for status, count in verification_result_stat_dict.items(): - print(f" {status}: {count}") - - print(rule(50, "-")) - - @staticmethod - def print_final_verification_summary(results: List) -> Any: - """ - Print final verification results summary and return the overall verdict. - - Args: - results: List of VerifyStatus for all verified samples - - Returns: - Overall VerifyStatus (CERTIFIED if all safe, FALSIFIED if any unsafe, UNKNOWN otherwise) - """ - # VerifyStatus is now defined at module level in this file - - print("\n" + "🏆" + rule(70) + "🏆") - print("📊 Final verification results summary") - print("🏆" + rule(70) + "🏆") - - for idx, result in enumerate(results): - print(f"Sample {idx+1}: {result.name}") - - print(rule(60, "-")) - - certified_count = sum(1 for r in results if r == VerifyStatus.CERTIFIED) - falsified_count = sum(1 for r in results if r == VerifyStatus.FALSIFIED) - model_infer_failure_count = sum(1 for r in results if r == VerifyStatus.MODEL_INFER_FAILURE) - unknown_count = sum(1 for r in results if r == VerifyStatus.UNKNOWN) - total_count = len(results) - - valid_count = total_count - model_infer_failure_count - - print("Verification statistics:") - print(f"Total samples: {total_count}") - print(f"✅ CERTIFIED (safe): {certified_count} ") - print(f"❌ FALSIFIED (unsafe): {falsified_count} ") - print(f"⚠️ MODEL_INFER_FAILURE (model inference failed): {model_infer_failure_count} ") - print(f"❓ UNKNOWN: {unknown_count} ") - print(f"🔍 Valid verification samples: {valid_count} ") - - if valid_count > 0: - certified_percentage = (certified_count / valid_count) * 100 - falsified_percentage = (falsified_count / valid_count) * 100 - print(f"📊 CERTIFIED over valid samples: {certified_percentage:.2f}% ({certified_count}/{valid_count})") - print(f"📊 FALSIFIED over valid samples: {falsified_percentage:.2f}% ({falsified_count}/{valid_count})") - else: - print(" ⚠️ No valid verification samples") - - if total_count > 0: - certified_total_percentage = (certified_count / total_count) * 100 - falsified_total_percentage = (falsified_count / total_count) * 100 - model_infer_failure_percentage = (model_infer_failure_count / total_count) * 100 - print(f"📊 CERTIFIED over total: {certified_total_percentage:.2f}% ({certified_count}/{total_count})") - print(f"📊 FALSIFIED over total: {falsified_total_percentage:.2f}% ({falsified_count}/{total_count})") - print(f"📊 MODEL_INFER_FAILURE over total: {model_infer_failure_percentage:.2f}% ({model_infer_failure_count}/{total_count})") - - print(rule(60, "-")) - - if all(r == VerifyStatus.CERTIFIED for r in results): - final_result = VerifyStatus.CERTIFIED - print("✅ Final Result: CERTIFIED - all samples verified safe") - elif any(r == VerifyStatus.FALSIFIED for r in results): - final_result = VerifyStatus.FALSIFIED - print("❌ Final Result: FALSIFIED - at least one sample violates the property") - else: - final_result = VerifyStatus.UNKNOWN - print("❓ Final Result: UNKNOWN - inconclusive") - - print("🏆" + rule(70) + "🏆") - return final_result - - -class ACTLog: - """ - Centralized logging utilities for the Abstract Constraints Transformer (ACT). - - Provides structured logging methods for prediction validation, constraint application, - and verification processes with consistent formatting and emoji indicators. - """ - - @staticmethod - def log_correct_prediction( - predicted_label: int, - ground_truth_label: int, - sample_index: int, - verbose: bool = True - ) -> None: - """Log successful prediction with verbose output.""" - if verbose: - print(f"✅ Sample {sample_index + 1} Clean Prediction: Correct " - f"(pred: {predicted_label}, true: {ground_truth_label})") - - @staticmethod - def log_incorrect_prediction( - predicted_label: int, - ground_truth_label: int, - sample_index: int, - verbose: bool = True - ) -> None: - """Log failed prediction with verbose output.""" - if verbose: - print(f"❌ Sample {sample_index + 1} Clean Prediction: Incorrect " - f"(pred: {predicted_label}, true: {ground_truth_label})") - print("⚠️ Skipping verification - clean prediction already incorrect") - - @staticmethod - def log_prediction_failure(error: Exception) -> None: - """Log prediction validation failure.""" - print(f"❌ Clean prediction validation failed: {error}") - - @staticmethod - def log_constraint_application_start(cache_type_name: str, constraint_count: int, verbose: bool = True) -> None: - """Log the start of constraint application process.""" - if verbose: - print(f"[Bounds Fix] Applying {constraint_count} ReLU constraints " - f"to layer bounds cache") - print(f"[Bounds Fix] Using {cache_type_name} bounds cache") - - @staticmethod - def log_constraint_processing( - layer_name: str, - neuron_index: int, - activation_type: str, - verbose: bool = True - ) -> None: - """Log individual constraint processing.""" - if verbose: - print(f"[Bounds Fix] Processing constraint: " - f"{layer_name}[{neuron_index}] = {activation_type}") - - @staticmethod - def log_constraint_application_complete(verbose: bool = True) -> None: - """Log successful completion of constraint application.""" - if verbose: - print("✅ [Bounds Fix] ReLU constraint application done") - - @staticmethod - def log_no_bounds_cache_warning(verbose: bool = True) -> None: - """Log warning when no bounds cache is available.""" - if verbose: - print("⚠️ No layer bounds cache found, skipping constraint application") - - @staticmethod - def log_label_warning(sample_index: int) -> None: - """Log warning when ground truth label cannot be found.""" - print(f"⚠️ Could not get true label for sample {sample_index+1}, using default label 0") - - @staticmethod - def log_relu_constraints_set(constraints: List, verbose: bool = True) -> None: - """Log ReLU constraint configuration.""" - if verbose: - if constraints: - print(f"Set ReLU constraints: {len(constraints)} constraints") - for constraint in constraints: - print(f"{constraint['layer']}[{constraint['neuron_idx']}] = {constraint['constraint_type']}") - else: - print(f"Cleared ReLU constraints") - - @staticmethod - def log_constraint_error(message: str, verbose: bool = True) -> None: - """Log constraint processing errors.""" - if verbose: - print(f"⚠️ {message}") - - @staticmethod - def log_constraint_bound_update( - constraint_type: str, - neuron_index: int, - original_bound: float, - new_bound: float, - verbose: bool = True - ) -> None: - """Log individual bound constraint updates.""" - if verbose: - print(f"{constraint_type} constraint: neuron {neuron_index} " - f"{'ub' if constraint_type == 'inactive' else 'lb'}: {original_bound:.6f} → {new_bound:.6f}") - - @staticmethod - def log_verification_info(message: str) -> None: - """Log general verification information.""" - print(f"{message}") - - @staticmethod - def log_verification_warning(message: str) -> None: - """Log verification warnings.""" - print(f"Warning: {message}") - - @staticmethod - def log_bab_start(sample_idx: int) -> None: - """Log BaB specification refinement start.""" - print(f"Starting generic BaB specification refinement verification (sample {sample_idx})") - print(f"Framework: theoretically-aligned specification refinement") - - @staticmethod - def log_bab_results( - status_name: str, - total_subproblems: int, - spurious_count: int, - has_real_counterexample: bool, - max_depth: int, - total_time: float - ) -> None: - """Log BaB verification results.""" - print(f"Specification refinement verification finished: {status_name}") - print(f"Total subproblems: {total_subproblems}") - print(f"Spurious counterexamples: {spurious_count}") - print(f"Real counterexample: {'Yes' if has_real_counterexample else 'No'}") - print(f"Max depth: {max_depth}") - print(f"Total time: {total_time:.2f}s") - - @staticmethod - def log_bab_error(error: Exception) -> None: - """Log BaB verification errors.""" - print(f"⚠️ Specification refinement verification error: {error}") \ No newline at end of file diff --git a/vnncomp/act_run_instance.py b/vnncomp/act_run_instance.py index 91555e0ff..e89ef689c 100755 --- a/vnncomp/act_run_instance.py +++ b/vnncomp/act_run_instance.py @@ -23,6 +23,7 @@ import argparse import sys import time +import traceback from pathlib import Path # This runner lives in /vnncomp/; the repo root (which holds the `act` @@ -50,7 +51,8 @@ def run_vnncomp_instance(args) -> None: t0 = time.time() initialize_device(args.device, args.dtype) - from act.back_end.bab.bab import clear_violation_check_module_cache, verify_bab_batched + from act.back_end.bab.bab import verify_bab_batched + from act.back_end.bab.violation import clear_violation_check_module_cache from act.config.config import build_vnncomp_bab_config from act.back_end.solver.solver_torchlp import TorchLPSolver from act.front_end.model_synthesis import merge_split_relus @@ -65,12 +67,7 @@ def run_vnncomp_instance(args) -> None: def remaining() -> float: return args.timeout - (time.time() - t0) - args.margin - try: - sr = create_specs_from_paths(args.onnx, args.vnnlib) - except SystemExit as exc: - print(f"[load failed] {exc}", file=sys.stderr) - write_vnncomp_result(args.output, "unknown") - return + sr = create_specs_from_paths(args.onnx, args.vnnlib) raw_model = sr[2] param = next(raw_model.parameters(), None) @@ -103,24 +100,27 @@ def raw_forward(x): # exactly (per-model budgets collapse to remaining()). wrapped_models = list(synthesize_models_from_specs([sr]).values()) n_models = len(wrapped_models) + if n_models == 0: + raise RuntimeError( + f"No wrapped models were synthesized for ONNX {args.onnx!r} and VNNLIB {args.vnnlib!r}" + ) if args.fuzzing_seconds > 0 and remaining() > 1.0: per_model_fuzz = min(args.fuzzing_seconds, remaining() / n_models) for wm in wrapped_models: if remaining() <= 1.0: break - try: - ce, _ = pgd_preattack(wm, sr[3], min(per_model_fuzz, remaining()), args.fuzzing_scale) - except Exception as exc: - ce = None - print(f"[attack skipped] {exc}", file=sys.stderr) + ce, _ = pgd_preattack( + wm, sr[3], min(per_model_fuzz, remaining()), args.fuzzing_scale + ) if ce is not None: x = ce.input if hasattr(ce, "input") else ce if io_decls is None: - write_vnncomp_result(args.output, "unknown") - else: - write_vnncomp_result(args.output, "sat", x=x, y=raw_forward(x), - in_decl=io_decls[0], out_decl=io_decls[1]) + raise RuntimeError( + f"Cannot write a counterexample because {args.vnnlib!r} has no VNNLIB 2.0 I/O declarations" + ) + write_vnncomp_result(args.output, "sat", x=x, y=raw_forward(x), + in_decl=io_decls[0], out_decl=io_decls[1]) return if remaining() <= 1.0: @@ -149,7 +149,7 @@ def _verify(tier, budget): # certification permanently impossible for the run. cfg.branching_method = "width" cfg.multi_split_levels = 1 - cfg.reuse_root_bounds = False + cfg.root_bounds_reuse = "none" cfg.intermediate_refine = "none" cfg.frontier_cap = 0 clear_violation_check_module_cache() @@ -182,30 +182,26 @@ def _verify(tier, budget): # the clock -> 'timeout'. statuses = [] falsified_ce = None - try: - unfinished = n_models - for wm in wrapped_models: - if remaining() <= 1.0: - break - deadline = time.time() + remaining() / max(1, unfinished) - res = _verify_wrapped(wm, deadline) - unfinished -= 1 - statuses.append(res.status) - if res.status == VerifyStatus.FALSIFIED and res.counterexample is not None: - falsified_ce = res.counterexample - break - except Exception as exc: - print(f"[verify error] {exc}", file=sys.stderr) - write_vnncomp_result(args.output, "unknown") - return + unfinished = n_models + for wm in wrapped_models: + if remaining() <= 1.0: + break + deadline = time.time() + remaining() / max(1, unfinished) + res = _verify_wrapped(wm, deadline) + unfinished -= 1 + statuses.append(res.status) + if res.status == VerifyStatus.FALSIFIED and res.counterexample is not None: + falsified_ce = res.counterexample + break if falsified_ce is not None: if io_decls is None: - write_vnncomp_result(args.output, "unknown") - else: - write_vnncomp_result(args.output, "sat", x=falsified_ce, - y=raw_forward(falsified_ce), - in_decl=io_decls[0], out_decl=io_decls[1]) + raise RuntimeError( + f"Cannot write a counterexample because {args.vnnlib!r} has no VNNLIB 2.0 I/O declarations" + ) + write_vnncomp_result(args.output, "sat", x=falsified_ce, + y=raw_forward(falsified_ce), + in_decl=io_decls[0], out_decl=io_decls[1]) elif len(statuses) == n_models and all(s == VerifyStatus.CERTIFIED for s in statuses): write_vnncomp_result(args.output, "unsat") elif any(s not in (VerifyStatus.CERTIFIED, VerifyStatus.TIMEOUT) for s in statuses): @@ -214,7 +210,7 @@ def _verify(tier, budget): write_vnncomp_result(args.output, "timeout") -def main() -> None: +def main() -> int: ap = argparse.ArgumentParser(description="ACT VNN-COMP 2026 single-instance runner") ap.add_argument("onnx") ap.add_argument("vnnlib") @@ -242,8 +238,16 @@ def main() -> None: "input-domain splitting with full per-node bound recomputation " "(the ACAS Xu regime); 0 disables the profile") args = ap.parse_args() - run_vnncomp_instance(args) + try: + run_vnncomp_instance(args) + except BaseException: + traceback.print_exc(file=sys.stderr) + from act.front_end.vnnlib_loader.vnnlib_parser import write_vnncomp_result + + write_vnncomp_result(args.output, "unknown") + return 1 + return 0 if __name__ == "__main__": - main() + raise SystemExit(main())