Repository: google/jax

Issue:
Raise OverflowError in X32 mode for large Python integers
Python integers are arbitrary width, meaning that they can represent values that are not representable by numpy's `int64`. In these cases, numpy raises an overflow error. For example:
```python
>>> import numpy as np
>>> np.array(1 << 65, dtype=int)
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-36-bc194570c3ad> in <module>
----> 1 np.array(1 << 65, dtype=int)

OverflowError: Python int too large to convert to C long
```
In X64 mode, JAX inherits this property via numpy:
```python
>>> import jax.numpy as jnp
>>> from jax import config; config.update('jax_enable_x64', True)
>>> np.array(1 << 65, dtype=int)
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-3-e4d8a251ed7e> in <module>
----> 1 jnp.array(1 << 65, dtype=int)

~/github/google/jax/jax/_src/numpy/lax_numpy.py in array(object, dtype, copy, order, ndmin)
   2874 
   2875   if _can_call_numpy_array(object):
-> 2876     object = _np_array(object, dtype=dtype, ndmin=ndmin, copy=False)
   2877 
   2878   assert type(object) not in dtypes.python_scalar_dtypes

~/github/google/jax/jax/_src/numpy/lax_numpy.py in _np_array(obj, dtype, **kwargs)
    228   uses Jax's default dtypes.
    229   """
--> 230   arr = np.array(obj, dtype=dtype, **kwargs)
    231   obj_dtype = getattr(obj, 'dtype', None)
    232   arr_dtype = np.dtype(arr.dtype).type

OverflowError: Python int too large to convert to C long
```
With X64 disabled, however, values that are too large to be stored in `int32` are silently truncated to zero:
```python
>>> import jax.numpy as jnp
>>> jnp.array(1 << 33)
DeviceArray(0, dtype=int32)
```
Such silent data loss is problematic, particularly because X32 is the default mode.

This PR adds a value check when Python integers are coerced to arrays, so that an `OverflowError` is raised as appropriate for the default integer type:
```python
>>> import jax.numpy as jnp
>>> jnp.array(1 << 33)
---------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
<ipython-input-2-50cb794a933c> in <module>
----> 1 jnp.array(1 << 33)

~/github/google/jax/jax/_src/numpy/lax_numpy.py in array(object, dtype, copy, order, ndmin)
   2872 
   2873   if _can_call_numpy_array(object):
-> 2874     object = dtypes.coerce_to_array(object, dtype=dtype, ndmin=ndmin, copy=False)
   2875 
   2876   assert type(object) not in dtypes.python_scalar_dtypes

~/github/google/jax/jax/dtypes.py in coerce_to_array(x, dtype, ndmin, copy)
    125     info = np.iinfo(dtype)
    126     if not info.min <= x <= info.max:
--> 127       raise OverflowError(f"Python int {x} too large to convert to {dtype}")
    128   return np.array(x, dtype=dtype, ndmin=ndmin, copy=copy)
    129 

OverflowError: Python int 8589934592 too large to convert to int32
```
This also required a minor change to `random.PRNGseed` as well, as the previous behavior was built on the assumption that out-of-bounds inputs were truncated rather than leading to an error.


Discussion:
@hawkinsp – I requested your review even though this still has a number of downstream impacts. If this change is something we want, I can start fixing/mitigating those. Thanks!

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -59,6 +59,11 @@ def PRNGKey(seed: int) -> jnp.ndarray:
     key is constructed from a 64-bit seed by effectively bit-casting to a pair
     of uint32 values (or from a 32-bit seed by first padding out with zeros).
   """
+  # Avoid overflowerror in X32 mode by first converting ints to int64.
+  # This breaks JIT invariance of PRNGKey for large ints, but supports the
+  # common use-case of instantiating PRNGKey with Python hashes in X32 mode.
+  if isinstance(seed, int):
+    seed = np.int64(seed)
   seed_arr = jnp.asarray(seed)
   if seed_arr.shape:
     raise TypeError(f"PRNGKey seed must be a scalar; got {seed!r}.")
@@ -279,7 +284,7 @@ def fold_in(key: jnp.ndarray, data: int) -> jnp.ndarray:
     A new PRNGKey that is a deterministic function of the inputs and is
     statistically safe for producing a stream of new pseudo-random values.
   """
-  return _fold_in(key, data)
+  return _fold_in(key, jnp.uint32(data))
 
 @jit
 def _fold_in(key, data):
