Repository: google/jax

Issue:
jax.random.randint cannot generate values equal to DTYPE_MAX
According to randint docstring: `Sample uniform random values in [minval, maxval) with given shape/dtype.`

However because of the conversion here https://github.com/google/jax/blob/5bbb449ae5849c508194c8eb5b10c101f1fa22ae/jax/_src/random.py#L439 passing `maxval={dtype max value + 1}` will generate an all-zero array.
The behavior in numpy is different (and correct):
```
jax.random.randint(jax.random.PRNGKey(0), shape=(10,), minval=0, maxval=256, dtype=jnp.uint8)
>> DeviceArray([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=uint8)

np.random.randint(size=(10,), low=0, high=256, dtype=np.uint8)
>> array([ 87, 136, 195, 254,  95,  94, 113,   7,  66,  13], dtype=uint8)
```

A workaround is to not pass the dtype to randint (or use a dtype with a higher maximum value), and cast the output:
```
jax.random.randint(jax.random.PRNGKey(0), shape=(10,), minval=0, maxval=256).astype(jnp.uint8)
>> DeviceArray([178, 233,  99, 238,  39, 149,  46, 139,  68, 142], dtype=uint8)
```


Discussion:
Thinking about this a bit... this is really tricky for JAX to handle in a general way. Consider the case of `max = 1 << 64`, which is one greater than the maximum uint64. Numpy handles it as expected:
```python
>>> np.random.randint(0, 1<<64, dtype='uint64')
17392714807819977315
```
But if you increase this bound by one, numpy returns an error:
```
>>> np.random.randint(0, 1<<64 + 1, dtype='uint64')                                                                                                                         
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-30-e02298560197> in <module>
----> 1 np.random.randint(0, 1<<64 + 1, dtype='uint64')

mtrand.pyx in numpy.random.mtrand.RandomState.randint()

_bounded_integers.pyx in numpy.random._bounded_integers._rand_uint64()

ValueError: high is out of bounds for uint64
```
Neither of these numbers (`1 << 64` and `1 << 64 + 1`) are representable as 64-bit integers or floats; numpy handles these by making use of Python's built-in arbitrary precision integers, as mentioned in the note here: https://github.com/numpy/numpy/blob/5cae51e794d69dd553104099305e9f92db237c53/numpy/random/_bounded_integers.pyx.in#L286-L292

JAX cannot do this, because it compiles to XLA which only has fixed-precision integers, and has no native way to represent larger numbers such as `1 << 64`.

How might we work around this? Some possibilities:

1. Follow NumPy and do our arithmetic in Python. The downside of this is that it would only work for concrete values of minval and maxval; i.e. you could no longer jit-compile code that passes minval and maxval as JAX arrays. You might imagine getting around this by only using the Python int logic in non-traced code; the problem there is that this would cause `randint` to act differently depending on whether it is wrapped in JIT, which is something we like to avoid.

2. Use some sort of extended representation (i.e. a pair of int64s) to represent the minval and maxval in order to handle larger bounds. However, this would require deeper changes to JAX's JIT dispatch mechanism, because currently a Python integer passed to a jitted function is **always** converted to a dtyped integer before the function can do any logic on it, because that's the only way JAX knows how to handle numerical data. I'm not sure how I'd approach changing this for cases where larger integers are required.

Neither of these is really a viable solution, and I can't think of any other alternatives.
The above shouldn't stop us from fixing the issue for smaller types, though: for example things like `random.randint(key, shape, minval=0, maxval=256, dtype='uint8')` should be workable with a few tweaks. But we would not be able to error on `maxval > 256` as numpy does, because JAX cannot branch on non-static values.
I started some work on this in #5868... it gets quite complicated to handle all the possible corner cases!
Could we fix this by adding an `inclusive`/`exclusive` option on the API?
> Could we fix this by adding an inclusive/exclusive option on the API?

