kfac-jax (pypi): A Jax package for approximate curvature estimation and optimization using KFAC.
Keywords: bayesian-deep-learning|machine-learning|optimization