Repository: jax-ml/jax
Pull request #27286 (2025-03-20 13:19:11)

Add support for configurable error checking categories
Add support for configurable error checking categories



Patch:
diff --git a/jax/_src/config.py b/jax/_src/config.py
index cf6a07834..5b8b87be2 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -245,7 +245,10 @@ def trace_context():
           pgle_profiling_runs.value,
           enable_pgle.value,
           use_shardy_partitioner.value,
-          use_high_dynamic_range_gumbel.value)
+          use_high_dynamic_range_gumbel.value,
+          error_checking_behavior_nan.value,
+          error_checking_behavior_divide.value,
+          error_checking_behavior_oob.value)
 
 config = Config()
 
@@ -1317,6 +1320,41 @@ disallow_mesh_context_manager = bool_state(
     ),
 )
 
+# TODO(ayx): Move these 3 flags out of config once we have a user-level
+# extension mechanism for adding contexts to which the jit cache is sensitive.
+error_checking_behavior_nan = enum_state(
+    name='jax_error_checking_behavior_nan',
+    enum_values=['ignore', 'raise'],
+    default='ignore',
+    help=(
+        'Specify the behavior when a NaN is encountered. Options are "ignore"'
+        ' or "raise".'
+    ),
+    include_in_jit_key=True,
+)
+
+error_checking_behavior_divide = enum_state(
+    name='jax_error_checking_behavior_divide',
+    enum_values=['ignore', 'raise'],
+    default='ignore',
+    help=(
+        'Specify the behavior when a divide by zero is encountered. Options are'
+        ' "ignore" or "raise".'
+    ),
+    include_in_jit_key=True,
+)
+
+error_checking_behavior_oob = enum_state(
+    name='jax_error_checking_behavior_oob',
+    enum_values=['ignore', 'raise'],
+    default='ignore',
+    help=(
+        'Specify the behavior when an out of bounds access is encountered.'
+        ' Options are "ignore" or "raise".'
+    ),
+    include_in_jit_key=True,
+)
+
 def _update_x64_global(val):
   jax_jit.global_state().enable_x64 = val
 
diff --git a/jax/_src/error_check.py b/jax/_src/error_check.py
index e78b9bc82..9d493c1f3 100644
--- a/jax/_src/error_check.py
+++ b/jax/_src/error_check.py
@@ -14,15 +14,18 @@
 
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 from functools import partial
 import json
 import threading
 import traceback as tb_lib
 from types import TracebackType
+from typing import Literal
 import warnings
 
 import jax
+from jax._src import config
 from jax._src import core
 from jax._src import source_info_util
 from jax._src import traceback_util
@@ -115,39 +118,56 @@ class error_checking_context:
     _error_storage.ref = self.old_ref
 
 
-def set_error_if(pred: jax.Array, /, msg: str) -> None:
+# TODO(ayx): Move all category-related logic into the jax.numpy integration once
+# it is ready. This logic is specific to how jax.numpy decides when to call
+# `set_error_if`, and doesn't belong in the core error-checking library itself.
+# The responsibility for deciding whether to predicate an error should lie with
+# the user or the higher-level library (like jax.numpy), not with
+# `set_error_if`.
+Category = Literal["nan", "divide", "oob"]
+
+
+def _is_category_disabled(
+    category: Category | None,
+) -> bool:
+  """Check if the error checking behavior for the given category is disabled."""
+  if category is None:
+    return False
+  if category == "nan":
+    return config.error_checking_behavior_nan.value == "ignore"
+  if category == "divide":
+    return config.error_checking_behavior_divide.value == "ignore"
+  if category == "oob":
+    return config.error_checking_behavior_oob.value == "ignore"
+  raise ValueError(f"Invalid category: {category}")
+
+
+def _set_error_if_with_category(
+    pred: jax.Array,
+    /,
+    msg: str,
+    category: Category | None = None,
+) -> None:
   """Set the internal error state if any element of `pred` is `True`.
 
-  This function is used inside JAX computations to detect runtime errors without
-  immediately halting execution. When this function is traced (e.g., inside
-  :func:`jax.jit`), the corresponding error message and its traceback are
-  recorded. At execution time, if `pred` contains any `True` values, the error
-  state is set, but execution continues without interruption. The recorded error
-  can later be raised using :func:`raise_if_error`.
-
-  If the error state has already been set, subsequent errors are ignored and
-  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
-  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.
+  This function is similar to :func:`set_error_if`, but it also takes a category
+  argument. The category can be "nan", "divide", or "oob". The error checking
+  behavior for each category can be configured using
+  :func:`set_error_checking_behavior`. If not provided, there will be no
+  category.
 
-  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.
+  This function is intended for use in JAX internal APIs (e.g., `jax.numpy`)
+  to perform category-specific runtime checks tied to the operation being
+  performed.
   """
   if _error_storage.ref is None:
     with core.eval_context():
       _initialize_error_code_ref()
     assert _error_storage.ref is not None
 
+  if _is_category_disabled(category):
+    return
+
   # Get the traceback.
   traceback = source_info_util.current().traceback
   assert traceback is not None
@@ -199,6 +219,37 @@ def set_error_if(pred: jax.Array, /, msg: str) -> None:
   _error_storage.ref[...] = error_code
 
 