That's an option, although unless users know to think about it, things like `randint(key, shape, 0, 256, 'uint8')` will continue to silently return zeros.
We should report (everywhere) when we see integer overflow casting Python integers to typed integers. Someone even started on that in the past.
That would be useful - I think currently we only get that via numpy when we overflow int64.
Also, even if we go the `inclusive / exclusive` route, I've come across some other corner cases that we'll need to fix as well, for example:
```python
from jax import random

random.randint(random.PRNGKey(0), (10,), 254, 254, dtype='uint8')
# Buffer([254, 254, 254, 254, 254, 254, 254, 254, 254, 254], dtype=uint8)

random.randint(random.PRNGKey(0), (10,), 255, 255, dtype='uint8')
# Buffer([111, 139,  85, 222,   6, 223, 108,  27,  20, 191], dtype=uint8)
```
Related: https://github.com/google/jax/issues/2006#issuecomment-575273556

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -330,6 +330,47 @@ def _promote_args_inexact(fun_name, *args):
   _check_no_float0s(fun_name, *args)
   return _promote_shapes(fun_name, *_promote_dtypes_inexact(*args))
 
+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)
+
+
 def _constant_like(x, const):
   return np.array(const, dtype=_dtype(x))
 
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/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
@@ -441,20 +441,28 @@ def randint(key: jnp.ndarray,
 def _randint(key, shape, minval, maxval, dtype):
   _check_shape("randint", shape, np.shape(minval), np.shape(maxval))
   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)
+    raise TypeError(f"randint only accepts integer dtypes, got {dtype}")
+
+  minval = _asarray(minval)
+  maxval = _asarray(maxval)
+  if not jnp.issubdtype(minval.dtype, np.integer):
+    minval = minval.astype(int)
+  if not jnp.issubdtype(maxval.dtype, np.integer):
+    maxval = maxval.astype(int)
+
+  # Flag where maxval is greater than the maximum value of dtype
+  # in order to handle cases like randint(key, shape, 0, 256, 'uint8')
+  maxval_out_of_range = lax.gt(
+    maxval, _convert_and_clip_integer(jnp.array(jnp.iinfo(dtype).max, dtype), 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
 
   if nbits not in (8, 16, 32, 64):
-    raise TypeError("randint only accepts 8-, 16-, 32-, or 64-bit dtypes.")
-
-  # if we don't have minval < maxval, just always return minval
-  # https://github.com/google/jax/issues/222
-  maxval = lax.max(lax.add(minval, np.array(1, dtype)), maxval)
+    raise TypeError(f"randint only accepts 8-, 16-, 32-, or 64-bit dtypes, got {dtype}")
 
   # This algorithm is biased whenever (maxval - minval) is not a power of 2.
   # We generate double the number of random bits required by the dtype so as to
@@ -466,6 +474,18 @@ def _randint(key, shape, minval, maxval, dtype):
   unsigned_dtype = _UINT_DTYPES[nbits]
   span = lax.convert_element_type(maxval - minval, unsigned_dtype)
 
+  # Ensure that span=1 when maxval <= minval, so minval is always returned;
+  # https://github.com/google/jax/issues/222
+  span = lax.select(maxval <= minval, lax.full_like(span, 1), span)
+
+  # When maxval is out of range, the span has to be one larger.
+  # If span is already the maximum representable value, this will wrap to zero,
+  # causing remainders below to have no effect, which is the correct semantics.
+  span = lax.select(
+    maxval_out_of_range & (maxval > minval),
+    lax.add(span, lax._const(span, 1)),
+    span)
+
   # To compute a remainder operation on an integer that might have twice as many
   # bits as we can represent in the native unsigned dtype, we compute a
   # multiplier equal to 2**nbits % span. To avoid overflow, we use the identity:


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/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,34 @@ 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)
+
+  def test_randint_out_of_range(self):
+    key = random.PRNGKey(0)
+
+    r = random.randint(key, (10,), 255, 256, np.uint8)
+    self.assertAllClose(r, jnp.full_like(r, 255))
+
+    r = random.randint(key, (1000,), -128, 128, np.int8)
+    self.assertGreater((r == -128).sum(), 0)
+    self.assertGreater((r == 127).sum(), 0)
+
+    r = random.randint(key, (1000,), -1000, 1000, np.uint8)
+    self.assertGreater((r == 0).sum(), 0)
+    self.assertGreater((r == 255).sum(), 0)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())