From 01a070620991f9c5ec25898cf2c936497abbf3bb Mon Sep 17 00:00:00 2001
From: Ronald Tse <ronald.tse@gmail.com>
Date: Tue, 1 Sep 2026 07:44:27 +0200
Subject: [PATCH] =?UTF-8?q?feat(parity):=20per-pair=20flip-position=20dump?=
 =?UTF-8?q?=20=E2=80=94=20the=20unit=20the=20flip=20bootstrap=20needs?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

run_margin_analysis(dump_positions=...) writes one JSONL record per
pair (tokens, flip positions, their reference margins) beside the
aggregate margins JSON. MarginReport stays schema-identical; the dump
is diagnostic. Enables the TODO.training-work/05 shipped-vs-head32
flip CIs; aggregate flip counts in the dump must reconcile with the
report (tested).
---
 src/imf/parity.py        | 16 +++++++++++++++-
 tests/test_imf_parity.py | 16 ++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/imf/parity.py b/src/imf/parity.py
index 02521c0..c9b98e5 100644
--- a/src/imf/parity.py
+++ b/src/imf/parity.py
@@ -253,7 +253,7 @@ def _onnx_forced_logits(enc_sess, dec_sess, source: str, target_ids: list[int]):
     return dict(zip(out_names, out, strict=True))["logits"][0]
 
 
-def run_margin_analysis(model, zip_path, pairs, max_len: int = 256) -> MarginReport:
+def run_margin_analysis(model, zip_path, pairs, max_len: int = 256, dump_positions: Path | str | None = None) -> MarginReport:
     """pairs: iterable of (source_text, gold_target) — the same probe set
     the CER parity gate uses. Teacher-forces both sides and measures the
     argmax flip rate, reference top1−top2 margin quantiles, and KL
@@ -286,6 +286,20 @@ def run_margin_analysis(model, zip_path, pairs, max_len: int = 256) -> MarginRep
 
     margins = np.concatenate(margin_chunks)
     flips = np.concatenate(flip_chunks)
+    if dump_positions is not None:
+        # per-position flip vectors, pair-indexed — the unit the
+        # paired bootstrap needs for shipped-vs-rebuild comparisons
+        import json as _json
+
+        dump = Path(dump_positions)
+        dump.parent.mkdir(parents=True, exist_ok=True)
+        with dump.open("w", encoding="utf-8") as out:
+            for pair_idx, (m, f) in enumerate(zip(margin_chunks, flip_chunks)):
+                out.write(_json.dumps(
+                    {"pair": pair_idx, "tokens": int(f.size),
+                     "flip_positions": [int(x) for x in np.nonzero(f)[0]],
+                     "flip_margins": [round(float(m[x]), 4) for x in np.nonzero(f)[0]]}
+                ) + "\n")
     p1, p10, p50 = (float(np.quantile(margins, q)) for q in (0.01, 0.10, 0.50))
     n_flips = int(flips.sum())
     return MarginReport(
diff --git a/tests/test_imf_parity.py b/tests/test_imf_parity.py
index 3510cef..56a57a6 100644
--- a/tests/test_imf_parity.py
+++ b/tests/test_imf_parity.py
@@ -166,3 +166,19 @@ def test_margin_report_json_roundtrip(gated_zip: Path, fixture_model, tmp_path:
     }
     assert data["samples"] == report.samples
     assert data["flip_rate"] == report.flip_rate
+
+
+def test_margin_analysis_dumps_per_pair_positions(gated_zip: Path, fixture_model, tmp_path: Path) -> None:
+    """TODO.training-work/05: the flip bootstrap needs per-pair token and
+    flip counts, not just aggregates."""
+    dump = tmp_path / "positions.jsonl"
+    report = run_margin_analysis(
+        fixture_model, gated_zip, PAIRS, max_len=12, dump_positions=dump,
+    )
+    rows = [json.loads(line) for line in dump.read_text().splitlines() if line]
+    assert rows and len(rows) <= len(PAIRS)
+    assert all({"pair", "tokens", "flip_positions", "flip_margins"} <= set(r) for r in rows)
+    total_flips = sum(len(r["flip_positions"]) for r in rows)
+    assert total_flips == report.flipped_tokens
+    total_tokens = sum(r["tokens"] for r in rows)
+    assert total_tokens == report.tokens