kfac_jax:

- Adding capability to switch between using Cholesky-based inversions and regular ones. The former are deterministic on GPUs and somewhat faster to compute, but tend to perform a bit worse and not tolerate very low damping values.
- Changing internal implementation of Kronecker factored blocks to prune out singleton dimensions instead of using special case logic for them pi_adjusted_kronecker_factors.
- Going back to using exp-mean-log formula for c_k since this will better avoid over/underflows for blocks with many factors
PiperOrigin-RevId: 578240057
Change-Id: I20ab5268599214a31f4206a9358b76d99434d6c8

ferminet/curvature_tags_and_blocks.py

@@ -24,7 +24,7 @@
 Scalar = kfac_jax.utils.Scalar
 Numeric = kfac_jax.utils.Numeric
 
-vmap_psd_inv_cholesky = jax.vmap(kfac_jax.utils.psd_inv_cholesky, (0, None), 0)
+vmap_psd_inv_cholesky = jax.vmap(kfac_jax.utils.psd_inv, (0, None), 0)
 vmap_matmul = jax.vmap(jnp.matmul, in_axes=(0, 0), out_axes=0)