feat(sparse_probing): add leakage-safe k-sparse probing over activation tensors

[janmenjayap]

---

Description

Adds transformer_lens/tools/analysis/sparse_probing.py, a dependency-free, model-free module for fitting binary k sparse probes to a supplied activation matrix X and label vector y, following the method in Gurnee et al., "Finding Neurons in a Haystack: Case Studies with Sparse Probing" (TMLR 2023).

TransformerLens already exposes activations via run_with_cache, but has no maintained probing primitive: users currently rebuild splitting, feature selection, fitting, and control logic themselves, and it is easy to leak test information into that process. This PR adds a leakage-safe, dependency-free core for that workflow. Activation extraction, model wrappers, plotting, notebooks, multiclass probing, and optimal (MIP) selection are deliberately deferred to follow-up work.

Fixes #1728

---

What's included

• fit_sparse_probe(X, y, k, ...) — deterministic stratified train/test split computed before any learned statistic, train-only raw mean-difference feature selection, and a CPU-float64 Torch LBFGS balanced-logistic fit with explicit objective/gradient-norm convergence diagnostics (raises rather than silently returning an unconverged fit).
• sweep_sparse_probe(X, y, ks, ...) — reuses one fixed split across a k-grid and reports raw random-coordinate and shuffled-training-label control distributions alongside each k, without assigning an automatic "significance" label.
• SparseProbeResult / SparseProbeSweep result dataclasses reporting held-out accuracy, precision, recall, and F1 (F1 primary), selected indices/scores, coefficients/intercept, preprocessing metadata, split indices/class counts, and control distributions.
• A regression test asserting that MPS-resident features are moved to CPU before the float64 fit (float64 is unsupported on MPS).
• A jaxtyping>=0.3 compatibility fix: the rebase onto origin/dev pulled in CI Warnings Cleanup #1732, which now raises jaxtyping.TypeCheckError instead of letting BeartypeCallHintParamViolation propagate, and reclassifies float8_e4m3fn as a valid Float dtype at the annotation level. Updated test_runtime_typecheck_rejects_invalid_tensor_contracts to use tests/typecheck_errors.TYPECHECK_ERRORS, and moved the float8 case to this module's own ValueError dtype-rejection table, where it's now actually caught.
• Exports from transformer_lens.tools.analysis, plus docs/source/content/sparse_probing.md (contracts, claim boundaries, and a run_with_cache composition example) linked from docs/source/index.md.

---

Design decisions

Issue #1728 asked maintainers to weigh in on five open questions before implementation. This PR takes the issue's own recommended position on each:

1. Repository fit / model-free scope — kept model-free; extraction and plotting are out of scope for this PR.
2. Optimizer — uses an explicit CPU-float64 Torch LBFGS objective (no new sklearn dependency), and reports final gradient/objective metadata, raising on non-convergence.
3. Controls — random-coordinate and shuffled-training-label distributions are first-class sweep outputs, disabled only when their repeat count is zero.
4. Preprocessing — defaults to "none" (matching the paper's reference code); optional train-only "standardize" is documented as intentionally changing the L2 objective.
5. Device/dtype — score reductions run in at least float32 on the input device; only the selected [example, k] matrices move to CPU float64 for the deterministic LBFGS fit.

---

Non-goals (tracked as follow-ups)

Activation extraction and position-reduction helpers, model wrappers/downloads, multiclass/one-vs-rest probing, plotting and notebooks, an optimal/MIP selector, feature batteries, SAE-latent composition, and causal validation are all explicitly out of scope for this PR.

---

Claim boundaries

A high held-out F1 means the labeled feature is linearly decodable from the supplied activations — it does not by itself establish that the model uses that feature, that a selected coordinate is monosemantic, or that a smooth k-sweep curve is evidence of superposition. The docs and result contracts call this out explicitly, and the module exposes no .plot() method or automatic "significant" label.

---

Test plan

• uv run pytest tests/unit/tools/test_sparse_probing.py tests/mps/test_mps_basic.py -q — 68 passed.
• make test-pr (unit + docstring + acceptance + integration): unit 5,827 passed, docstring 18 passed, acceptance 209 passed, integration 1,459 passed / 1 failed. The one failure, test_granite_eager_scan_device_correctness[mps], is an unrelated Granite MoE Hybrid eager/fused-scan divergence — reproduced identically on a clean origin/dev checkout in an isolated worktree, confirming it's pre-existing on trunk and not introduced by this branch.
• make format and uv run mypy . clean.
• docs/source/index.md includes the new guide exactly once; docs build succeeds.

---

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

---

Checklist:

• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

---

[jlarson4]

Thanks for putting this together @janmenjayap. A few comments below:

[jlarson4]

The objective is mean-reduced and scale-invariant, but its gradient grows with |X|, so the absolute 1e-7 default rejects fits that match a float64 Newton solve to a few ulp: on n=1000, d=768 activations with per-coordinate std 60, every default fit_sparse_probe raises and sweep_sparse_probe dies on its first such fit. Because line 359 derives tolerance_change from gradient_tolerance**2, relaxing the tolerance also stops LBFGS earlier and max_iter changes nothing. Please make the acceptance test scale-relative (for example gradient_inf_norm <= gradient_tolerance * max(1.0, initial_gradient_inf_norm)), set tolerance_change=0.0, and add a guard test that fits the defaults on a matrix with a coordinate at std in the hundreds.

[jlarson4]

sweep_sparse_probe validates once with k=ks[-1] (line 668), so this line stamps that value on every result: for ks=[1, 2, 4] all three read k=4 while their supports are sized 1, 2, 4, and any (r.k, r.metrics.f1) curve comes out flat in k. Please pass the loop's k into _fit_result explicitly and assert tuple(r.k for r in sweep.results) == sweep.ks in the sweep test.

[jlarson4]

This range check is the only direct assertion on a returned metric, and every other metric check compares two paths mutated the same way: scoring the training rows instead of the held-out rows, permuting the labels the control arms are scored against, ignoring class_weight, holding the split seed constant, swapping precision and recall, and moving the decision threshold to 0.5 each leave all 53 tests green. Please add a test that recomputes the confusion counts from features[result.test_indices], labels[result.test_indices] and the returned coefficients and intercept and asserts equality with result.metrics, plus a shuffle-arm check that the scored labels are labels[test_indices].

[jlarson4]

SparseProbeResult has objective, gradient_inf_norm, iterations and function_evaluations, and a fit that misses the threshold raises rather than returning a flag, so result.converged is an AttributeError. Please replace "and convergence flag" with "and function-evaluation count" and say that a fit missing the threshold raises.

[jlarson4]

The row-level split does not know that several positions of one prompt share a label, so the common way of building features puts positions of the same prompt on both sides: on label-independent grouped data it scores 0.78 held-out accuracy where a group-held-out split scores 0.50. Please add one sentence here saying rows that share a source prompt must not straddle the split.