Removed eager_pmap config option

It defaults to True and is not flipped to False by any internal JAX users.

PiperOrigin-RevId: 745067361

Author: superbobry

CHANGELOG.md

@@ -23,6 +23,7 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * Removed the `config.jax_data_dependent_tracing_fallback` config option,
     which was added temporarily in v0.4.36 to allow users to opt out of the
     new "stackless" tracing machinery.
+  * Removed the `config.jax_eager_pmap` config option.
 
 * Changes
   * The minimum CuDNN version is v9.8.

jax/_src/config.py

@@ -1514,13 +1514,6 @@ def _update_disable_jit_thread_local(val):
           'compute when encountering OOM errors. However, you are '
           'likely to get better results manually with jax.checkpoint'))
 
-# TODO(sharadmv,mattjj): set default to True, then remove
-eager_pmap = bool_state(
-    name='jax_eager_pmap',
-    default=True,
-    upgrade=True,
-    help='Enable eager-mode pmap when jax_disable_jit is activated.')
-
 no_tracing = bool_state(
     name='jax_no_tracing',
     default=False,

jax/_src/interpreters/pxla.py

@@ -338,8 +338,8 @@ def xla_pmap_impl_lazy(
     donated_invars: Sequence[bool],
     is_explicit_global_axis_size: bool,
 ) -> Callable:
-  if (config.disable_jit.value and config.eager_pmap.value and
-      not is_explicit_global_axis_size and not any(d for d in donated_invars)):
+  if (config.disable_jit.value and
+      not is_explicit_global_axis_size and not any(donated_invars)):
     def _emap_apply_fn(*args):
       return _emap_impl(fun, *args, backend=backend, axis_name=axis_name,
                         axis_size=axis_size, global_axis_size=global_axis_size,

tests/pmap_test.py

@@ -3189,7 +3189,7 @@ class EagerPmapMixin:
   def setUp(self):
     super().setUp()
     stack = contextlib.ExitStack()
-    stack.enter_context(jtu.thread_local_config_context(jax_disable_jit=True, jax_eager_pmap=True))
+    stack.enter_context(jtu.thread_local_config_context(jax_disable_jit=True))
     stack.enter_context(jtu.ignore_warning(
         message="Some donated buffers were not usable", category=UserWarning))
     self.addCleanup(stack.close)