Repository: google/jax

Issue:
Possible bug in `jax.lax.index_take`
The `jax.lax.index_take` will fail when `axes` is a list, for example,

```py
import jax
def fn(src):
    idxs = jax.numpy.array([[2, 2], [31, 0]], dtype=jax.numpy.int32)
    axes = [0, 1]
    return jax.lax.index_take(src, idxs, axes)
genkey = jax.random.PRNGKey(90376501)
src = jax.random.randint(genkey, [3, 4, 5], -32, 32, jax.numpy.int32)
fn(src)
``` 

```
File ~/torch-1.12.1/lib/python3.9/site-packages/jax/_src/lax/slicing.py:281, in gather(operand, start_indices, dimension_numbers, slice_sizes, unique_indices, indices_are_sorted, mode, fill_value)
    279 else:
    280   fill_value = None
--> 281 return gather_p.bind(
...
    220   return f(*args, **kwargs)
    221 else:
--> 222   return cached(config._trace_context(), *args, **kwargs)

TypeError: unhashable type: 'list'
```

However, the document said `axes` is a sequence of integer. Thus, I think it should succeed with integer list.

---

Plus, its `jit` version can succeed with the same input

```py
import jax
def fn(src):
    idxs = jax.numpy.array([[2, 2], [31, 0]], dtype=jax.numpy.int32)
    axes = [0, 1]
    return jax.lax.index_take(src, idxs, axes)
genkey = jax.random.PRNGKey(90376501)
src = jax.random.randint(genkey, [3, 4, 5], -32, 32, jax.numpy.int32)
print(jax.jit(fn)(src))
```

```
[[-17 -28   3  16 -31]
 [ 14 -28 -21 -18  17]]
```

but I think the indices here is not valid, `31` is obviously out of range, and jax just uses `31 % 4` as the index to take the elements. Maybe it needs a index check.

---

jax version: 0.3.17


Discussion:
Thanks - we should be able to fix this.

Resolving patch:
diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -14,7 +14,7 @@
 
 import enum
 from functools import partial
-from typing import Any, Callable, NamedTuple, Optional, Sequence, Union
+from typing import Any, Callable, NamedTuple, Optional, Sequence, Tuple, Union
 import weakref
 
 import numpy as np
@@ -173,9 +173,9 @@ class GatherDimensionNumbers(NamedTuple):
   implicit; there is always an index vector dimension and it must always be the
   last dimension. To gather scalar indices, add a trailing dimension of size 1.
   """
-  offset_dims: Sequence[int]
-  collapsed_slice_dims: Sequence[int]
-  start_index_map: Sequence[int]
+  offset_dims: Tuple[int, ...]
+  collapsed_slice_dims: Tuple[int, ...]
+  start_index_map: Tuple[int, ...]
 
 
 class GatherScatterMode(enum.Enum):
@@ -612,15 +612,17 @@ def scatter(
 
 def index_take(src: Array, idxs: Array, axes: Sequence[int]) -> Array:
   indices = lax.concatenate([lax.expand_dims(i, (1,)) for i in idxs], 1)
-  indices = indices % np.array([src.shape[ax] for ax in axes])
+  max_idx = lax.expand_dims(np.array([src.shape[ax] for ax in axes]),
+                            tuple(range(indices.ndim - 1)))
+  indices = indices % max_idx
   slice_sizes = list(src.shape)
   for ax in axes:
     slice_sizes[ax] = 1
   offset_dims = tuple(range(1, src.ndim - indices.shape[1] + 1))
   dnums = GatherDimensionNumbers(
       offset_dims=offset_dims,
-      collapsed_slice_dims=axes,
-      start_index_map=axes)
+      collapsed_slice_dims=tuple(axes),
+      start_index_map=tuple(axes))
   return gather(src, indices, dimension_numbers=dnums,
                 slice_sizes=tuple(slice_sizes))
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2287,9 +2287,8 @@ def collapse_first_two(x):
           [(3, 4, 5), (np.array([0, 2, 1]),), (0,)],
           [(3, 4, 5), (np.array([-1, -2]),), (0,)],
           [(3, 4, 5), (np.array([0, 2]), np.array([1, 3])), (0, 1)],
-          [(3, 4, 5), (np.array([0, 2]), np.array([1, 3])), (0, 2)],
+          [(3, 4, 5), (np.array([0, 2]), np.array([1, 3])), [0, 2]],
       ]))
-  @jax.numpy_rank_promotion('allow')  # Test explicitly exercises implicit rank promotion.
   def testIndexTake(self, shape, dtype, idxs, axes):
     rng = jtu.rand_default(self.rng())
     rand_idxs = lambda: tuple(rng(e.shape, e.dtype) for e in idxs)