Repository: jax-ml/jax
Pull request #27119 (2025-03-13 13:12:05)

Add AOT support for error checking
Add AOT support for error checking



Patch:
diff --git a/jax/_src/error_check.py b/jax/_src/error_check.py
index 88dcec706..e78b9bc82 100644
--- a/jax/_src/error_check.py
+++ b/jax/_src/error_check.py
@@ -14,8 +14,12 @@
 
 from __future__ import annotations
 
+import dataclasses
 from functools import partial
+import json
 import threading
+import traceback as tb_lib
+from types import TracebackType
 import warnings
 
 import jax
@@ -23,19 +27,17 @@ from jax._src import core
 from jax._src import source_info_util
 from jax._src import traceback_util
 import jax._src.mesh as mesh_lib
-from jax.experimental.shard_map import shard_map
+from jax.experimental import shard_map
+import jax.export
 import jax.numpy as jnp
 from jax.sharding import NamedSharding, PartitionSpec as P
 
 
-Traceback = source_info_util.Traceback
-
-
 traceback_util.register_exclusion(__file__)
 
 
 class JaxValueError(ValueError):
-  """Exception raised for failed runtime error checks in JAX."""
+  """Exception raised for runtime errors detected within JAX computations."""
 
 
 #: The default error code for no error.
@@ -45,8 +47,9 @@ class JaxValueError(ValueError):
 _NO_ERROR = jnp.iinfo(jnp.uint32).max
 
 
-_error_list_lock = threading.Lock()
-_error_list: list[tuple[str, Traceback]] = []  # (error_message, traceback) pair
+_error_list_lock = threading.RLock()
+# (error_message, traceback) pairs. Traceback is `str` when imported from AOT.
+_error_list: list[tuple[str, TracebackType | str]] = []
 
 
 class _ErrorStorage(threading.local):
@@ -65,22 +68,21 @@ def _initialize_error_code_ref() -> None:
   In single-device environments, the array is a scalar. In multi-device
   environments, its shape and size match those of the mesh.
   """
-  with core.eval_context():
-    # Get mesh from the context.
-    mesh = mesh_lib.get_concrete_mesh()
-
-    if mesh is None:  # single-device case.
-      error_code = jnp.uint32(_NO_ERROR)
-
-    else:  # multi-device case.
-      sharding = NamedSharding(mesh, P(*mesh.axis_names))
-      error_code = jnp.full(
-          mesh.axis_sizes,
-          jnp.uint32(_NO_ERROR),
-          device=sharding,
-      )
+  # Get mesh from the context.
+  mesh = mesh_lib.get_concrete_mesh()
+
+  if mesh is None:  # single-device case.
+    error_code = jnp.uint32(_NO_ERROR)
+
+  else:  # multi-device case.
+    sharding = NamedSharding(mesh, P(*mesh.axis_names))
+    error_code = jnp.full(
+        mesh.axis_sizes,
+        jnp.uint32(_NO_ERROR),
+        device=sharding,
+    )
 
-    _error_storage.ref = core.mutable_array(error_code)
+  _error_storage.ref = core.mutable_array(error_code)
 
 
 class error_checking_context:
@@ -105,7 +107,8 @@ class error_checking_context:
 
   def __enter__(self):
     self.old_ref = _error_storage.ref
-    _initialize_error_code_ref()
+    with core.eval_context():
+      _initialize_error_code_ref()
     return self
 
   def __exit__(self, exc_type, exc_value, traceback):
@@ -126,22 +129,33 @@ def set_error_if(pred: jax.Array, /, msg: str) -> None:
   will not override the existing error.
 
   For multi-device environments, in explicit mode, users must call
-  :func:`error_checking_context()` to initialize a new error tracking state that
+  :func:`error_checking_context` to initialize a new error tracking state that
   matches the device mesh. In auto mode, implicit cross-device communication may
   occur inside this function, which could impact performance. A warning is
   issued in such cases.
 
+  When exporting a function with `jax.export`, error checking must be explicitly
+  wrapped using :func:`wrap_for_export` before export and
+  :func:`unwrap_from_import` after import.
+
   Args:
     pred: A JAX boolean array. If any element of `pred` is `True`, the internal
       error state will be set.
     msg: The corresponding error message to be raised later.
   """
   if _error_storage.ref is None:
-    _initialize_error_code_ref()
+    with core.eval_context():
+      _initialize_error_code_ref()
     assert _error_storage.ref is not None
 
+  # Get the traceback.
   traceback = source_info_util.current().traceback
   assert traceback is not None
