Port many uses of contextlib.contextdecorator to explicit context manager classes.

contextdecorator turns out to be slower than just writing a decorator class explicitly. Since we use many decorators per-equation, this causes a measurable speed difference in certain benchmarks.

PiperOrigin-RevId: 730939406

Author: hawkinsp

jax/_src/api.py

@@ -24,7 +24,7 @@
 
 import atexit
 import collections
-from collections.abc import Callable, Generator, Hashable, Iterable, Sequence
+from collections.abc import Callable, Hashable, Iterable, Sequence
 from functools import partial, lru_cache
 import inspect
 import math
@@ -2736,10 +2736,9 @@ def named_call(
   return source_info_util.extend_name_stack(name)(fun)
 
 
-@contextmanager
 def named_scope(
     name: str,
-  ) -> Generator[None, None, None]:
+  ) -> source_info_util.ExtendNameStackContextManager:
   """A context manager that adds a user specified name to the JAX name stack.
 
   When staging out computations for just-in-time compilation to XLA (or other
@@ -2786,8 +2785,7 @@ def named_scope(
   """
   if not isinstance(name, str):
     raise TypeError("named_scope name argument must be a string.")
-  with source_info_util.extend_name_stack(name):
-    yield
+  return source_info_util.extend_name_stack(name)
 
 def effects_barrier():
   """Waits until existing functions have completed any side-effects."""

jax/_src/config.py

@@ -304,31 +304,8 @@ def _set(self, value: _T) -> None:
     if self._update_global_hook:
       self._update_global_hook(value)
 
-  @contextlib.contextmanager
   def __call__(self, new_val: Any = no_default):
