Add faithful J-lens (jlens_v3): W_U-probed residual Jacobian per paper; both-ways comparison vs log-softmax proxy; 3-model adversarial reviews

1 files changed, 32 insertions, 0 deletions

diff --git a/smoke_test_v3.py b/smoke_test_v3.py
new file mode 100644
index 0000000..7e83e1c
--- /dev/null
+++ b/smoke_test_v3.py
@@ -0,0 +1,32 @@
+"""CPU smoke test for jlens_v3 (faithful J-lens machinery). Random tiny model."""
+import sys, os
+sys.path.insert(0, '.')
+sys.path.insert(0, 'src')
+import torch, numpy as np
+from model import GPT, GPTConfig
+import jlens_v3
+
+torch.manual_seed(0)
+
+cfg = GPTConfig(n_layer=3, n_head=4, n_embd=16, block_size=32, bias=False,
+                vocab_size=65, dropout=0.0)
+model = GPT(cfg).eval()
+
+data = np.random.randint(0, 65, 2000).astype(np.uint16)
+batches = jlens_v3.make_batches(data, 32, 4, 3, 'cpu')
+
+J = jlens_v3.compute_faithful_jlens(model, 1, batches, 'cpu', chunk=8)
+print("faithful vecs shape:", J.shape, "finite:", torch.isfinite(J).all().item())
+print("row norms (first 5):", [round(v, 4) for v in J.norm(dim=1)[:5].tolist()])
+
+pn = jlens_v3.compute_proxy_norms(model, 1, batches, 'cpu', 65)
+print("proxy norms finite:", all(np.isfinite(v) for v in pn.values()))
+
+fn = {k: J[k].norm().item() for k in range(65)}
+freq = np.bincount(data, minlength=65).astype(float)
+freq = freq / freq.sum() * 100
+r_p = np.corrcoef(np.array([pn[k] for k in range(65)]), freq)[0, 1]
+r_f = np.corrcoef(np.array([fn[k] for k in range(65)]), freq)[0, 1]
+print(f"random-init sanity: proxy r={r_p:+.3f} faithful r={r_f:+.3f} (should be ~0, model is random)")
+assert J.shape == (65, 16) and torch.isfinite(J).all()
+print("SMOKE TEST PASSED")