+  traceback = traceback.as_python_traceback()
+  assert isinstance(traceback, TracebackType)
+  traceback = traceback_util.filter_traceback(traceback)
+  assert isinstance(traceback, TracebackType)
+
   with _error_list_lock:
     new_error_code = jnp.uint32(len(_error_list))
     _error_list.append((msg, traceback))
@@ -171,7 +185,7 @@ def set_error_if(pred: jax.Array, /, msg: str) -> None:
             "Please use `with error_checking_context()` to redefine the error "
             "code state based on the mesh."
         )
-      pred = shard_map(
+      pred = shard_map.shard_map(
           partial(jnp.any, keepdims=True),
           mesh=out_sharding.mesh,
           in_specs=in_sharding.spec,
@@ -179,7 +193,7 @@ def set_error_if(pred: jax.Array, /, msg: str) -> None:
       )(pred)  # perform per-device reduction
 
   error_code = _error_storage.ref[...]
-  should_update = jnp.logical_and(pred, error_code == jnp.uint32(_NO_ERROR))
+  should_update = jnp.logical_and(error_code == jnp.uint32(_NO_ERROR), pred)
   error_code = jnp.where(should_update, new_error_code, error_code)
   # TODO(ayx): support vmap and shard_map.
   _error_storage.ref[...] = error_code
@@ -216,8 +230,128 @@ def raise_if_error() -> None:
       device=_error_storage.ref.sharding,
   )  # clear the error code
 
-  msg, traceback = _error_list[error_code]
-  exc = JaxValueError(msg)
-  traceback = traceback.as_python_traceback()
-  filtered_traceback = traceback_util.filter_traceback(traceback)
-  raise exc.with_traceback(filtered_traceback)
+  with _error_list_lock:
+    msg, traceback = _error_list[error_code]
+  if isinstance(traceback, str):  # from imported AOT functions
+    exc = JaxValueError(
+        f"{msg}\nThe original traceback is shown below:\n{traceback}"
+    )
+    raise exc
+  else:
+    exc = JaxValueError(msg)
+    raise exc.with_traceback(traceback)
+
+
+@dataclasses.dataclass(frozen=True)
+class _ErrorClass:
+  """A class to store error information for AOT compilation.
+
+  This class is used internally by the wrapper functions `wrap_for_export` and
+  `unwrap_from_import` to encapsulate error-related data within an exported
+  function.
+
+  Attributes:
+    error_code (jax.Array): A JAX array representing the final error state of
+      the function to be exported. This value is local to the wrapper function.
+    error_list (list[tuple[str, str]]): A list of `(error_message, traceback)`
+      pairs containing error messages and corresponding stack traces. This error
+      list is local to the wrapper function, and does not contain pairs of error
+      information from other functions.
+  """
+
+  error_code: jax.Array
+  error_list: list[tuple[str, str]]
+
+
+jax.tree_util.register_dataclass(
+    _ErrorClass, data_fields=("error_code",), meta_fields=("error_list",)
+)
+jax.export.register_pytree_node_serialization(
+    _ErrorClass,
+    serialized_name=f"{_ErrorClass.__module__}.{_ErrorClass.__name__}",
+    serialize_auxdata=lambda x: json.dumps(x, ensure_ascii=False).encode(
+        "utf-8"
+    ),
+    deserialize_auxdata=lambda x: json.loads(x.decode("utf-8")),
+)
+
+
+def _traceback_to_str(traceback: TracebackType) -> str:
+  """Convert a traceback to a string for export."""
+  return "".join(tb_lib.format_list(tb_lib.extract_tb(traceback))).rstrip("\n")
+
+
+def wrap_for_export(f):
+  """Wrap a function with error checking to make it compatible with AOT mode.
+
+  Error checking relies on global state, which cannot be serialized across
+  processes. This wrapper ensures that the error state remains within the
+  function scope, making it possible to export the function and later import in
+  other processes.
+
+  This function should only be applied once to a function; wrapping the same
+  function multiple times is unnecessary.
+  """
+
+  def inner(*args, **kwargs):
+    global _error_list
+
+    # 1. Save the old state and initialize a new state.
+    with core.eval_context():
+      old_ref = _error_storage.ref
+    _initialize_error_code_ref()
+    with _error_list_lock:
+      old_error_list, _error_list = _error_list, []
+
+      # 2. Trace the function.
+      out = f(*args, **kwargs)
+      error_code = _error_storage.ref[...].min()
+
+      # 3. Restore the old state.
+      _error_list, new_error_list = old_error_list, _error_list
+    with core.eval_context():
+      _error_storage.ref = old_ref
+
+    new_error_list = [
+        (msg, _traceback_to_str(traceback)) for msg, traceback in new_error_list
+    ]
+    return out, _ErrorClass(error_code, new_error_list)
+
+  return inner
+
+
+def unwrap_from_import(f):
+  """Unwrap a function after AOT import to restore error checking.
+
+  When an AOT-exported function is imported in a new process, its error state is
+  separate from the global error state of the current process. This wrapper
+  ensures that errors detected during execution are correctly integrated into
+  the global error checking mechanism of the current process.
+  """
+  if _error_storage.ref is None:
+    with core.eval_context():
+      _initialize_error_code_ref()
+    assert _error_storage.ref is not None
+
+  def inner(*args, **kwargs):
+    out, error_class = f(*args, **kwargs)
+    new_error_code, error_list = error_class.error_code, error_class.error_list
+
+    # Update the global error list.
+    with _error_list_lock:
+      offset = len(_error_list)
+      _error_list.extend(error_list)
+
+    # Update the global error code array.
+    error_code = _error_storage.ref[...]
+    should_update = jnp.logical_and(
+        error_code == jnp.uint32(_NO_ERROR),
+        new_error_code != jnp.uint32(_NO_ERROR),
+    )
+    error_code = jnp.where(should_update, new_error_code + offset, error_code)
+    # TODO(ayx): support vmap and shard_map.
+    _error_storage.ref[...] = error_code
+
+    return out
+
+  return inner