-    if new_val is no_default:
-      if self._default_context_manager_value is not no_default:
-        new_val = self._default_context_manager_value  # default_context_manager_value provided to constructor
-      else:
-        # no default_value provided to constructor and no value provided as an
-        # argument, so we raise an error
-        raise TypeError(f"Context manager for {self.__name__} config option "
-                        "requires an argument representing the new value for "
-                        "the config option.")
-    if self._validator:
-      self._validator(new_val)
-    prev_val = self.swap_local(new_val)
-    if self._update_thread_local_hook:
-      self._update_thread_local_hook(new_val)
-    try:
-      yield
-    finally:
-      self.set_local(prev_val)
-      if self._update_thread_local_hook:
-        if prev_val is config_ext.unset:
-          self._update_thread_local_hook(None)
-        else:
-          self._update_thread_local_hook(cast(Optional[Any], prev_val))
+    return StateContextManager(self, new_val)
 
   def _add_hooks(self, update_global_hook, update_thread_local_hook):
     """Private method that adds hooks to an existing context-manager.
@@ -339,6 +316,40 @@ def _add_hooks(self, update_global_hook, update_thread_local_hook):
     update_global_hook(self.get_global())
 
 
+class StateContextManager(contextlib.ContextDecorator):
+  __slots__ = ['state', 'new_val', 'prev']
+
+  def __init__(self, state, new_val):
+    self.state = state
+    self.new_val = new_val
+
+    if new_val is no_default:
+      if state._default_context_manager_value is not no_default:
+        new_val = state._default_context_manager_value  # default_context_manager_value provided to constructor
+      else:
+        # no default_value provided to constructor and no value provided as an
+        # argument, so we raise an error
+        raise TypeError(f"Context manager for {state.__name__} config option "
+                        "requires an argument representing the new value for "
+                        "the config option.")
+    if state._validator:
+      state._validator(new_val)
+
+
+  def __enter__(self):
+    self.prev = self.state.swap_local(self.new_val)
+    if self.state._update_thread_local_hook:
+      self.state._update_thread_local_hook(self.new_val)
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    self.state.set_local(self.prev)
+    if self.state._update_thread_local_hook:
+      if self.prev is config_ext.unset:
+        self.state._update_thread_local_hook(None)
+      else:
+        self.state._update_thread_local_hook(cast(Optional[Any], self.prev))
+
+
 UPGRADE_BOOL_HELP = (
     " This will be enabled by default in future versions of JAX, at which "
     "point all uses of the flag will be considered deprecated (following "

jax/_src/core.py

@@ -1136,51 +1136,77 @@ def update_thread_local_jit_state(self):
 trace_ctx = TracingContext()
 
 
-@contextmanager
-def take_current_trace():
-  prev = trace_ctx.trace
-  try:
+class TakeCurrentTraceContextManager:
+  __slots__ = ['prev']
+
+  def __enter__(self):
+    self.prev = trace_ctx.trace
     trace_ctx.set_trace(eval_trace)
-    yield prev
-  finally:
-    trace_ctx.set_trace(prev)
+    return self.prev
 
-@contextmanager
-def set_current_trace(trace, check_leaks=False):
-  prev = trace_ctx.trace
-  try:
-    trace_ctx.set_trace(trace)
-    yield
-  finally:
-    trace_ctx.set_trace(prev)
-    if check_leaks and config.check_tracer_leaks.value:
-      trace.invalidate()
-      trace_ref = ref(trace)
-      del trace
+  def __exit__(self, exc_type, exc_value, traceback):
+    trace_ctx.set_trace(self.prev)
+
+take_current_trace = TakeCurrentTraceContextManager
+
+
+class SetCurrentTraceContextManager:
+  __slots__ = ['trace', 'check_leaks', 'prev']
+
+  def __init__(self, trace, check_leaks=False):
+    self.trace = trace
+    self.check_leaks = check_leaks
+
+  def __enter__(self):
+    self.prev = trace_ctx.trace
+    trace_ctx.set_trace(self.trace)
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    trace_ctx.set_trace(self.prev)
+    if self.check_leaks and config.check_tracer_leaks.value:
+      self.trace.invalidate()
+      trace_ref = ref(self.trace)
+      del self.trace
       live_trace = trace_ref()
       if live_trace is not None:
         leaked_tracers = maybe_find_leaked_tracers(live_trace)
         if leaked_tracers:
           raise leaked_tracer_error("trace", live_trace, leaked_tracers)
 
-@contextmanager
-def extend_axis_env_nd(name_size_pairs : Iterable[tuple[AxisName, int]]):
-  prev = trace_ctx.axis_env
-  try:
-    trace_ctx.set_axis_env(prev.extend_pure(name_size_pairs))
-    yield
-  finally:
-    trace_ctx.set_axis_env(prev)
+set_current_trace = SetCurrentTraceContextManager
+
+class ExtendAxisEnvNdContextManager:
+  __slots__ = ['prev', 'name_size_pairs']
+
+  def __init__(self, name_size_pairs: Iterable[tuple[AxisName, int]]):
+    self.name_size_pairs = name_size_pairs
+
+  def __enter__(self):
+    self.prev = trace_ctx.axis_env
+    trace_ctx.set_axis_env(self.prev.extend_pure(self.name_size_pairs))
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    trace_ctx.set_axis_env(self.prev)
+
+extend_axis_env_nd = ExtendAxisEnvNdContextManager
+
+
+class AddSpmdAxisNamesContextManager:
+  __slots__ = ['prev', 'axis_names']
+
+  def __init__(self, axis_names: AxisName | None):
+    self.axis_names = axis_names
+
+  def __enter__(self):
+    self.prev = trace_ctx.axis_env
+    if self.axis_names is not None:
+      trace_ctx.set_axis_env(self.prev.add_spmd_axis_names(self.axis_names))
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    trace_ctx.set_axis_env(self.prev)
+
+add_spmd_axis_names = AddSpmdAxisNamesContextManager
 
-@contextmanager
-def add_spmd_axis_names(axis_names: AxisName | None):
-  prev = trace_ctx.axis_env
-  try:
-    if axis_names is not None:
-      trace_ctx.set_axis_env(prev.add_spmd_axis_names(axis_names))
-    yield
-  finally:
-    trace_ctx.set_axis_env(prev)
 
 def get_axis_env():
   return trace_ctx.axis_env
@@ -1233,6 +1259,7 @@ def ensure_no_leaks(trace:Trace):
       if leaked_tracers:
         raise leaked_tracer_error("trace", live_trace, leaked_tracers)
 
+
 def maybe_find_leaked_tracers(trace: Trace) -> list[Tracer]:
   """Find the leaked tracers holding a reference to the Trace
   """

jax/_src/dispatch.py

@@ -17,7 +17,6 @@
 
 import atexit
 from collections.abc import Sequence
-import contextlib
 import dataclasses
 import enum
 from functools import partial
@@ -170,22 +169,32 @@ def wait_for_tokens():
   runtime_tokens.block_until_ready()
 
 
-@contextlib.contextmanager
-def log_elapsed_time(fmt: str, fun_name: str, event: str | None = None):
-  if _on_exit:
-    yield
-  else:
-    log_priority = logging.WARNING if config.log_compiles.value else logging.DEBUG
-    start_time = time.time()
-    yield
+class LogElapsedTimeContextManager:
+  __slots__ = ['fmt', 'fun_name', 'event', 'start_time']
+
+  def __init__(self, fmt: str, fun_name: str, event: str | None = None):
+    self.fmt = fmt
+    self.fun_name = fun_name
+    self.event = event
+
+  def __enter__(self):
+    self.start_time = time.time()
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    if _on_exit:
+      return
+
     end_time = time.time()
-    elapsed_time = end_time - start_time
+    elapsed_time = end_time - self.start_time
+    log_priority = logging.WARNING if config.log_compiles.value else logging.DEBUG
     if logger.isEnabledFor(log_priority):
-      logger.log(log_priority, fmt.format(
-          fun_name=fun_name, elapsed_time=elapsed_time))
-    if event is not None:
-      record_event_duration_secs(event, elapsed_time)
-      record_event_time_span(event, start_time, end_time)
+      logger.log(log_priority, self.fmt.format(
+          fun_name=self.fun_name, elapsed_time=elapsed_time))
+    if self.event is not None:
+      record_event_duration_secs(self.event, elapsed_time)
+      record_event_time_span(self.event, self.start_time, end_time)
+
+log_elapsed_time = LogElapsedTimeContextManager
 
 
 def should_tuple_args(num_args: int, platform: str) -> bool:

jax/_src/mesh.py

@@ -549,14 +549,20 @@ def _extremely_unsafe_enter_tracing_context(mesh: AbstractMesh):
 def _raise_value_error(name):
   raise ValueError(f"AbstractMesh does not implement {name}")
 
+class SetAbstractMeshContextManager:
+  __slots__ = ['mesh', 'prev']
+
+  def __init__(self, mesh: AbstractMesh):
+    self.mesh = mesh
+
+  def __enter__(self):
+    self.prev = jax_config.abstract_mesh_context_manager.swap_local(self.mesh)
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    jax_config.abstract_mesh_context_manager.set_local(self.prev)
+
+set_abstract_mesh = SetAbstractMeshContextManager
 
-@contextlib.contextmanager
-def set_abstract_mesh(mesh: AbstractMesh):
-  prev_val = jax_config.abstract_mesh_context_manager.swap_local(mesh)
-  try:
-    yield
-  finally:
-    jax_config.abstract_mesh_context_manager.set_local(prev_val)
 
 empty_abstract_mesh = AbstractMesh(())