google-deepmind/kfac-jax v0.0.4

kfac_jax 0.0.4

Repository: google-deepmind/kfac-jax

Tag: v0.0.4

Published: 2023-04-27T12:32:13Z

Prerelease: no

Release notes:

What's Changed

• - Adding use_exact_inverses argument to optimizer by @copybara-service in https://github.com/deepmind/kfac-jax/pull/60
• Better PSM tests. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/66
• * Minor fix of pi-adjusted-inverse by @copybara-service in https://github.com/deepmind/kfac-jax/pull/71
• Fixing minor type error. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/76
• * Minor fix of how jax scopes are named (':' is not valid). by @copybara-service in https://github.com/deepmind/kfac-jax/pull/80
• - Improving docstring for optimizer. In particular regarding the damping parameter and LR/momentum/damping adaptation methods. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/84
• Changing examples code so that dataset functions directly return iterators instead of TF datasets. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/86
• - Removing unused put_stop_grad_on_loss_factor argument in multiply_fisher_factor by @copybara-service in https://github.com/deepmind/kfac-jax/pull/92
• Use jax.device_put_replicated to broadcast to local devices by @copybara-service in https://github.com/deepmind/kfac-jax/pull/89
• feat(ci): enable pip caching in CI by @SauravMaheshkar in https://github.com/deepmind/kfac-jax/pull/94
• Updates the code to always create variables and computations of the same dtype as the its inputs. Previously, if float64 was enabled, some of the results would be (potentially incorrectly) promoted to higher precision. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/82
• Adding data seen to the reported statistics on the evaluator in the examples. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/101
• - Modifying examples to only use label smoothing and L2 reg loss when training by @copybara-service in https://github.com/deepmind/kfac-jax/pull/102
• - Minor fixes/improvements to optimizer docstrings by @copybara-service in https://github.com/deepmind/kfac-jax/pull/108
• Replaces references to jax.numpy.DeviceArray with jax.Array. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/110
• Suppress pytype errors. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/112
• Reenable some types that were previously disabled due to pytype crashes. by @copybara-service in https://github.com/deepmind/kfac-jax/pull/116

New Contributors

• @SauravMaheshkar made their first contribution in https://github.com/deepmind/kfac-jax/pull/94

Full Changelog: https://github.com/deepmind/kfac-jax/compare/v0.0.3...v0.0.4