test_util: avoid overflow errors in NumPy 2.0

jax-ml · Feb 7, 2024 · 77f030c · 77f030c
1 parent 7004497
commit 77f030c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jax/_src/test_util.py b/jax/_src/test_util.py
@@ -575,9 +575,9 @@ def _rand_dtype(rand, shape, dtype, scale=1., post=lambda x: x):
     to rand but scaled, converted to the appropriate dtype, and post-processed.
   """
   if _dtypes.issubdtype(dtype, np.unsignedinteger):
-    r = lambda: np.asarray(scale * abs(rand(*_dims_of_shape(shape))), dtype)
+    r = lambda: np.asarray(scale * abs(rand(*_dims_of_shape(shape)))).astype(dtype)
   else:
-    r = lambda: np.asarray(scale * rand(*_dims_of_shape(shape)), dtype)
+    r = lambda: np.asarray(scale * rand(*_dims_of_shape(shape)).astype(dtype)
   if _dtypes.issubdtype(dtype, np.complexfloating):
     vals = r() + 1.0j * r()
   else: