f-dangel / curvlinops / 23988044085

91%

push

github

[OPT] Make GGN and MC-Fisher compile-friendly and benchmark (#287)

* [ADD] Make GGN and MC-Fisher compile-friendly (0 graph breaks)

Exact GGN and EF already had 0 graph breaks after the _init_mp
refactor. MC-Fisher had graph breaks from torch.Generator.

Fix: remove Generator from the MC sampling path, use global RNG
instead. Seed with torch.manual_seed in __matmul__ (before dynamo
traces). The generator parameter in ggn_utils is kept optional
(None = global RNG) for backward compatibility.

Add compile tests for GGN (exact), EF, and MC-Fisher.

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

* [REF] Inline vmap call for batched_grad_output_fn

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

* [FIX] Use fork_rng in GGN.__matmul__ to avoid RNG side effects

Replace bare torch.manual_seed with fork_rng + manual_seed to
isolate the RNG state for MC sampling. Placed in __matmul__
(not _matmat) because _matmat is inside torch.compile's trace.

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

* [REF] Parametrize GGN compile test; remove unrelated EF test

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

* [REF] Use direct import: from pytest import mark

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

* [REF] Use direct imports for manual_seed and fork_rng

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

* [DOC] Improve note on RNG in make_batch_ggn_mc_vector_product

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

* [FIX] Move fork_rng seeding from __matmul__ to _matmat

Seeding in __matmul__ missed internal _matmat callers like
_ChainPyTorchLinearOperator. Move fork_rng + manual_seed into
_matmat so all code paths get deterministic MC samples.

fork_rng causes 2 graph breaks (context manager not traceable),
but the per-batch computation is still compiled. Test asserts
breaks are from fork_rng and verifie... (continued)

11 of 11 new or added lines in 1 file covered. (100.0%)

2407 of 2640 relevant lines covered (91.17%)

0.91 hits per line