f-dangel / curvlinops

scipy linear operators for the Hessian, Fisher/GGN, and more in PyTorch
https://curvlinops.readthedocs.io/en/latest/
MIT License
17 stars 5 forks source link

Allow for matrix-matrix and matrix-vector products with `KFACLinearOperator` and `KFACInverseLinearOperator` without converting to numpy #91

Closed runame closed 6 months ago

runame commented 6 months ago

Addresses #71, but only for the KFACLinearOperator and KFACInverseLinearOperator.

KFACLinearOperator/KFACInverseLinearOperator is arguably the only linear operator here that is likely to be used for preconditioned-gradient methods for large scale neural networks. Therefore, adding torch_matmat and torch_matvec methods seem like a simple solution to avoid unnecessary device transfers, which are a bottleneck for this use case. However, this doesn't address the issue in general.

coveralls commented 6 months ago

Pull Request Test Coverage Report for Build 8360156180

Details


Changes Missing Coverage Covered Lines Changed/Added Lines %
curvlinops/inverse.py 16 17 94.12%
curvlinops/kfac.py 36 41 87.8%
<!-- Total: 52 58 89.66% -->
Totals Coverage Status
Change from base Build 8250604944: 0.2%
Covered Lines: 1017
Relevant Lines: 1138

💛 - Coveralls