Repository: google/jax
Issue: Boolean index with ellipsis

Boolean indexing the 'last' index of a 1D array using the ellipsis notation does currently no longer work in the version of JAX I'm using (jax==0.2.26, jaxlib==0.1.75). This is not the intended behavior as the ellipsis should generalize to 0 extra dimensions in front as well (check by for example swapping `jnp` with `np` in the example below).

Here's a minimal example:

```python
arr = jnp.array([1, 2, 3])
mask = jnp.array([True, False, True])
arr[..., mask]
```

This is the full traceback:

```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
/tmp/ipykernel_36135/3397629872.py in <module>
      1 arr = jnp.array([1, 2, 3])
      2 mask = jnp.array([True, False, True])
----> 3 arr[..., mask]

~/.anaconda/envs/sax/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in _rewriting_take(arr, idx, indices_are_sorted, unique_indices, mode, fill_value)
   5701   # followed by an optional reverse and broadcast_in_dim.
   5702   arr = asarray(arr)
-> 5703   treedef, static_idx, dynamic_idx = _split_index_for_jit(idx, arr.shape)
   5704   return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,
   5705                  unique_indices, mode, fill_value)

~/.anaconda/envs/sax/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in _split_index_for_jit(idx, shape)
   5780   # Expand any (concrete) boolean indices. We can then use advanced integer
   5781   # indexing logic to handle them.
-> 5782   idx = _expand_bool_indices(idx, shape)
   5783 
   5784   leaves, treedef = tree_flatten(idx)

~/.anaconda/envs/sax/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in _expand_bool_indices(idx, shape)
   6092         expected_shape = shape[start: start + _ndim(i)]
   6093         if i_shape != expected_shape:
-> 6094           raise IndexError("boolean index did not match shape of indexed array in index "
   6095                            f"{dim_number}: got {i_shape}, expected {expected_shape}")
   6096         out.extend(np.where(i))

IndexError: boolean index did not match shape of indexed array in index 1: got (3,), expected ()
```


Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_indexing_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 73d92dba5..ea67fcbec 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -6090,7 +6090,7 @@ def _expand_bool_indices(idx, shape):
     else:
       out.append(i)
     if i is Ellipsis:
-      ellipsis_offset = _max(0, len(shape) - total_dims - 1)
+      ellipsis_offset = len(shape) - total_dims - 1
   return tuple(out)
 
 
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index 26272e26e..d8bb9dbf7 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -757,6 +757,21 @@ class IndexingTest(jtu.JaxTestCase):
     expected = x[idx]
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testBoolean1DIndexingWithEllipsis2(self):
+    # Regression test for https://github.com/google/jax/issues/9050
+    x = np.arange(3)
+    idx = (..., np.array([True, False, True]))
+    ans = jnp.array(x)[idx]
+    expected = x[idx]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testBoolean1DIndexingWithEllipsis3(self):
+    x = np.arange(6).reshape(2, 3)
+    idx = (0, ..., np.array([True, False, True]))
+    ans = jnp.array(x)[idx]
+    expected = x[idx]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testBoolean2DIndexingWithEllipsis(self):
     x = np.arange(24).reshape(4, 3, 2)
     idx = (..., np.array([[True, False], [True, False], [False, False]]))