Merge pull request #19701 from jakevdp:fix-jtu

PiperOrigin-RevId: 605124959
jax-ml · Feb 8, 2024 · bbeddbd · bbeddbd
2 parents 1e31c15 + bc5da40
commit bbeddbd
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: