From f22bbeb0d2f7202922eb7039cc99040f7fc2ff7a Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 10 Jul 2026 12:52:18 +0800
Subject: [PATCH] parity(U1): cached-J pullback vs direct VJP, all layers
 cos>0.999

min cos 0.999801 (layer 8), rising to 0.999996; gap is fp16 cache
storage as expected. Gate PASS.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
---
 artifacts/parity_u1.txt | 24 +++++++++++++
 scripts/parity_u1.py    | 77 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 artifacts/parity_u1.txt
 create mode 100644 scripts/parity_u1.py

diff --git a/artifacts/parity_u1.txt b/artifacts/parity_u1.txt
new file mode 100644
index 0000000..fdbf86e
--- /dev/null
+++ b/artifacts/parity_u1.txt
@@ -0,0 +1,24 @@
+U1 parity: cached-Jacobian pullback vs direct VJP (word=happy/joy)
+model=Qwen/Qwen3-0.6B  prompts=8  layers=[8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
+
+|   layer |       cos |      |vA| |      |vB| |
+|--------:|----------:|----------:|----------:|
+|       8 | +0.999801 | +1.000000 | +1.000000 |
+|       9 | +0.999862 | +1.000000 | +1.000000 |
+|      10 | +0.999899 | +1.000000 | +1.000000 |
+|      11 | +0.999904 | +1.000000 | +1.000000 |
+|      12 | +0.999923 | +1.000000 | +1.000000 |
+|      13 | +0.999940 | +1.000000 | +1.000000 |
+|      14 | +0.999953 | +1.000000 | +1.000000 |
+|      15 | +0.999967 | +1.000000 | +1.000000 |
+|      16 | +0.999976 | +1.000000 | +1.000000 |
+|      17 | +0.999982 | +1.000000 | +1.000000 |
+|      18 | +0.999985 | +1.000000 | +1.000000 |
+|      19 | +0.999990 | +1.000000 | +1.000000 |
+|      20 | +0.999992 | +1.000000 | +1.000000 |
+|      21 | +0.999994 | +1.000000 | +1.000000 |
+|      22 | +0.999996 | +1.000000 | +1.000000 |
+|      23 | +0.999996 | +1.000000 | +1.000000 |
+|      24 | +0.999996 | +1.000000 | +1.000000 |
+
+min cos = +0.999801   GATE (>0.999): PASS
diff --git a/scripts/parity_u1.py b/scripts/parity_u1.py
new file mode 100644
index 0000000..f33e510
--- /dev/null
+++ b/scripts/parity_u1.py
@@ -0,0 +1,77 @@
+"""U1 parity gate: cached-Jacobian pullback == direct VJP, per layer.
+
+(authored by Claude)
+
+The two paths are linear-identical:  mean_p(J_p)^T w == mean_p(J_p^T w).
+Path A pulls the word cotangent through the CACHED pooled Jacobian.
+Path B contracts the same cotangent inside per-prompt backward passes.
+The only expected gap is fp16 storage in the cache, so per-layer cosine must
+exceed 0.999. A failure is a WIRING bug (layer index, position mask, pooling),
+not a threshold to tune.
+
+Run:
+    uv run python scripts/parity_u1.py 2>&1 | tee /tmp/claude-1000/jsteer_parity.log
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+from loguru import logger
+from tabulate import tabulate
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from jsteer import Jacobian, word_vector_vjp
+
+# Claude: repo root on path so `scripts.smoke` imports whether run as a file or -m.
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from scripts.smoke import CACHE, DEVICE, DTYPE, MODEL, PROMPTS  # same inputs  # noqa: E402
+
+WORDS = ["happy", "joy"]
+OUT = "artifacts/parity_u1.txt"
+
+
+def _unit_dir(vec, layer: int) -> torch.Tensor:
+    """Pull the per-layer unit direction out of a steering_lite Vector."""
+    return vec.stacked[layer]["v"].squeeze(0).float().cpu()  # [d]
+
+
+def main() -> None:
+    logger.info(f"loading {MODEL} ({DTYPE}) on {DEVICE}")
+    tok = AutoTokenizer.from_pretrained(MODEL)
+    model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=DTYPE).to(DEVICE).eval()
+
+    jac = Jacobian.load(CACHE)
+    layers = jac.layers  # exact int layers fitted by the smoke
+    logger.info(f"cached layers={layers}")
+
+    # Path A: cached pooled Jacobian pullback.
+    vA = jac.word_vector(model, tok, WORDS)
+    # Path B: direct per-prompt VJP over the SAME prompts / layers / skip_first / max_length.
+    vB = word_vector_vjp(model, tok, PROMPTS, WORDS, layers=layers, max_length=128)
+
+    rows = []
+    for l in layers:
+        a, b = _unit_dir(vA, l), _unit_dir(vB, l)
+        cos = float(torch.dot(a, b) / (a.norm() * b.norm()))
+        rows.append((l, cos, float(a.norm()), float(b.norm())))
+
+    table = tabulate(rows, headers=["layer", "cos", "|vA|", "|vB|"],
+                     tablefmt="pipe", floatfmt="+.6f")
+    min_cos = min(r[1] for r in rows)
+    gate = "PASS" if min_cos > 0.999 else "FAIL"
+    report = f"{table}\n\nmin cos = {min_cos:+.6f}   GATE (>0.999): {gate}\n"
+
+    logger.info("U1 parity table:\n" + report)
+    with open(OUT, "w") as f:
+        f.write("U1 parity: cached-Jacobian pullback vs direct VJP (word=happy/joy)\n")
+        f.write(f"model={MODEL}  prompts={len(PROMPTS)}  layers={layers}\n\n")
+        f.write(report)
+    logger.info(f"wrote {OUT}")
+    if gate == "FAIL":
+        raise SystemExit(f"U1 parity FAILED: min cos={min_cos:.6f} <= 0.999 (wiring bug)")
+
+
+if __name__ == "__main__":
+    main()