f-dangel / curvlinops / 23986024830

91%

Build Type

push

github

Committed by web-flow

Commit Message

[OPT] Make `HessianLinearOperator` compile-friendly and benchmark (#286)

* [ADD] Compile-friendly HessianLinearOperator with benchmark

Replace `cached_property` with eager `_init_mp()` in
`CurvatureLinearOperator` and all subclasses to eliminate
`torch.compile` graph breaks for `HessianLinearOperator`.

Add `trace_gradient_and_loss` and `make_compiled_gradient_and_loss`
to `curvlinops.examples`: uses `make_fx` to capture `autograd.grad`
backward as explicit forward ops, then compiles. Benchmark measures
both eager and compiled matvec/memory for compilable operators
(stored in same JSON). Compiled gradient reference enables fair
comparison.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [REF] Remove defensive guards in compiled visualization

The compiled reference and operator data must always be present.
Remove None guards so missing data crashes instead of silently
producing empty plots.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [DOC] Overlay eager/compiled bars with opacity; add reference axis

Plot compiled bars on top of eager bars with alpha=0.7 instead of
side-by-side. Use add_gradient_reference for the eager reference
(adds secondary axis showing multiples of gradient time). Compiled
gradient reference shown as dotted line.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [REF] Use torch.testing.assert_close instead of custom recursive comparison

Replaces hand-rolled _assert_allclose with torch.testing.assert_close
which handles nested structures (tuples, dicts, lists) and catches
length mismatches.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [TST] Use two batches of different size in Hessian compile test

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [REF] Extract _setup_problem in compile tests

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* [REF] Move data-to-device into compiled... (continued)

Coverage Stats

35 of 44 new or added lines in 6 files covered. (79.55%)

30 existing lines in 6 files now uncovered.

2406 of 2639 relevant lines covered (91.17%)

0.91 hits per line