Repository: google/jax

Issue:
Should JAX deprecate indexing with lists?
Since numpy 1.16, indexing with a list in place of a tuple has led to a `FutureWarning` (See https://github.com/numpy/numpy/pull/9686 for a discussion of the rationale for this):
```python
>>> import numpy as np
>>> x = np.arange(6).reshape(2, 3)
>>> idx = [[0], [1]]
>>> x[idx]
FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will be interpreted as an array index, `arr[np.array(seq)]`, which will result either in an error or a different result.
array([1])
```
As mentioned in the warning, the current behavior treats the indices as identical to a tuple:
```python
>>> x[tuple(idx)]
array([1])
```
while in the future, the indices will be treated as an array:
```python
>>> x[np.array(idx)]
array([[[0, 1, 2]],
       [[3, 4, 5]]])
```

JAX currently implements the old, deprecated behavior, without any warning:
```python
>>> import jax.numpy as jnp
>>> jnp.array(x)[idx]
DeviceArray([1], dtype=int32)
```

This is setting us up for a future where numpy and JAX have different indexing semantics for lists of indices. I would propose that we follow numpy and start warning about this behavior now, so that when a numpy release finally does deprecate this indexing behavior, jax will be ready to immediately follow suit.

Thoughts?


Discussion:
👍  we should definitely deprecate this. We can probably be even more aggressive about removing support for this given the experimental nature of JAX.

In the long term I would be inclined to make _all_ indexing with lists in JAX an error. Even the case where lists are treated like an array is at odds with how JAX disallows lists as arguments to functions like `jnp.sum()`. It's not much more painful to require inserting array() in expressions like `x[:, jnp.array([1, 2])]` and has the advantage of much more explicit conversion.
+1. Let's just disallow this now.

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
@@ -4232,10 +4232,28 @@ def _eliminate_deprecated_list_indexing(idx):
   # deprecated by NumPy and exists for backward compatibility.
   if not isinstance(idx, tuple):
     if isinstance(idx, Sequence) and not isinstance(idx, ndarray):
+      # As of numpy 1.16, some non-tuple sequences of indices result in a warning, while
+      # others are converted to arrays, based on a set of somewhat convoluted heuristics
+      # (See https://github.com/numpy/numpy/blob/v1.19.2/numpy/core/src/multiarray/mapping.c#L179-L343)
+      # In JAX, we raise a warning for *all* non-tuple sequences, and in the future will
+      # *always* raise a TypeError here, rather than silently converting to an array or tuple
+      # depending on the contents of the list as numpy will. "Explicit is better than implicit".
+      # TODO(jakevdp): raise a TypeError here.
       if _any(_should_unpack_list_index(i) for i in idx):
+        msg = ("Using a non-tuple sequence for multidimensional indexing is deprecated; "
+               "use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this will "
+               "result in a TypeError. See https://github.com/google/jax/issues/4564 "
+               "for discussion of why this type of indexing is being deprecated.")
         idx = tuple(idx)
       else:
+        msg = ("Using a non-tuple sequence for multidimensional indexing is deprecated; "
+               "use `arr[array(seq)]` instead of `arr[seq]`. In the future this will "
+               "result in a TypeError. See https://github.com/google/jax/issues/4564 "
+               "for discussion of why this type of indexing is being deprecated.")
         idx = (idx,)
+      # TODO(jakevdp): this stacklevel is appropriate for x[idx]; for ops.index_update
+      # we should use stacklevel=5; for x.at[idx].set() we should use stacklevel=6.
+      warnings.warn(msg, FutureWarning, stacklevel=4)
     else:
       idx = (idx,)
   return 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
@@ -14,6 +14,7 @@
 
 
 import collections
+from contextlib import contextmanager
 import enum
 from functools import partial
 import itertools
@@ -458,7 +459,7 @@ def testStaticIndexing(self, shape, dtype, indexer):
     jnp_fun = lambda x: jnp.asarray(x)[indexer]
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
-    self._CompileAndCheck(jnp_fun, args_maker)
+      self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters({
       "testcase_name":
@@ -562,7 +563,7 @@ def jnp_fun(x, unpacked_indexer):
     args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
-    self._CompileAndCheck(jnp_fun, args_maker)
+      self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(
       {"testcase_name": "{}_inshape={}_indexer={}"
@@ -613,7 +614,7 @@ def testAdvancedIntegerIndexing(self, shape, dtype, indexer):
     jnp_fun = lambda x, idx: jnp.asarray(x)[idx]
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
-    self._CompileAndCheck(jnp_fun, args_maker)
+      self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(
       {"testcase_name": "{}_inshape={}_indexer={}"
@@ -672,7 +673,8 @@ def testAdvancedIntegerIndexingGrads(self, shape, dtype, indexer):
     tol = 1e-2 if jnp.finfo(dtype).bits == 32 else None
     arg = rng(shape, dtype)
     fun = lambda x: jnp.asarray(x)[indexer]
-    check_grads(fun, (arg,), 2, tol, tol, eps=1.)
+    with suppress_deprecated_indexing_warnings():
+      check_grads(fun, (arg,), 2, tol, tol, eps=1.)
 
   @parameterized.named_parameters(
       {"testcase_name": "{}_inshape={}_indexer={}"
@@ -699,7 +701,7 @@ def np_fun(x, indexer_with_dummies):
 
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
-    self._CompileAndCheck(jnp_fun, args_maker)
+      self._CompileAndCheck(jnp_fun, args_maker)
 
   def testAdvancedIndexingManually(self):
     x = np.random.RandomState(0).randn(3, 4, 5)
@@ -752,7 +754,8 @@ def testBooleanIndexingArray1D(self):
   def testBooleanIndexingList1D(self):
     idx = [True, True, False]
     x = api.device_put(np.arange(3))
-    ans = x[idx]
+    with suppress_deprecated_indexing_warnings():
+      ans = x[idx]
     expected = np.arange(3)[idx]
     self.assertAllClose(ans, expected, check_dtypes=False)
 
@@ -766,7 +769,8 @@ def testBooleanIndexingArray2DBroadcast(self):
   def testBooleanIndexingList2DBroadcast(self):
     idx = [True, True, False, True]
     x = np.arange(8).reshape(4, 2)
-    ans = api.device_put(x)[idx]
+    with suppress_deprecated_indexing_warnings():
+      ans = api.device_put(x)[idx]
     expected = x[idx]
     self.assertAllClose(ans, expected, check_dtypes=False)
 
@@ -946,7 +950,7 @@ def testStaticIndexing(self, shape, dtype, update_shape, update_dtype,
       jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y)
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fn, jax_fn, args_maker)
-    self._CompileAndCheck(jax_fn, args_maker)
+      self._CompileAndCheck(jax_fn, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list({
       "testcase_name": "{}_inshape={}_indexer={}_update={}_sugared={}_op={}".format(
@@ -973,7 +977,7 @@ def testAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
       jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y, unique_indices=True)
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fn, jax_fn, args_maker)
-    self._CompileAndCheck(jax_fn, args_maker)
+      self._CompileAndCheck(jax_fn, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list({
       "testcase_name": "{}_inshape={}_indexer={}_update={}_sugared={}_op={}".format(
@@ -1002,7 +1006,7 @@ def testAdvancedIndexingSorted(self, shape, dtype, update_shape, update_dtype,
           op, indexer, x, y, indices_are_sorted=True, unique_indices=True)
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, check_dtypes=True)
-    self._CompileAndCheck(jax_fn, args_maker, check_dtypes=True)
+      self._CompileAndCheck(jax_fn, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list({
       "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}_sugared={}".format(
@@ -1029,7 +1033,7 @@ def testMixedAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
       jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y)
     with suppress_deprecated_indexing_warnings():
       self._CheckAgainstNumpy(np_fn, jax_fn, args_maker)
-    self._CompileAndCheck(jax_fn, args_maker)
+      self._CompileAndCheck(jax_fn, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list({
       "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
@@ -1105,6 +1109,40 @@ def testIndexDtypeError(self):
       jnp.zeros(5).at[::2].set(1)
       self.assertLen(w, 0)
 
+  @contextmanager
+  def assertNoWarnings(self):
+    with warnings.catch_warnings(record=True) as caught_warnings:
+      yield
+    self.assertEmpty(caught_warnings)
+
+  @parameterized.named_parameters(jtu.cases_from_list({
+      "testcase_name": "idx={}".format(idx), "idx": idx, "idx_type": idx_type}
+    for idx, idx_type in [
+      ([0], "array"),
+      ([0, 0], "array"),
+      ([[0, 0]], "tuple"),
+      ([0, [0, 1]], "tuple"),
+      ([0, np.arange(2)], "tuple"),
+      ([0, None], "tuple"),
+      ([0, slice(None)], "tuple"),
+    ]))
+  def testIndexSequenceDeprecationWarning(self, idx, idx_type):
+    msg = fr"Using a non-tuple sequence for multidimensional indexing is deprecated.*arr\[{idx_type}\(seq\)\]"
+    normalize = {"array": np.array, "tuple": tuple}[idx_type]
+    x = jnp.arange(6).reshape(3, 2)
+
+    with self.assertWarnsRegex(FutureWarning, msg):
+      idx_get = x[idx]
+    with self.assertNoWarnings():
+      idx_get_norm = x[normalize(idx)]
+    self.assertArraysEqual(idx_get, idx_get_norm)
+
+    with self.assertWarnsRegex(FutureWarning, msg):
+      idx_set = x.at[idx].set(0)
+    with self.assertNoWarnings():
+      idx_set_norm = x.at[normalize(idx)].set(0)
+    self.assertArraysEqual(idx_set, idx_set_norm)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
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
@@ -3696,7 +3696,7 @@ def testIssue764(self):
   def testIssue776(self):
     """Tests that the scatter-add transpose rule instantiates symbolic zeros."""
     def f(u):
-      y = jax.ops.index_add(np.ones(10,), [2, 4, 5], u)
+      y = jnp.ones(10).at[np.array([2, 4, 5])].add(u)
       # The transpose rule for lax.tie_in returns a symbolic zero for its first
       # argument.
       return lax.tie_in(y, 7.)
@@ -4008,7 +4008,7 @@ def testLinspaceEndpoints(self, dtype, rng_factory):
     rng = rng_factory(self.rng())
     endpoints = rng((2,), dtype)
     out = jnp.linspace(*endpoints, 10, dtype=dtype)
-    self.assertAllClose(out[[0, -1]], endpoints, rtol=0, atol=0)
+    self.assertAllClose(out[np.array([0, -1])], endpoints, rtol=0, atol=0)
 
   @parameterized.named_parameters(
       jtu.cases_from_list(