random.randint: clip rather than wrap out-of-bounds min/max

jax/_src/numpy/lax_numpy.py

@@ -333,6 +333,46 @@ def _promote_args_inexact(fun_name, *args):
 def _constant_like(x, const):
   return np.array(const, dtype=_dtype(x))
 
+def _convert_and_clip_integer(val, dtype):
+  """
+  Convert integer-typed val to specified integer dtype, clipping to dtype
+  range rather than wrapping.
+
+  Args:
+    val: value to be converted
+    dtype: dtype of output
+
+  Returns:
+    equivalent of val in new dtype
+
+  Examples
+  --------
+  Normal integer type conversion will wrap:
+
+  >>> val = jnp.uint32(0xFFFFFFFF)
+  >>> val.astype('int32')
+  DeviceArray(-1, dtype=int32)
+
+  This function clips to the values representable in the new type:
+
+  >>> _convert_and_clip_integer(val, 'int32')
+  DeviceArray(2147483647, dtype=int32)
+  """
+  val = val if isinstance(val, ndarray) else asarray(val)
+  dtype = dtypes.canonicalize_dtype(dtype)
+  if not (issubdtype(dtype, integer) and issubdtype(val.dtype, integer)):
+    raise TypeError("_convert_and_clip_integer only accepts integer dtypes.")
+
+  val_dtype = dtypes.canonicalize_dtype(val.dtype)
+  if val_dtype != val.dtype:
+    # TODO(jakevdp): this is a weird corner case; need to figure out how to handle it.
+    # This happens in X32 mode and can either come from a jax value created in another
+    # context, or a Python integer converted to int64.
+    pass
+  min_val = _constant_like(val, _max(iinfo(dtype).min, iinfo(val_dtype).min))
+  max_val = _constant_like(val, _min(iinfo(dtype).max, iinfo(val_dtype).max))
+  return clip(val, min_val, max_val).astype(dtype)
+
 ### implementations of numpy functions in terms of lax
 
 @_wraps(np.fmin)

jax/_src/random.py

@@ -25,7 +25,7 @@
 from jax import dtypes
 from jax.core import NamedShape
 from jax.api import jit, vmap
-from jax._src.numpy.lax_numpy import _constant_like, asarray
+from jax._src.numpy.lax_numpy import _constant_like, _convert_and_clip_integer, asarray
 from jax.lib import xla_bridge
 from jax.lib import xla_client
 from jax.lib import cuda_prng
@@ -443,8 +443,8 @@ def _randint(key, shape, minval, maxval, dtype):
   if not jnp.issubdtype(dtype, np.integer):
     raise TypeError("randint only accepts integer dtypes.")
 
-  minval = lax.convert_element_type(minval, dtype)
-  maxval = lax.convert_element_type(maxval, dtype)
+  minval = _convert_and_clip_integer(minval, dtype)
+  maxval = _convert_and_clip_integer(maxval, dtype)
   minval = lax.broadcast_to_rank(minval, len(shape))
   maxval = lax.broadcast_to_rank(maxval, len(shape))
   nbits = jnp.iinfo(dtype).bits

tests/random_test.py

@@ -26,6 +26,7 @@
 
 from jax import api
 from jax import core
+from jax import dtypes
 from jax import grad
 from jax import lax
 from jax import numpy as jnp
@@ -983,6 +984,20 @@ def test_random_split_doesnt_device_put_during_tracing(self):
       api.jit(random.split)(key)
     self.assertEqual(count[0], 1)  # 1 for the argument device_put
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": f"_dtype={dtype}", "dtype": dtype}
+      for dtype in int_dtypes + uint_dtypes))
+  def test_randint_bounds(self, dtype):
+    min = np.iinfo(dtype).min
+    max = np.iinfo(dtype).max
+    key = random.PRNGKey(1701)
+    shape = (10,)
+    if np.iinfo(dtype).bits < np.iinfo(dtypes.canonicalize_dtype(int)).bits:
+      expected = random.randint(key, shape, min, max, dtype)
+      self.assertArraysEqual(expected, random.randint(key, shape, min - 12345, max + 12345, dtype))
+    else:
+      self.assertRaises(OverflowError, random.randint, key, shape, min - 12345, max + 12345, dtype)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())