Use https://github.com/deepmind/kfac-jax to implement kfac-laplace
Use https://github.com/deepmind/kfac-jax to implement kfac-laplace