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

---