Adding JAX config change jax_pmap_shmap_merge=False to the top of the optimizer.py file. The new pmap implementation currently seems to slow down projects using kfac_jax.

PiperOrigin-RevId: 840716303

Author: james-martens

kfac_jax/_src/optimizer.py

@@ -25,6 +25,12 @@
 from kfac_jax._src import utils
 from typing_extensions import Self
 
+# For now we are opting out of JAX's new pmap simulation in favor using the old
+# pmap implementation. This is because the new simulation currently leads to
+# a compute performance regression in some experiments.
+if "jax_pmap_shmap_merge" in jax.config.values:
+  jax.config.update("jax_pmap_shmap_merge", False)
+
 
 # Types for annotation
 Array = utils.Array