Fix BGE-small export compatibility (#263)

* Fix BGE-small export compatibility

The documented BGE flow (scripts/download_bge_small_en_v1_5.py + the bge
examples) is broken on a fresh download: current optimum exports produce a
graph the tracer rejects at load with 'Unimplemented ONNX operator: Sqrt'.

Root causes (export drift):
- Fused LayerNormalization (opset >= 17) decomposes under tract into
  Div(x, Sqrt(var+eps)); tract only re-fuses Recip(Sqrt(.)) into Rsqrt (see
  declutter_recip), and the tracer has an Rsqrt handler but no Sqrt handler.
- The export now has two graph outputs (token_embeddings, sentence_embedding);
  the prover seeds an opening only for outputs()[0] (prover.rs output_claim),
  so the second output is never opened (EvalReduction EmptyInput).
- The export takes two inputs now (token_type_ids dropped), so the examples'
  three-input call panics with an out-of-bounds index at execution.

Fix, in the download script's normalize_graph():
- Pin --opset 14 (keeps LayerNorm decomposed) and rewrite
  Div(x, Sqrt(v)) -> Mul(x, Reciprocal(Sqrt(v))) so tract fuses to Rsqrt.
- Keep token_embeddings (the encoder's last hidden state) as the sole output
  and dead-node eliminate the pooling / L2-norm subgraph that only fed
  sentence_embedding. Pooling + normalization are a cheap public post-step.
- Update bge.rs / bge_generate.rs to the two-input signature.

Verified: scripts/download + 'cargo run -p jolt-atlas-core --example bge'
loads, traces, proves, and verifies end to end.

Note: keeping the full model's sentence_embedding output instead hits a
separate degenerate-instance panic in the Shout lookup prover (scalar Rsqrt
over the single pooled element -> RaPolynomial stuck in Round1 ->
final_claim); filed separately.

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Khalil Gibran Hassam <59589957+Forpee@users.noreply.github.com>

Author: 3 people

atlas-onnx-tracer/examples/bge_generate.rs

@@ -13,8 +13,7 @@
 /// | Index | Name            | Shape            | Description                           |
 /// |-------|-----------------|------------------|---------------------------------------|
 /// | 0     | input_ids       | [1, seq_len]     | Token IDs (Gather indices)            |
-/// | 1     | token_type_ids  | [1, seq_len]     | Segment IDs (all zeros for single)    |
-/// | 2     | attention_mask  | [1, seq_len]     | Binary mask (1 = attend)              |
+/// | 1     | attention_mask  | [1, seq_len]     | Binary mask (1 = attend)              |
 use atlas_onnx_tracer::model::{Model, RunArgs};
 
 fn main() {

jolt-atlas-core/examples/bge.rs

@@ -46,11 +46,7 @@ fn main() {
     let input_ids_data: Vec<i32> = (0..seq_len).map(|_| rng.gen_range(0..vocab_size)).collect();
     let input_ids = Tensor::new(Some(&input_ids_data), &[1, seq_len]).unwrap();
 
-    // Input 1: token_type_ids — segment IDs in {0, 1}. Use all zeros for single-segment input.
-    let token_type_ids_data: Vec<i32> = vec![0; seq_len];
-    let token_type_ids = Tensor::new(Some(&token_type_ids_data), &[1, seq_len]).unwrap();
-
-    // Input 2: attention_mask — all 1s (attend everywhere), in BERT-style binary form.
+    // Input 1: attention_mask — all 1s (attend everywhere), in BERT-style binary form.
     let attention_mask_data: Vec<i32> = vec![1; seq_len];
     let attention_mask = Tensor::new(Some(&attention_mask_data), &[1, seq_len]).unwrap();
 
@@ -61,7 +57,7 @@ fn main() {
     let timing = std::time::Instant::now();
     let (proof, io, _debug_info) = ONNXProof::<Fr, Blake2bTranscript, HyperKZG<Bn254>>::prove(
         &prover_preprocessing,
-        &[input_ids, token_type_ids, attention_mask],
+        &[input_ids, attention_mask],
     );
     println!("Proof generation took {:.2?}", timing.elapsed());
 

scripts/download_bge_small_en_v1_5.py

@@ -43,6 +43,8 @@ def export_model():
             "optimum.exporters.onnx",
             "--model",
             MODEL_ID,
+            "--opset",
+            "14",
             str(MODEL_DIR),
         ],
     )
@@ -53,6 +55,85 @@ def export_model():
     print("Export complete.")
 
 
+def normalize_graph():
+    """Rewrite export patterns the Atlas tracer/prover does not support.
+
+    The stock optimum feature-extraction export is incompatible with the Atlas
+    tracer in three ways; all are handled here so the documented flow works
+    end to end (load, trace, prove, verify).
+
+    1. Fused ``LayerNormalization`` (opset >= 17) decomposes under tract into
+       ``Div(x, Sqrt(var + eps))``, which tract cannot re-fuse: it only rewrites
+       ``Recip(Sqrt(.))`` to ``Rsqrt`` (see tract's ``declutter_recip``), and the
+       tracer has an ``Rsqrt`` handler but no ``Sqrt`` handler. We pin
+       ``--opset 14`` at export (keeps LayerNorm decomposed) and rewrite
+       ``Div(x, Sqrt(v))`` into ``Mul(x, Reciprocal(Sqrt(v)))`` so tract fuses it.
+
+    2. The export has two graph outputs: ``token_embeddings`` (the encoder's
+       last hidden state) and ``sentence_embedding`` (mean/CLS pooled and
+       L2-normalized). The Atlas prover seeds an opening only for
+       ``outputs()[0]`` (see prover.rs ``output_claim``), so a second output is
+       never opened. We keep ``token_embeddings`` as the sole output -- it is
+       the full 12-layer encoder result, and pooling + normalization are a
+       cheap public post-step -- then dead-node eliminate the pooling / L2-norm
+       subgraph that only fed ``sentence_embedding``.
+
+    Dropping that subgraph also removes its CLS-pooling ``Gather(axis=1)`` (the
+    tracer's Gather is axis-0 only) and its final scalar ``Rsqrt`` over a
+    single pooled element, which triggers a degenerate-instance panic in the
+    Shout lookup prover (see the referenced issue).
+    """
+    import onnx
+    from onnx import helper
+
+    path = MODEL_DIR / "model.onnx"
+    model = onnx.load(str(path))
+    g = model.graph
+
+    # -- 1. Div(x, Sqrt(v)) -> Mul(x, Reciprocal(Sqrt(v))) ------------------
+    sqrt_outs = {n.output[0] for n in g.node if n.op_type == "Sqrt"}
+    nodes, patched_div = [], 0
+    for n in g.node:
+        if n.op_type == "Div" and len(n.input) == 2 and n.input[1] in sqrt_outs:
+            recip_out = n.output[0] + "_recip"
+            nodes.append(helper.make_node(
+                "Reciprocal", [n.input[1]], [recip_out], name=n.name + "_recip"))
+            nodes.append(helper.make_node(
+                "Mul", [n.input[0], recip_out], [n.output[0]], name=n.name + "_mul"))
+            patched_div += 1
+        else:
+            nodes.append(n)
+
+    # -- 2. Single output = token_embeddings, then dead-node elimination ----
+    keep_output = "token_embeddings"
+    pruned_outputs = 0
+    if any(o.name == keep_output for o in g.output) and len(g.output) > 1:
+        for o in [o for o in g.output if o.name != keep_output]:
+            g.output.remove(o)
+            pruned_outputs += 1
+
+    producer = {out: n for n in nodes for out in n.output}
+    live = set()
+    frontier = [o.name for o in g.output]
+    while frontier:
+        name = frontier.pop()
+        n = producer.get(name)
+        if n is None or id(n) in live:
+            continue
+        live.add(id(n))
+        frontier.extend(n.input)
+    kept = [n for n in nodes if id(n) in live]
+    removed_nodes = len(nodes) - len(kept)
+
+    del g.node[:]
+    g.node.extend(kept)
+    onnx.save(model, str(path))
+    print(
+        f"Graph normalized: {patched_div} layernorm div(s) rewritten, "
+        f"{pruned_outputs} output(s) pruned, {removed_nodes} dead node(s) removed."
+    )
+
+
 def rename_network():
     """Rename model.onnx -> network.onnx for compatibility."""
     src = MODEL_DIR / "model.onnx"
@@ -70,6 +151,7 @@ def rename_network():
 def main():
     ensure_packages()
     export_model()
+    normalize_graph()
     rename_network()
     print(f"\nOK: {MODEL_ID} ONNX model is ready at {MODEL_DIR}")