f-dangel / curvlinops / 23122505326

91%

push

github

[ADD] Weight tying support for (E)KFAC make_fx backend (#264)

* [ADD] Weight tying support for KFAC make_fx backend

Support parameters used in multiple affine operations (weight tying)
in the FX-based KFAC backend. Per-usage inputs/gradients are formatted
individually, then concatenated along the shared dimension before
computing covariances, so the normalization is correct.

Also replace `append_ones_for_bias` with `bias_pad: int | None` in
`input_to_weight_sharing_format` to support weight-tied layers where
some usages have bias and others don't.

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

* [DOC] Add PR link to changelog entry

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

* [ADD] Weight tying support for EKFAC make_fx backend

Restructure eigenvalue correction in MakeFxEKFACComputer to collect
per-usage formatted tensors and concatenate along the shared dimension
before computing corrections, matching the KFAC backend approach.

Extract SplitConcatModel to test/utils.py for reuse. Add EKFAC
weight-tying test and parametrize both tests by loss reduction.

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

* [FIX] Address Copilot review: bias validation and formatting

- Filter NOT_A_PARAM from bias conflict check (could cause false positives)
- Add check that a tracked bias is not reused across different weights
- Extract _verify_weight_bias_compatibility to reduce complexity
- Fix line-too-long formatting in ekfac.py

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

* [FIX] Remove overly restrictive bias-reuse check

A bias appearing with multiple weights is valid when weights and biases
are treated separately (each becomes its own parameter group).

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

* [TST] Parametrize weight tying tests by bias and separate_weight_and_bias

Add bias parameter to SplitConcatModel and parametrize KFAC/EK... (continued)

53 of 55 new or added lines in 4 files covered. (96.36%)

2459 of 2689 relevant lines covered (91.45%)

0.91 hits per line