Repository: google/jax

Issue:
Add np.unravel_index
This is what I've been using. If it looks reasonable, will see about putting together a PR.

```python
def unravel_index(flat_index, shape):
    """
    To support abstract `flat_index` there is no bounds checking. For
    overflow some dimensions of the resulting index will be out of bounds.
    For negative `flat_index` the resulting index is valid, although the first
    dimension will be negative. 
    """

    if len(shape) == 0:
        return ()

    size_of_next_dimension = np.pad(shape[1:], (0, 1), constant_values=1)
    strides = np.cumprod(size_of_next_dimension[::-1], dtype='int32')[::-1]

    _, idx = jax.lax.scan(
        lambda remaining, stride: (
            remaining % stride,
            remaining // stride
        ),
        flat_index,
        strides
    )

    return tuple(idx)
```

Haven't looked at jax testing yet, would be helpful to know whether / how to incorporate these...

```python
def test_unravel_index():

    assert unravel_index(0, (2, 3)) == (0, 0)
    assert unravel_index(5, (2, 1, 3)) == (1, 0, 2)

    assert unravel_index(0, ()) == ()
    assert unravel_index(3, (2,)) == (3,)
    assert unravel_index(-1, (2, 1, 3)) == (-1, 0, 2)
```


Discussion:
While it's probably correct as is, I think it would be preferable to avoid the `scan`, along the following lines:

```
In [72]: np.unravel_index(692, [4,6,8,9])
Out[72]: (1, 3, 4, 8)

In [73]: index = 692

In [74]: shape = [4, 6, 8, 9]

In [75]: cp = np.cumprod(np.pad(shape, [(0, 1)], constant_values=1)[::-1])[::-1]

In [76]: index % cp[:-1] // cp[1:]
Out[76]: array([1, 3, 4, 8])
```

In general vectorizable operations will be faster than `scan`.

Want to send a PR? Tests would go in `lax_numpy_test.py`.

Thanks @hawkinsp 

Any preference on overflow behavior?

1. `np.unravel_index(3, (2,)) == (3,)`
2. `np.unravel_index(3, (2,)) == (2,)`
3. `np.unravel_index(3, (2,)) == (1,)`

Both (1) and (2) seem consistent with "JAX does not raise an error and instead returns the last value in the array". And (1) is FWIW what I expected to see when I first tried `jax.numpy.unravel_index`



Hi @mattwescott 

Are you still looking into this? I could use this feature as well.
@rdaems thanks for the motivation, will put together a PR shortly

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -390,7 +390,7 @@ def fmax(x1, x2):
   return where((x1 > x2) | isnan(x2), x1, x2)
 
 @_wraps(onp.finfo)
-def finfo(dtype): 
+def finfo(dtype):
   return dtypes.finfo(dtype)
 
 @_wraps(onp.issubdtype)
@@ -724,7 +724,7 @@ def _conv(x, y, mode, op, precision):
   if ndim(x) != 1 or ndim(y) != 1:
     raise ValueError(f"{op}() only support 1-dimensional inputs.")
   x, y = _promote_dtypes_inexact(x, y)
-  
+
   out_order = slice(None)
   if len(x) < len(y):
     x, y = y, x
@@ -1153,6 +1153,25 @@ def ravel(a, order="C"):
   return reshape(a, (size(a),), order)
 
 
+_UNRAVEL_INDEX_DOC = """\
+Unlike numpy's implementation of unravel_index, negative indices are accepted
+and out-of-bounds indices are clipped.
+"""
+
+@_wraps(onp.unravel_index, lax_description=_UNRAVEL_INDEX_DOC)
+def unravel_index(indices, shape):
+  indices = asarray(indices)
+  sizes = pad(shape, (0, 1), constant_values=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 flat_index
+  cumulative_sizes = cumulative_sizes.reshape([-1] + [1] * indices.ndim)
+  idx = clipped_indices % cumulative_sizes[:-1] // cumulative_sizes[1:]
+  return tuple(idx)
+
+
 @_wraps(onp.squeeze)
 def squeeze(a, axis=None):
   shape_a = shape(a)


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
@@ -1992,6 +1992,25 @@ def testRavel(self):
     args_maker = lambda: [rng.randn(3, 4).astype("float32")]
     self._CompileAndCheck(lambda x: x.ravel(), args_maker, check_dtypes=True)
 
+  @parameterized.parameters(
+    (0, (2, 1, 3)),
+    (5, (2, 1, 3)),
+    (0, ()),
+    ([0, 1, 2], (2, 2)),
+    ([[[0, 1], [2, 3]]], (2, 2)))
+  def testUnravelIndex(self, flat_index, shape):
+    self._CheckAgainstNumpy(
+      onp.unravel_index,
+      jnp.unravel_index,
+      lambda: (flat_index, shape),
+      check_dtypes=True
+    )
+
+  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,))
+
   def testAstype(self):
     rng = onp.random.RandomState(0)
     args_maker = lambda: [rng.randn(3, 4).astype("float32")]