google-deepmind / kfac-jax

Second Order Optimization and Curvature Estimation with K-FAC in JAX.
Apache License 2.0
235 stars 18 forks source link

Remove the `TwoKroneckerFactored` class and use the `KroneckerFactored` class instead. #257

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Remove the TwoKroneckerFactored class and use the KroneckerFactored class instead.

This is a no-op change. It just simplifies the code and makes it easier to add support for more than two parameters.