Async dispatch expensive computations on the JAX CPU backend.

By setting `jax.config.update('jax_cpu_enable_async_dispatch', False)`, one could opt out of the change and recover the old behavior.

PiperOrigin-RevId: 633264117

CHANGELOG.md

@@ -8,6 +8,12 @@ Remember to align the itemized text with the first line of an item within a list
 
 ## jax 0.4.29
 
+* Changes
+  * Computations on the CPU backend may now be dispatched asynchronously in
+    more cases. Previously non-parallel computations were always dispatched
+    synchronously. You can recover the old behavior by setting
+    `jax.config.update('jax_cpu_enable_async_dispatch', False)`.
+
 ## jaxlib 0.4.29
 
 ## jax 0.4.28 (May 9, 2024)

jax/_src/xla_bridge.py

@@ -113,11 +113,9 @@
          '"gloo" or "mpi"'
 )
 
-# TODO(yueshengys): turn default back to True after resolving memory increase
-# issue.
 _CPU_ENABLE_ASYNC_DISPATCH = config.DEFINE_bool(
     name="jax_cpu_enable_async_dispatch",
-    default=False,
+    default=True,
     help="Only applies to non-parallel computations. If False, run computations"
     "inline without async dispatch.",
 )