1 files changed, 69 insertions, 0 deletions
diff --git a/tests/test_jlens_v3.py b/tests/test_jlens_v3.py
new file mode 100644
index 0000000..5ac3f67
--- /dev/null
+++ b/tests/test_jlens_v3.py
@@ -0,0 +1,69 @@
+"""Minimal tests for the faithful J-lens (jlens_v3) correctness.
+
+Run inside the meru container where torch exists:
+    python3 tests/test_jlens_v3.py
+Skips gracefully when torch is unavailable.
+
+Key assertion: at the LAST layer the residual-to-residual Jacobian is the
+identity, so the faithful J-lens vectors must equal the unembedding rows
+(cosine similarity ~1.0). This is the ruler-check that validates the
+W_U-probed VJP machinery.
+"""
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
+
+try:
+    import numpy as np
+    import torch
+    import jlens_v3
+except ImportError:
+    print("SKIP: torch not available (run inside the meru container)")
+    sys.exit(0)
+
+from model import GPT, GPTConfig  # noqa: E402
+
+DEV = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+def _tiny_model(n_layer=2, d=16, block=32):
+    torch.manual_seed(0)
+    cfg = GPTConfig(n_layer=n_layer, n_head=2, n_embd=d, block_size=block,
+                    bias=False, vocab_size=65, dropout=0.0)
+    return GPT(cfg).eval()
+
+
+def _tiny_data(n=2000, block=32):
+    return np.random.RandomState(0).randint(0, 65, n).astype(np.uint16)
+
+
+def test_shapes_and_finite():
+    model = _tiny_model()
+    data = _tiny_data()
+    batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
+    fv = jlens_v3.compute_faithful_jlens(model, 1, batches, DEV, chunk=8)
+    assert fv.shape == (65, 16), f"faithful vecs shape {fv.shape}"
+    assert torch.isfinite(fv).all(), "faithful vecs not finite"
+    pn = jlens_v3.compute_proxy_norms(model, 1, batches, DEV, 65, chunk=16)
+    assert all(np.isfinite(v) for v in pn.values()), "proxy norms not finite"
+
+
+def test_last_layer_identity():
+    """At the final layer J must be identity: faithful vectors == W_U rows."""
+    n_layer, d = 2, 16
+    model = _tiny_model(n_layer=n_layer, d=d)
+    data = _tiny_data()
+    batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
+    fv = jlens_v3.compute_faithful_jlens(model, n_layer - 1, batches, DEV,
+                                         chunk=8)
+    wu = model.lm_head.weight.detach().float()
+    sims = torch.nn.functional.cosine_similarity(fv.float(), wu.cpu().float(),
+                                                 dim=1)
+    assert sims.mean().item() > 0.99, f"identity check failed: {sims.mean():.4f}"
+
+
+if __name__ == '__main__':
+    test_shapes_and_finite()
+    test_last_layer_identity()
+    print(f"PASS ({DEV})")