Repository: google/jax

Issue:
Failure while indexing across both dimensions of a 2D jax.lax.FilledConstant array
The same indexing works when applied to a numpy.ndarray. According to the README, all NumPy indexing (excepting indexed assignments) should work.

Screenshots:

Indexing works in original numpy but fails with jax's numpy
https://drive.google.com/file/d/1bpj4RWV9_oFajeEDAN43IdiN9VHeRHJE/view?usp=sharing

Indexing over one dimension works, but not over two:
https://drive.google.com/file/d/10nCNyIIj2Gr9vOcSKV0p203hfKqEuQd_/view?usp=sharing

Any idea why this is happening?


Discussion:
Thanks for reporting this! This must be something I broke in #140. I'll take a look...
Nevermind, it's not #140, and actually is just a bug in our indexing implementation. In particular, this fails the same way:

```python
import numpy as onp
from jax import jit
jit(lamda x: x[[0, 2, 4], [0, 2, 4]])(onp.ones((5, 5))
```
This is just a case we never handled correctly, due to one incorrect `isinstance(idx, list)` that should have been an `isinstance(idx, (list, tuple))`.

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
@@ -1510,7 +1510,7 @@ def _rewriting_take(arr, idx, axis=0):
   # Handle integer array indexing *without* ellipsis/slices/nones
   # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#integer-array-indexing
   if _is_advanced_int_indexer_without_slices(idx):
-    if isinstance(idx, list):
+    if isinstance(idx, (tuple, list)):
       if _any(_shape(e) for e in idx):
         # At least one sequence element in the index list means broadcasting.
         idx = broadcast_arrays(*idx)
@@ -1521,7 +1521,7 @@ def _rewriting_take(arr, idx, axis=0):
       # The indexer is just a single integer array.
       idx = [idx]
 
-    flat_idx = tuple(mod(ravel(x), arr.shape[i]) for i, x in enumerate(idx))
+    flat_idx = tuple([mod(ravel(x), arr.shape[i]) for i, x in enumerate(idx)])
     out = lax.index_take(arr, flat_idx, tuple(range(len(idx))))
     return lax.reshape(out, idx[0].shape + _shape(arr)[len(idx):])
 


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -398,6 +398,10 @@ def fun(unpacked_indexer, x):
            [IndexSpec(shape=(3, 4, 5), indexer=[[0, 1]]),
             IndexSpec(shape=(3, 4, 5), indexer=[[[0], [-1]], [[2, 3, 0, 3]]]),
             ]),
+          ("TupleOfListsOfPythonInts",
+           [IndexSpec(shape=(3, 4, 5), indexer=([0, 1])),
+            IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0, 3]])),
+            ]),
           ("ListOfPythonIntsAndIntArrays",
            [IndexSpec(shape=(3, 4, 5), indexer=[0, onp.array([0, 1])]),
             IndexSpec(shape=(3, 4, 5), indexer=[0, 1,
@@ -630,6 +634,15 @@ def testBooleanIndexingDynamicShapeError(self):
     i = onp.array([True, True, False])
     self.assertRaises(IndexError, lambda: api.jit(lambda x, i: x[i])(x, i))
 
+  def testIssue187(self):
+    x = lnp.ones((5, 5))
+    x[[0, 2, 4], [0, 2, 4]]  # doesn't crash
+
+    x = onp.arange(25).reshape((5, 5))
+    ans = api.jit(lambda x: x[[0, 2, 4], [0, 2, 4]])(x)
+    expected = x[[0, 2, 4], [0, 2, 4]]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == "__main__":
   absltest.main()