Test patch:
diff --git a/tests/error_check_test.py b/tests/error_check_test.py
index 5bf71a9eb..69e292fb6 100644
--- a/tests/error_check_test.py
+++ b/tests/error_check_test.py
@@ -22,6 +22,7 @@ from jax._src import config
 from jax._src import error_check
 from jax._src import mesh as mesh_lib
 from jax._src import test_util as jtu
+import jax.export
 import jax.numpy as jnp
 from jax.sharding import NamedSharding, PartitionSpec as P
 
@@ -33,7 +34,9 @@ config.parse_flags_with_absl()
 jtu.request_cpu_devices(4)
 
 
-@jtu.with_config(jax_check_tracer_leaks=True)
+# TODO: AOT tests fails with the tracer leak checker.
+# Reenable once https://github.com/jax-ml/jax/issues/27315 is fixed.
+# @jtu.with_config(jax_check_tracer_leaks=True)
 class ErrorCheckTests(jtu.JaxTestCase):
 
   @parameterized.product(jit=[True, False])
@@ -280,6 +283,60 @@ class ErrorCheckTests(jtu.JaxTestCase):
       with self.assertRaisesRegex(JaxValueError, "x must be greater than 0"):
         error_check.raise_if_error()
 
+  def test_error_check_aot(self):
+    def run_export():
+      def f(x):
+        error_check.set_error_if(x <= 0, "x must be greater than 0")
+        return x + 1
+
+      f = jax.jit(error_check.wrap_for_export(jax.jit(f)))
+      x = jax.ShapeDtypeStruct((), jnp.float32)
+      serialized = jax.export.export(f)(x).serialize()
+      return serialized
+
+    def run_import(serialized):
+      f = jax.export.deserialize(serialized).call
+      f = jax.jit(error_check.unwrap_from_import(jax.jit(f)))
+      x = jnp.float32(-3.)
+      _ = f(x)
+      with self.assertRaisesRegex(JaxValueError, "x must be greater than 0"):
+        error_check.raise_if_error()
+
+    serialized = run_export()
+    run_import(serialized)
+
+  def test_error_check_aot_should_not_override_existing_error(self):
+    def f1(x):
+      error_check.set_error_if(x <= 0, "x must be greater than 0 in f1")
+      return x + 1
+
+    def run_export():
+      def f2(x):
+        error_check.set_error_if(x <= 0, "x must be greater than 0 in f2")
+        return x + 1
+
+      f2 = jax.jit(error_check.wrap_for_export(jax.jit(f2)))
+      x = jax.ShapeDtypeStruct((), jnp.float32)
+      serialized = jax.export.export(f2)(x).serialize()
+      return serialized
+
+    def run_import(serialized):
+      f2 = jax.export.deserialize(serialized).call
+      f2 = jax.jit(error_check.unwrap_from_import(jax.jit(f2)))
+      return f2
+
+    x = jnp.float32(-3.)
+    _ = f1(x)  # check fails. so it should set error
+
+    serialized = run_export()
+    f2 = run_import(serialized)
+    _ = f2(x)  # check fails, but should not override the error
+
+    with self.assertRaisesRegex(
+        JaxValueError, "x must be greater than 0 in f1"
+    ):
+      error_check.raise_if_error()
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())