Repository: google/jax

Issue:
Unexpected behaviour in `jnp.unravel_index`
I'm getting erroneous values for `jnp.unravel_index`. For an index of `0`, some  shapes give different results than the expected `(0,0)`. Some examples:

```python
import jax.numpy as jnp

side = 10
jnp.unravel_index(0, (side, side))
# (DeviceArray(0, dtype=int32), DeviceArray(0, dtype=int32))

side = 201_995
jnp.unravel_index(0, (side, side))
# (DeviceArray(0, dtype=int32), DeviceArray(0, dtype=int32))

side = 201_996
jnp.unravel_index(0, (side, side))
# (DeviceArray(-1, dtype=int32), DeviceArray(130531, dtype=int32))

side = 205_000
jnp.unravel_index(0, (side, side))
# (DeviceArray(-1, dtype=int32), DeviceArray(82039, dtype=int32))

side = 500_000
jnp.unravel_index(0, (side, side))
# (DeviceArray(0, dtype=int32), DeviceArray(0, dtype=int32))
```

I'm using Jax version `0.3.6`.


Discussion:
This is due to int32 arithmetic overflowing for arrays this large. In your incorrect example the total size, i.e. square of `side`, is greater than the maximal value representable by a 32-bit integer. If you want to do those calculations on the host, then it's better to stick to regular `numpy` (then you also won't unnecessarily do scalar computations on the device). If you really need to perform those operations on the device, please try [running with X64 mode](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision) and perform this compute with larger integer types.
Agreed with that synopsis, but I wonder if we could make the implementation safer with respect to potential integer overflow? After all, neither the inputs nor the expected outputs in this case are out of the representable range.
Another potential improvement: in the common case of a static `shape` argument, we should be able to warn or error at trace time if there is not enough range to represent the output.

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
@@ -808,22 +808,22 @@ def ravel_multi_index(multi_index, dims, mode='raise', order='C'):
 
 _UNRAVEL_INDEX_DOC = """\
 Unlike numpy's implementation of unravel_index, negative indices are accepted
-and out-of-bounds indices are clipped.
+and out-of-bounds indices are clipped into the valid range.
 """
 
 @_wraps(np.unravel_index, lax_description=_UNRAVEL_INDEX_DOC)
 def unravel_index(indices, shape):
   _check_arraylike("unravel_index", indices)
-  sizes = append(array(shape), 1)
-  cumulative_sizes = cumprod(sizes[::-1])[::-1]
-  total_size = cumulative_sizes[0]
-  # Clip so raveling and unraveling an oob index will not change the behavior
-  clipped_indices = clip(indices, -total_size, total_size - 1)
-  # Add enough trailing dims to avoid conflict with clipped_indices
-  cumulative_sizes = expand_dims(cumulative_sizes, range(1, 1 + _ndim(indices)))
-  clipped_indices = expand_dims(clipped_indices, axis=0)
-  idx = clipped_indices % cumulative_sizes[:-1] // cumulative_sizes[1:]
-  return tuple(idx)
+  shape = atleast_1d(shape)
+  if shape.ndim != 1:
+    raise ValueError("unravel_index: shape should be a scalar or 1D sequence.")
+  out_indices = [None] * len(shape)
+  for i, s in reversed(list(enumerate(shape))):
+    indices, out_indices[i] = divmod(indices, s)
+  oob_pos = indices > 0
+  oob_neg = indices < -1
+  return tuple(where(oob_pos, s - 1, where(oob_neg, 0, i))
+               for s, i in zip(shape, out_indices))
 
 @_wraps(np.resize)
 @partial(jit, static_argnames=('new_shape',))


Test patch:
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
@@ -4216,22 +4216,26 @@ def jnp_fun(a, c):
     else:
       self._CompileAndCheck(jnp_fun, args_maker)
 
-  @parameterized.parameters(
-    (0, (2, 1, 3)),
-    (5, (2, 1, 3)),
-    (0, ()),
-    (np.array([0, 1, 2]), (2, 2)),
-    (np.array([[[0, 1], [2, 3]]]), (2, 2)))
-  def testUnravelIndex(self, flat_index, shape):
-    args_maker = lambda: (flat_index, shape)
-    np_fun = jtu.with_jax_dtype_defaults(np.unravel_index, use_defaults=not hasattr(flat_index, 'dtype'))
-    self._CheckAgainstNumpy(np_fun, jnp.unravel_index, args_maker)
-    self._CompileAndCheck(jnp.unravel_index, args_maker)
-
-  def testUnravelIndexOOB(self):
-    self.assertEqual(jnp.unravel_index(2, (2,)), (1,))
-    self.assertEqual(jnp.unravel_index(-2, (2, 1, 3,)), (1, 0, 1))
-    self.assertEqual(jnp.unravel_index(-3, (2,)), (0,))
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_shape={}_idx={}".format(shape,
+        jtu.format_shape_dtype_string(idx_shape, dtype)),
+       "shape": shape, "idx_shape": idx_shape, "dtype": dtype}
+      for shape in nonempty_nonscalar_array_shapes
+      for dtype in int_dtypes
+      for idx_shape in all_shapes))
+  def testUnravelIndex(self, shape, idx_shape, dtype):
+    size = prod(shape)
+    rng = jtu.rand_int(self.rng(), low=-((2 * size) // 3), high=(2 * size) // 3)
+
+    def np_fun(index, shape):
+      # Adjust out-of-bounds behavior to match jax's documented behavior.
+      index = np.clip(index, -size, size - 1)
+      index = np.where(index < 0, index + size, index)
+      return np.unravel_index(index, shape)
+    jnp_fun = jnp.unravel_index
+    args_maker = lambda: [rng(idx_shape, dtype), shape]
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
+    self._CompileAndCheck(jnp_fun, args_maker)
 
   def testAstype(self):
     rng = self.rng()