Repository: jax-ml/jax
Issue category: Bug Report

Issue:
Difference between numpy and jax.numpy in advanced indexing axes order
### Description

According to #15653, jax should match numpy with respect to advanced indexing.
There is a difference with empty ellipses.


```python
import numpy as np
a=np.ones((3,4,5))
print(a[:,(0,1),...,(0,1)].shape)
# (2, 3) # numpy

import jax.numpy as jnp
a=jnp.ones((3,4,5))
print(a[:,(0,1),...,(0,1)].shape)
# (3, 2) # jax

import torch
a=torch.ones(3,4,5)
print(a[:,(0,1),...,(0,1)].shape)
# torch.Size([3, 2]) # torch

print(np.__version__)
# 1.26.4. same result on 2.1.3

import jax
print(jax.__version__)
# 0.4.33
```



Numpy treats advanced indices separated by an ellipsis that does not consume any axes not as contiguous, whereas jax does  and subsequently does not move the new indices to the front..

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.33
jaxlib: 0.4.33
numpy:  1.26.4
python: 3.10.12 (main, Nov  6 2024, 20:22:13) [GCC 11.4.0]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Linux', node='151e18e640a6', release='6.1.85+', version='#1 SMP PREEMPT_DYNAMIC Thu Jun 27 21:05:47 UTC 2024', machine='x86_64')
```


Hints:
Thanks for the report! I'll take a look.
The issue here is that we determine contiguous advanced indices here: https://github.com/jax-ml/jax/blob/c9a5902216b61bddb88415b9da6bbcf589ef12ea/jax/_src/numpy/lax_numpy.py#L12010-L12019

But we've already removed ellipses a few lines before: https://github.com/jax-ml/jax/blob/c9a5902216b61bddb88415b9da6bbcf589ef12ea/jax/_src/numpy/lax_numpy.py#L11981-L11982

We'll have to rework this logic to account for the fact that unnecessary ellipses do impact contiguousness of advanced indices.

Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 5af8c6ddad19..3d99405428de 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -11971,6 +11971,14 @@ def _int(aval):
 
 def _index_to_gather(x_shape: Sequence[int], idx: Sequence[Any],
                      normalize_indices: bool = True) -> _Indexer:
+  # Check whether advanced indices are contiguous. We must do this before
+  # removing ellipses (https://github.com/jax-ml/jax/issues/25109)
+  # If advanced idexing axes do not appear contiguously, NumPy semantics
+  # move the advanced axes to the front.
+  is_advanced, = np.nonzero([isinstance(e, (int, Sequence, Array, np.ndarray))
+                             or isscalar(e) for e in idx])
+  advanced_axes_are_contiguous = np.all(np.diff(is_advanced) == 1)
+
   # Remove ellipses and add trailing slice(None)s.
   idx = _canonicalize_tuple_index(len(x_shape), idx)
 
@@ -11987,10 +11995,6 @@ def _index_to_gather(x_shape: Sequence[int], idx: Sequence[Any],
   # Check for advanced indexing:
   # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing
 
-  # Do the advanced indexing axes appear contiguously? If not, NumPy semantics
-  # move the advanced axes to the front.
-  advanced_axes_are_contiguous = False
-
   advanced_indexes: Sequence[Array | np.ndarray] | None = None
 
   # The positions of the advanced indexing axes in `idx`.
@@ -12009,7 +12013,6 @@ def _index_to_gather(x_shape: Sequence[int], idx: Sequence[Any],
       advanced_pairs = ((_normalize_index(e, x_shape[j]), i, j)
                         for e, i, j in advanced_pairs)
     advanced_indexes, idx_advanced_axes, x_advanced_axes = zip(*advanced_pairs)
-    advanced_axes_are_contiguous = bool(np.all(np.diff(idx_advanced_axes) == 1))
 
   x_axis = 0  # Current axis in x.
   y_axis = 0  # Current axis in y, before collapsing. See below.


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index 392af2688c1d..ab625d10b4d8 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -399,6 +399,14 @@ def check_grads(f, args, order, atol=None, rtol=None, eps=None):
     IndexSpec(shape=(3, 4), indexer=(Ellipsis, np.array(1, dtype=np.int32)),
               out_shape=(3,)),
   ]),
+  ("EllipsisWithArrayIndices", [
+    IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 1]), ..., np.array([0, 1])),
+              out_shape=(2, 4)),
+    IndexSpec(shape=(3, 4, 5), indexer=(slice(None), np.array([0, 1]), ..., np.array([0, 1])),
+              out_shape=(2, 3)),
+    IndexSpec(shape=(3, 4, 5), indexer=(slice(None), ..., np.array([0, 1]), np.array([0, 1])),
+              out_shape=(3, 2)),
+  ]),
 ]