Repository: jax-ml/jax
Pull request #27366 (2025-03-24 14:31:59)

Add nan checking to jax.numpy functions
Add nan checking to jax.numpy functions



Patch:
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
index 1df973039..0ea2992c9 100644
--- a/jax/_src/numpy/ufuncs.py
+++ b/jax/_src/numpy/ufuncs.py
@@ -575,7 +575,9 @@ def log1p(x: ArrayLike, /) -> Array:
     >>> jnp.expm1(jnp.log(x1+1))  # doctest: +SKIP
     Array([1.000166e-04, 9.536743e-07, 0.000000e+00], dtype=float32)
   """
-  return lax.log1p(*promote_args_inexact('log1p', x))
+  out = lax.log1p(*promote_args_inexact('log1p', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -607,7 +609,9 @@ def sin(x: ArrayLike, /) -> Array:
     ...   print(jnp.sin(x))
     [ 0.707  1.     0.707 -0.   ]
   """
-  return lax.sin(*promote_args_inexact('sin', x))
+  out = lax.sin(*promote_args_inexact('sin', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -638,7 +642,9 @@ def cos(x: ArrayLike, /) -> Array:
     ...   print(jnp.cos(x))
     [ 0.707 -0.    -0.707 -0.866]
   """
-  return lax.cos(*promote_args_inexact('cos', x))
+  out = lax.cos(*promote_args_inexact('cos', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -669,7 +675,9 @@ def tan(x: ArrayLike, /) -> Array:
     ...   print(jnp.tan(x))
     [ 0.     0.577  1.    -1.    -0.577]
   """
-  return lax.tan(*promote_args_inexact('tan', x))
+  out = lax.tan(*promote_args_inexact('tan', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -711,7 +719,9 @@ def arcsin(x: ArrayLike, /) -> Array:
     ...   jnp.arcsin(3+4j)
     Array(0.634+2.306j, dtype=complex64, weak_type=True)
   """
-  return lax.asin(*promote_args_inexact('arcsin', x))
+  out = lax.asin(*promote_args_inexact('arcsin', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -754,7 +764,9 @@ def arccos(x: ArrayLike, /) -> Array:
     ...   jnp.arccos(4-1j)
     Array(0.252+2.097j, dtype=complex64, weak_type=True)
   """
-  return lax.acos(*promote_args_inexact('arccos', x))
+  out = lax.acos(*promote_args_inexact('arccos', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -1008,6 +1020,7 @@ def arccosh(x: ArrayLike, /) -> Array:
   # Note: arccosh is multi-valued for complex input, and lax.acosh
   # uses a different convention than np.arccosh.
   result = lax.acosh(*promote_args_inexact("arccosh", x))
+  jnp_error._set_error_if_nan(result)
   if dtypes.issubdtype(result.dtype, np.complexfloating):
     result = _where(real(result) < 0, lax.neg(result), result)
   return result
@@ -1113,7 +1126,9 @@ def arctanh(x: ArrayLike, /) -> Array:
     ...   jnp.arctanh(x1)
     Array([-0.549+1.571j,  0.347+1.571j,  0.239-1.509j], dtype=complex64)
   """
-  return lax.atanh(*promote_args_inexact('arctanh', x))
+  out = lax.atanh(*promote_args_inexact('arctanh', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -1146,7 +1161,9 @@ def sqrt(x: ArrayLike, /) -> Array:
     >>> jnp.sqrt(-1)
     Array(nan, dtype=float32, weak_type=True)
   """
-  return lax.sqrt(*promote_args_inexact('sqrt', x))
+  out = lax.sqrt(*promote_args_inexact('sqrt', x))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -1215,7 +1232,11 @@ def add(x: ArrayLike, y: ArrayLike, /) -> Array:
     Array([10, 11, 12, 13], dtype=int32)
   """
   x, y = promote_args("add", x, y)
-  return lax.add(x, y) if x.dtype != bool else lax.bitwise_or(x, y)
+  if x.dtype == bool:
+    return lax.bitwise_or(x, y)
+  out = lax.add(x, y)
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 def _multiply_at(a: Array, indices: Any, b: ArrayLike) -> Array:
@@ -1544,7 +1565,9 @@ def subtract(x: ArrayLike, y: ArrayLike, /) -> Array:
     >>> x - 10
     Array([-10,  -9,  -8,  -7], dtype=int32)
   """
-  return lax.sub(*promote_args("subtract", x, y))
+  out = lax.sub(*promote_args("subtract", x, y))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -1768,7 +1791,9 @@ def float_power(x: ArrayLike, y: ArrayLike, /) -> Array:
     >>> jnp.float_power(-3, 1.7)
     Array(nan, dtype=float32, weak_type=True)
   """
-  return lax.pow(*promote_args_inexact("float_power", x, y))
+  out = lax.pow(*promote_args_inexact("float_power", x, y))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -2446,7 +2471,9 @@ def true_divide(x1: ArrayLike, x2: ArrayLike, /) -> Array:
     :func:`jax.numpy.floor_divide` for integer division
   """
   x1, x2 = promote_args_inexact("true_divide", x1, x2)
-  return lax.div(x1, x2)
+  out = lax.div(x1, x2)
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -2648,7 +2675,9 @@ def power(x1: ArrayLike, x2: ArrayLike, /) -> Array:
       return lax.integer_pow(x1, x2)
 
   # Handle cases #2 and #3 under a jit:
-  return _power(x1, x2)
+  out = _power(x1, x2)
+  jnp_error._set_error_if_nan(out)
+  return out
 
 @export
 def pow(x1: ArrayLike, x2: ArrayLike, /) -> Array:
@@ -2774,7 +2803,9 @@ def log2(x: ArrayLike, /) -> Array:
     im = lax.imag(r)
     ln2 = lax.log(_constant_like(re, 2))
     return lax.complex(lax.div(re, ln2), lax.div(im, ln2))
-  return lax.div(lax.log(x), lax.log(_constant_like(x, 2)))
+  out = lax.div(lax.log(x), lax.log(_constant_like(x, 2)))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -2804,7 +2835,9 @@ def log10(x: ArrayLike, /) -> Array:
     im = lax.imag(r)
     ln10 = lax.log(_constant_like(re, 10))
     return lax.complex(lax.div(re, ln10), lax.div(im, ln10))
-  return lax.div(lax.log(x), lax.log(_constant_like(x, 10)))
+  out = lax.div(lax.log(x), lax.log(_constant_like(x, 10)))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -3064,7 +3097,9 @@ def remainder(x1: ArrayLike, x2: ArrayLike, /) -> Array:
   trunc_mod_not_zero = lax.ne(trunc_mod, zero)
   do_plus = lax.bitwise_and(
       lax.ne(lax.lt(trunc_mod, zero), lax.lt(x2, zero)), trunc_mod_not_zero)
-  return lax.select(do_plus, lax.add(trunc_mod, x2), trunc_mod)
+  out = lax.select(do_plus, lax.add(trunc_mod, x2), trunc_mod)
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export
@@ -3112,7 +3147,9 @@ def fmod(x1: ArrayLike, x2: ArrayLike, /) -> Array:
   check_arraylike("fmod", x1, x2)
   if dtypes.issubdtype(dtypes.result_type(x1, x2), np.integer):
     x2 = _where(x2 == 0, lax._ones(x2), x2)
-  return lax.rem(*promote_args_numeric("fmod", x1, x2))
+  out = lax.rem(*promote_args_numeric("fmod", x1, x2))
+  jnp_error._set_error_if_nan(out)
+  return out
 
 
 @export


Test patch:
diff --git a/tests/jax_numpy_error_test.py b/tests/jax_numpy_error_test.py
index c2883f200..08917aeed 100644
--- a/tests/jax_numpy_error_test.py
+++ b/tests/jax_numpy_error_test.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import operator
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -112,16 +114,68 @@ class JaxNumpyErrorTests(jtu.JaxTestCase):
           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
+  @staticmethod
+  def op_cases(cases):
+    for jit in (True, False):
+      for func, operands in cases:
+        if not isinstance(operands, tuple):
+          operands = (operands,)
+
+        jit_str = "jit" if jit else "nojit"
+        func_str = f"{func.__module__}.{func.__name__}"
+        name = f"_{jit_str}_{func_str}"
+
+        yield name, jit, func, operands
+
+  @parameterized.named_parameters(
+      op_cases((
+          # list of all NaN-producing jax.numpy functions
+          # go/keep-sorted start
+          (jnp.acos, 2.0),
+          (jnp.acosh, 0.5),
+          (jnp.add, (jnp.inf, -jnp.inf)),
+          (jnp.arccos, 2.0),
+          (jnp.arccosh, 0.5),
+          (jnp.arcsin, -2.0),
+          (jnp.arctanh, -2.0),
+          (jnp.asin, -2.0),
+          (jnp.atanh, -2.0),
+          (jnp.cos, jnp.inf),
+          (jnp.divide, (0.0, 0.0)),
+          (jnp.divmod, (1.0, 0.0)),
+          (jnp.float_power, (-1.0, 0.5)),
+          (jnp.fmod, (1.0, 0.0)),
+          (jnp.log, -1.0),
+          (jnp.log10, -1.0),
+          (jnp.log1p, -1.5),
+          (jnp.log2, -1.0),
+          (jnp.mod, (1.0, 0.0)),
+          (jnp.pow, (-1.0, 0.5)),
+          (jnp.power, (-1.0, 0.5)),
+          (jnp.remainder, (1.0, 0.0)),
+          (jnp.sin, jnp.inf),
+          # TODO(https://github.com/jax-ml/jax/issues/27470): Not yet supported.
+          # (jnp.sinc, jnp.inf),
+          (jnp.sqrt, -4.0),
+          (jnp.subtract, (jnp.inf, jnp.inf)),
+          (jnp.tan, jnp.inf),
+          (jnp.true_divide, (0.0, 0.0)),
+          (operator.add, (jnp.inf, -jnp.inf)),
+          (operator.mod, (1.0, 0.0)),
+          (operator.pow, (-1.0, 0.5)),
+          (operator.sub, (jnp.inf, jnp.inf)),
+          (operator.truediv, (0.0, 0.0)),
+          # go/keep-sorted end
+      ))
+  )
+  def test_can_raise_nan_error(self, jit, f, operands):
+    operands = [jnp.float32(x) for x in operands]
 
-    f = jnp.log
     if jit:
       f = jax.jit(f)
 
     with jnp_error.error_checking_behavior(nan="raise"):
-      f(x)
+      f(*operands)
       with self.assertRaisesRegex(JaxValueError, "NaN"):
         error_check.raise_if_error()