+def set_error_if(pred: jax.Array, /, msg: str) -> None:
+  """Set the internal error state if any element of `pred` is `True`.
+
+  This function is used inside JAX computations to detect runtime errors without
+  immediately halting execution. When this function is traced (e.g., inside
+  :func:`jax.jit`), the corresponding error message and its traceback are
+  recorded. At execution time, if `pred` contains any `True` values, the error
+  state is set, but execution continues without interruption. The recorded error
+  can later be raised using :func:`raise_if_error`.
+
+  If the error state has already been set, subsequent errors are ignored and
+  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
+  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.
+  """
+  _set_error_if_with_category(pred, msg)
+
+
 def raise_if_error() -> None:
   """Raise an exception if the internal error state is set.
 
@@ -355,3 +406,53 @@ def unwrap_from_import(f):
     return out
 
   return inner
+
+
+Behavior = Literal["ignore", "raise"]
+
+
+class error_checking_behavior:
+  """A context manager to set the error checking behavior.
+
+  If both `all` and a category are provided, the category will override the
+  `all` setting.
+
+  When the error checking behavior is set to "ignore", all errors will be
+  ignored. When set to "raise", errors will be detected and recorded, but an
+  exception will not be raised immediately. Users must call
+  :func:`raise_if_error` to at the end of the computation to raise the
+  exception.
+  """
+
+  def __init__(
+      self,
+      *,
+      all: Behavior | None = None,
+      nan: Behavior | None = None,
+      divide: Behavior | None = None,
+      oob: Behavior | None = None,
+  ) -> None:
+    new_settings = {}
+    if all is not None:
+      new_settings["nan"] = new_settings["divide"] = new_settings["oob"] = all
+    if nan is not None:
+      new_settings["nan"] = nan
+    if divide is not None:
+      new_settings["divide"] = divide
+    if oob is not None:
+      new_settings["oob"] = oob
+    self.new_settings = new_settings
+    self.stack = contextlib.ExitStack()
+
+  def __enter__(self):
+    config_flags = {
+        "nan": config.error_checking_behavior_nan,
+        "divide": config.error_checking_behavior_divide,
+        "oob": config.error_checking_behavior_oob,
+    }
+    for key, value in self.new_settings.items():
+        self.stack.enter_context(config_flags[key](value))
+    return self
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    self.stack.close()


Test patch:
diff --git a/tests/error_check_test.py b/tests/error_check_test.py
index 69e292fb6..7f75eeb62 100644
--- a/tests/error_check_test.py
+++ b/tests/error_check_test.py
@@ -337,6 +337,92 @@ class ErrorCheckTests(jtu.JaxTestCase):
     ):
       error_check.raise_if_error()
 
+  @parameterized.product(jit=[True, False])
+  def test_error_category_nan_check(self, jit):
+    def f(x):
+      error_check._set_error_if_with_category(
+          jnp.isnan(x), "x is NaN", category="nan"
+      )
+      return x
+
+    if jit:
+      f = jax.jit(f)
+
+    x = jnp.full((4,), jnp.nan, dtype=jnp.float32)
+
+    with error_check.error_checking_behavior(nan="ignore"):
+      _ = f(x)
+      error_check.raise_if_error()  # should not raise error
+
+    with error_check.error_checking_behavior(nan="raise"):
+      _ = f(x)
+      with self.assertRaisesRegex(JaxValueError, "x is NaN"):
+        error_check.raise_if_error()
+
+  @parameterized.product(jit=[True, False])
+  def test_error_category_divide_check(self, jit):
+    def f(x, y):
+      error_check._set_error_if_with_category(
+          y == 0.0, "division by zero", category="divide"
+      )
+      return x / y
+
+    if jit:
+      f = jax.jit(f)
+
+    x = jnp.arange(4, dtype=jnp.float32) + 1
+    y = jnp.arange(4, dtype=jnp.float32)
+
+    with error_check.error_checking_behavior(divide="ignore"):
+      _ = f(x, y)
+      error_check.raise_if_error()  # should not raise error
+
+    with error_check.error_checking_behavior(divide="raise"):
+      _ = f(x, y)
+      with self.assertRaisesRegex(JaxValueError, "division by zero"):
+        error_check.raise_if_error()
+
+  @parameterized.product(jit=[True, False])
+  def test_error_category_oob_check(self, jit):
+    def f(x, start_indices, slice_sizes):
+      error_check._set_error_if_with_category(
+          jnp.logical_or(
+              start_indices < 0,
+              start_indices + jnp.array(slice_sizes, dtype=jnp.int32)
+              >= jnp.array(x.shape, dtype=jnp.int32),
+          ),
+          "Out of bounds in dynamic_slice",
+          category="oob",
+      )
+      y = jax.lax.dynamic_slice(
+          x, start_indices, slice_sizes, allow_negative_indices=False
+      )
+      return y
+
+    if jit:
+      f = jax.jit(f, static_argnums=(2,))
+
+    x = jnp.arange(12).reshape(3, 4)
+    start_indices = jnp.array([0, -1], dtype=jnp.int32)
+    slice_sizes = (3, 4)
+
+    with error_check.error_checking_behavior(oob="ignore"):
+      _ = f(x, start_indices, slice_sizes)
+      error_check.raise_if_error()  # should not raise error
+
+    with error_check.error_checking_behavior(oob="raise"):
+      _ = f(x, start_indices, slice_sizes)
+      with self.assertRaisesRegex(
+          JaxValueError, "Out of bounds in dynamic_slice",
+      ):
+        error_check.raise_if_error()
+
+  def test_error_category_invalid_category(self):
+    with self.assertRaisesRegex(ValueError, "Invalid category"):
+      error_check._set_error_if_with_category(
+          jnp.isnan(jnp.float32(1.0)), "x is NaN", category="invalid"
+      )
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())