Time the stages of the Hessian-inverse pipeline

[a-rahimi]

Summary

Stacks on #2 (now merged). Adds lightweight wall-clock instrumentation so each stage of the Hessian-inverse pipeline is visible:

• src/timing.py (new, stdlib-only): a module-level registry with a record(name) context manager (time.perf_counter), reset(), and report() which renders an aligned table sorted by total time.
• src/hessian.py: stages setup/derivatives (the functorch pass), solve/assemble-K, and per-solver stages — splu: to-csc / factorize / substitute / pack; block: pivot / factorize / substitute (Tridiagonal.solve spelled out verbatim at the call site so the two phases time separately; block_partitioned_matrices.py untouched).
• src/train_newton.py: train/gradient and train/trs-subproblem stages; prints the summary table to stderr at the end of training.

No numerical behavior changes — instrumentation is context managers around existing statements.

Example (trust-region, --tr-solver efficient --batch-size 32 --num-steps 5)

stage                   calls    total s    mean ms
train/trs-subproblem        5     24.281    4856.28
solve/splu/factorize       56     10.039     179.27
setup/derivatives           5      7.077    1415.39
solve/splu/to-csc          56      5.287      94.42
solve/assemble-K           56      1.297      23.16
solve/splu/substitute      56      0.165       2.95
solve/splu/pack           112      0.020       0.18
train/gradient             10      0.002       0.25

Already actionable: the secular loop runs ~11 solves per step, and to-csc + assemble-K (~118 ms/solve) are recomputed per trial λ even though only the diagonal changes — caching the CSC structure and patching the diagonal would cut ~30% off the subproblem.

Tests

137 passed (133 + 4 new for the timing registry).

🤖 Generated with Claude Code