- Adding feature to distribute preconditioner application and computation operations across devices.

- Device syncing of factors in TwoKroneckerFactored and Full blocks now happens after every call to update_curvature_matrix_estimate. This fixes a bug which caused devices to possibly desync when using damping adaptation with the approximate quadratic model. However, device communication will now happen more often, which may degrade performance.
- Changing a bunch of methods to be pure functions that don't modify their input state trees. This shouldn't change any behavior and is just a code quality improvement.
- Using jax.pure_callback instead of host_callback.call to better support TPUs.
- Adding newlines between lines of code to improve readability.

PiperOrigin-RevId: 476188081

Author: james-martens

kfac_jax/__init__.py

@@ -24,7 +24,7 @@
 from kfac_jax._src import utils
 
 
-__version__ = "0.0.2"
+__version__ = "0.0.3"
 
 # Patches Second Moments
 patches_moments = patches_second_moment.patches_moments