Merge pull request #19387 from jakevdp:searchsorted-test

PiperOrigin-RevId: 598935247
jax-ml · Jan 16, 2024 · 08837a9 · 08837a9
2 parents 7f661de + 5bad0db
commit 08837a9
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
@@ -2560,11 +2560,15 @@ def testSearchsortedDtype(self):
     if config.enable_x64.value:
       out_int64 = jax.eval_shape(jnp.searchsorted, a_int64, v)
       self.assertEqual(out_int64.dtype, np.int64)
-    else:
+    elif jtu.numpy_version() < (2, 0, 0):
       with self.assertWarnsRegex(UserWarning, "Explicitly requested dtype int64"):
         with jtu.ignore_warning(category=DeprecationWarning,
                                 message="NumPy will stop allowing conversion.*"):
           out_int64 = jax.eval_shape(jnp.searchsorted, a_int64, v)
+    else:
+      with self.assertWarnsRegex(UserWarning, "Explicitly requested dtype int64"):
+        with self.assertRaisesRegex(OverflowError, "Python integer 2147483648 out of bounds.*"):
+          out_int64 = jax.eval_shape(jnp.searchsorted, a_int64, v)
 
   @jtu.sample_product(
     dtype=inexact_dtypes,