Repository: jax-ml/jax
Pull request #27287 (2025-03-20 13:39:22)

Prototype of adding error checking to jax.numpy functions
Prototype of adding error checking to jax.numpy functions



Patch:
diff --git a/jax/_src/error_check.py b/jax/_src/error_check.py
index 9d493c1f3..e78b9bc82 100644
--- a/jax/_src/error_check.py
+++ b/jax/_src/error_check.py
@@ -14,18 +14,15 @@
 
 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
@@ -118,56 +115,39 @@ class error_checking_context:
     _error_storage.ref = self.old_ref
 
 
-# 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:
+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 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.
+  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 intended for use in JAX internal APIs (e.g., `jax.numpy`)
-  to perform category-specific runtime checks tied to the operation being
-  performed.
+  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:
     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
@@ -219,37 +199,6 @@ def _set_error_if_with_category(
   _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.
 
@@ -406,53 +355,3 @@ 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()
diff --git a/jax/_src/numpy/error.py b/jax/_src/numpy/error.py
new file mode 100644
index 000000000..52b996a0b
--- /dev/null
+++ b/jax/_src/numpy/error.py
@@ -0,0 +1,131 @@
+# Copyright 2025 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+from typing import Literal
+
+import jax
+from jax._src import config
+
+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":
+    raise ValueError("nan is deprecated. Use `_set_error_if_nan` instead.")
+  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 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.
+
+  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 _is_category_disabled(category):
+    return
+
+  # TODO(mattjj): fix the circular import issue.
+  from jax._src import error_check as error_check_lib
+  error_check_lib.set_error_if(pred, msg)
+
+
+def _set_error_if_nan(pred: jax.Array, /):
+  """Set the internal error state if any element of `pred` is `NaN`.
+
+  This function is disabled if the `jax_error_checking_behavior_nan` flag is
+  set to "ignore".
+  """
+  if config.error_checking_behavior_nan.value == "ignore":
+    return
+
+  # TODO(mattjj): fix the circular import issue.
+  import jax.numpy as jnp
+  if not jnp.issubdtype(pred.dtype, jnp.floating):  # only check floats
+    return
+
+  # TODO(mattjj): fix the circular import issue.
+  from jax._src import error_check as error_check_lib
+  error_check_lib.set_error_if(jnp.isnan(pred), "NaN encountered")
+
+
+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()
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index 91191d24a..1df973039 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -32,12 +32,13 @@ from jax._src.custom_derivatives import custom_jvp
 from jax._src.lax import lax
 from jax._src.lax import other as lax_other
 from jax._src.typing import Array, ArrayLike
+from jax._src.numpy import error as jnp_error
+from jax._src.numpy import reductions
+from jax._src.numpy.ufunc_api import ufunc
 from jax._src.numpy.util import (
    check_arraylike, promote_args, promote_args_inexact,
    promote_args_numeric, promote_dtypes_inexact, promote_dtypes_numeric,
    promote_shapes, _where, check_no_float0s)
-from jax._src.numpy.ufunc_api import ufunc
-from jax._src.numpy import reductions
 from jax._src.util import set_module
 
 
@@ -486,7 +487,9 @@ def log(x: ArrayLike, /) -> Array:
     >>> jnp.allclose(jnp.log(x1*x2), jnp.log(x1)+jnp.log(x2))
     Array(True, dtype=bool)
   """
-  return lax.log(*promote_args_inexact('log', x))
+  out = lax.log(*promote_args_inexact('log', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index 2e03f3317..1baeb4f83 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -1165,6 +1165,11 @@ jax_multiplatform_test(
     srcs = ["error_check_test.py"],
 )
 
+jax_multiplatform_test(
+    name = "jax_numpy_error_test",
+    srcs = ["jax_numpy_error_test.py"],
+)
+
 jax_multiplatform_test(
     name = "stax_test",
     srcs = ["stax_test.py"],
diff --git a/tests/error_check_test.py b/tests/error_check_test.py
index af3f35c7a..0c77989b8 100644
--- a/tests/error_check_test.py
+++ b/tests/error_check_test.py
@@ -372,92 +372,6 @@ 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())
diff --git a/tests/jax_numpy_error_test.py b/tests/jax_numpy_error_test.py
new file mode 100644
index 000000000..c2883f200
--- /dev/null
+++ b/tests/jax_numpy_error_test.py
@@ -0,0 +1,130 @@
+# Copyright 2025 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax._src import config
+from jax._src import error_check
+from jax._src import test_util as jtu
+from jax._src.numpy import error as jnp_error
+import jax.numpy as jnp
+
+config.parse_flags_with_absl()
+
+
+JaxValueError = error_check.JaxValueError
+
+
+class JaxNumpyErrorTests(jtu.JaxTestCase):
+  @parameterized.product(jit=[True, False])
+  def test_set_error_if_nan(self, jit):
+    def f(x):
+      jnp_error._set_error_if_nan(x)
+      return x
+
+    if jit:
+      f = jax.jit(f)
+
+    x = jnp.full((4,), jnp.nan, dtype=jnp.float32)
+
+    with jnp_error.error_checking_behavior(nan="ignore"):
+      _ = f(x)
+      error_check.raise_if_error()  # should not raise error
+
+    with jnp_error.error_checking_behavior(nan="raise"):
+      _ = f(x)
+      with self.assertRaisesRegex(JaxValueError, "NaN"):
+        error_check.raise_if_error()
+
+  @parameterized.product(jit=[True, False])
+  def test_error_category_divide_check(self, jit):
+    def f(x, y):
+      jnp_error._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 jnp_error.error_checking_behavior(divide="ignore"):
+      _ = f(x, y)
+      error_check.raise_if_error()  # should not raise error
+
+    with jnp_error.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):
+      jnp_error._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 jnp_error.error_checking_behavior(oob="ignore"):
+      _ = f(x, start_indices, slice_sizes)
+      error_check.raise_if_error()  # should not raise error
+
+    with jnp_error.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"):
+      jnp_error._set_error_if_with_category(
+          jnp.isnan(jnp.float32(1.0)), "x is NaN", category="invalid"
+      )
+
+  @parameterized.product(jit=[True, False])
+  def test_can_raise_nan_error(self, jit):
+    x = jnp.arange(4, dtype=jnp.float32) - 1
+
+    f = jnp.log
+    if jit:
+      f = jax.jit(f)
+
+    with jnp_error.error_checking_behavior(nan="raise"):
+      f(x)
+      with self.assertRaisesRegex(JaxValueError, "NaN"):
+        error_check.raise_if_error()
+
+
+if __name__ == "__main__":
+  absltest.main(testLoader=jtu.JaxTestLoader())