diff --git a/jax/dtypes.py b/jax/dtypes.py
--- a/jax/dtypes.py
+++ b/jax/dtypes.py
@@ -121,8 +121,7 @@ def _scalar_type_to_dtype(typ: type, value: Any = None):
   ---------------------------------------------------------------------------
   OverflowError: Python int 9223372036854775808 too large to convert to int64
   """
-  dtype = python_scalar_dtypes[typ]
-  # TODO(jakevdp): use proper overflow for int32.
+  dtype = canonicalize_dtype(python_scalar_dtypes[typ])
   if typ is int and value is not None:
     if value < np.iinfo(dtype).min or value > np.iinfo(dtype).max:
       raise OverflowError(f"Python int {value} too large to convert to {dtype}")


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -5031,30 +5031,28 @@ def test_partial_eval(self):
       for jit_type in [None, "python", "cpp"]
       if not (jit_type is None and func == 'identity')))
   def test_integer_overflow(self, jit_type, func):
-    def jit(f, **kwargs):
-      if jit_type is None:
-        return f
-      elif jit_type == "python":
-        return api._python_jit(f, **kwargs)
-      elif jit_type == "cpp":
-        return api._cpp_jit(f, **kwargs)
-      else:
-        raise ValueError(f"invalid jit_type={jit_type}")
-    func = jit({
+    if jit_type == "cpp" and not config.x64_enabled and jax.lib.version < (0, 1, 65):
+      self.skipTest("int32 overflow not yet implemented in CPP JIT.")
+    funcdict = {
       'identity': lambda x: x,
       'asarray': jnp.asarray,
-      'device_put': api.device_put
-    }[func])
-
-    int64_max = np.iinfo(np.int64).max
-    int64_min = np.iinfo(np.int64).min
-
-    int_dtype = dtypes.canonicalize_dtype(np.int64)
-
-    self.assertEqual(func(int64_max).dtype, int_dtype)
-    self.assertEqual(func(int64_min).dtype, int_dtype)
-    self.assertRaises(OverflowError, func, int64_max + 1)
-    self.assertRaises(OverflowError, func, int64_min - 1)
+      'device_put': api.device_put,
+    }
+    jit = {
+      'python': api._python_jit,
+      'cpp': api._cpp_jit,
+      None: lambda x: x,
+    }
+    f = jit[jit_type](funcdict[func])
+
+    int_dtype = dtypes.canonicalize_dtype(jnp.int_)
+    int_max = np.iinfo(int_dtype).max
+    int_min = np.iinfo(int_dtype).min
+
+    self.assertEqual(f(int_max).dtype, int_dtype)
+    self.assertEqual(f(int_min).dtype, int_dtype)
+    self.assertRaises(OverflowError, f, int_max + 1)
+    self.assertRaises(OverflowError, f, int_min - 1)
 
 
 if __name__ == '__main__':
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -3275,17 +3275,17 @@ def testArrayUnsupportedDtypeError(self):
       jnp.array(3, [('a','<i4'),('b','<i4')])
 
   def testArrayFromInteger(self):
-    # TODO(jakevdp): implement X32 overflow and canonicalize these
-    int_max = jnp.iinfo(jnp.int64).max
-    int_min = jnp.iinfo(jnp.int64).min
+    int_dtype = dtypes.canonicalize_dtype(jnp.int64)
+    int_max = jnp.iinfo(int_dtype).max
+    int_min = jnp.iinfo(int_dtype).min
 
     # Values at extremes are converted correctly.
     for val in [int_min, 0, int_max]:
-      self.assertEqual(jnp.array(val).dtype, dtypes.canonicalize_dtype('int64'))
+      self.assertEqual(jnp.array(val).dtype, int_dtype)
 
     # out of bounds leads to an OverflowError
     val = int_max + 1
-    with self.assertRaisesRegex(OverflowError, f"Python int {val} too large to convert to int64"):
+    with self.assertRaisesRegex(OverflowError, f"Python int {val} too large to convert to {int_dtype.name}"):
       jnp.array(val)
 
     # explicit uint64 should work
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -946,6 +946,9 @@ def f(x):
     ]
   ))
   def test_prng_seeds_and_keys(self, seed, type, jit, key):
+    if (jit and type is int and not config.x64_enabled and
+        (seed < np.iinfo('int32').min or seed > np.iinfo('int32').max)):
+      self.skipTest("Expected failure: integer out of range for jit.")
     seed = type(seed)
     if jit:
       actual = api.jit(random.PRNGKey)(seed)
@@ -961,6 +964,8 @@ def test_prng_seeds_and_keys(self, seed, type, jit, key):
   def test_prng_jit_invariance(self, seed, type):
     if type == "int" and seed == (1 << 64) - 1:
       self.skipTest("Expected failure: Python int too large.")
+    if not config.x64_enabled and seed > np.iinfo(np.int32).max:
+      self.skipTest("Expected failure: Python int too large.")
     type = {"int": int, "np.array": np.array, "jnp.array": jnp.array}[type]
     args_maker = lambda: [type(seed)]
     self._CompileAndCheck(random.PRNGKey, args_maker)