Repository: google/jax
Issue: Boolean indexing does not compatible with Ellipsis

It seems that the following snippet does not work anymore.

```python
import jax
import jax.numpy as jnp
import numpy as np

jnp.ones((4, 3, 2))[..., np.array([True, False])]  # fail
jnp.ones((4, 3, 2))[:, :, np.array([True, False])]  # work
np.ones((4, 3, 2))[..., np.array([True, False])]  # work
```

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 fe5f900a0..6f00fff42 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -5996,16 +5996,32 @@ def _eliminate_deprecated_list_indexing(idx):
       idx = (idx,)
   return idx
 
+def _is_boolean_index(i):
+  try:
+    abstract_i = core.get_aval(i)
+  except TypeError:
+    abstract_i = None
+  return (isinstance(abstract_i, ShapedArray) and issubdtype(abstract_i.dtype, bool_)
+          or isinstance(i, list) and i and _all(_is_scalar(e)
+          and issubdtype(_dtype(e), np.bool_) for e in i))
+
 def _expand_bool_indices(idx, shape):
   """Converts concrete bool indexes into advanced integer indexes."""
   out = []
+  total_dims = len(shape)
+  num_ellipsis = _sum(e is Ellipsis for e in idx)
+  if num_ellipsis > 1:
+    raise IndexError("an index can only have a single ellipsis ('...')")
+  elif num_ellipsis == 1:
+    total_dims = _sum(_ndim(e) if _is_boolean_index(e) else 1 for e in idx
+                      if e is not None and e is not Ellipsis)
+  ellipsis_offset = 0
   for dim_number, i in enumerate(idx):
     try:
       abstract_i = core.get_aval(i)
     except TypeError:
       abstract_i = None
-    if (isinstance(abstract_i, ShapedArray) and issubdtype(abstract_i.dtype, bool_)
-        or isinstance(i, list) and i and _all(_is_scalar(e) and issubdtype(_dtype(e), np.bool_) for e in i)):
+    if _is_boolean_index(i):
       if isinstance(i, list):
         i = array(i)
         abstract_i = core.get_aval(i)
@@ -6017,13 +6033,16 @@ def _expand_bool_indices(idx, shape):
         raise TypeError("JAX arrays do not support boolean scalar indices")
       else:
         i_shape = _shape(i)
-        expected_shape = shape[len(out): len(out) + _ndim(i)]
+        start = len(out) + ellipsis_offset
+        expected_shape = shape[start: start + _ndim(i)]
         if i_shape != expected_shape:
           raise IndexError("boolean index did not match shape of indexed array in index "
                            f"{dim_number}: got {i_shape}, expected {expected_shape}")
         out.extend(np.where(i))
     else:
       out.append(i)
+    if i is Ellipsis:
+      ellipsis_offset = _max(0, 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 1a62300e4..68b6e99e0 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -749,6 +749,28 @@ class IndexingTest(jtu.JaxTestCase):
     expected = x[idx]
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testBoolean1DIndexingWithEllipsis(self):
+    # Regression test for https://github.com/google/jax/issues/8412
+    x = np.arange(24).reshape(4, 3, 2)
+    idx = (..., np.array([True, False]))
+    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]]))
+    ans = jnp.array(x)[idx]
+    expected = x[idx]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testBoolean1DIndexingWithTrailingEllipsis(self):
+    x = np.arange(24).reshape(4, 3, 2)
+    idx = (np.array([True, False, True, False]), ...)
+    ans = jnp.array(x)[idx]
+    expected = x[idx]
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testBooleanIndexingDynamicShapeError(self):
     x = np.zeros(3)
     i = np.array([True, True, False])
@@ -760,6 +782,8 @@ class IndexingTest(jtu.JaxTestCase):
       jnp.arange(4)[True]
     with self.assertRaisesRegex(TypeError, msg):
       jnp.arange(4)[False]
+    with self.assertRaisesRegex(TypeError, msg):
+      jnp.arange(4)[..., True]
 
   def testIssue187(self):
     x = jnp.ones((5, 5))