f-dangel / curvlinops / 23211349482

91%

push

github

[ADD] Support callable `model_func` in (E)KFAC FX backend (#271)

Allow KFACLinearOperator and EKFACLinearOperator to accept a plain
callable (params_dict, X) -> prediction as model_func with
backend="make_fx", instead of requiring an nn.Module.

When model_func is a callable, params must be dict[str, Tensor].
Parameter groups are built from the IO collector's detections rather
than from Module traversal.

Add SUPPORTS_FUNCTIONAL class attribute to _EmpiricalRiskMixin (default
False, True for MakeFxKFACComputer/MakeFxEKFACComputer). Validate
model_func type (Module or callable) and params format (list or dict).

Co-authored-by: f-dangel <felix.dangel@concordia.ca>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

18 of 18 new or added lines in 3 files covered. (100.0%)

2345 of 2566 relevant lines covered (91.39%)

0.91 hits per line