Repository: google/jax (python)

Issue:
JAX indexing should support zero-dimensional boolean masks
NumPy supports this:

```python
>>> import numpy as np
>>> x = np.arange(3)
>>> m = np.array([], dtype=bool)
>>> x[m]
array([], dtype=int64)

>>> x = np.arange(12).reshape(3, 4)
>>> m = np.empty((3, 0), dtype=bool)
>>> x[m]
array([], dtype=int64)
```

JAX does not:
```python

>>> import jax.numpy as jnp
>>> x = jnp.arange(3)
>>> m = jnp.array([], dtype=bool)
>>> x[m]
IndexError: boolean index did not match shape of indexed array in index 0: got (0,), expected (3,)

>>> x = jnp.arange(12).reshape(3, 4)
>>> m = jnp.empty((3, 0), dtype=bool)
>>> x[m]
IndexError: boolean index did not match shape of indexed array in index 0: got (3, 0), expected (3, 4)
```
This behavior also appears to be required by the [Array API specification: Boolean Array Indexing](https://data-apis.org/array-api/latest/API_specification/indexing.html#boolean-array-indexing):

> - The size of each dimension in `B` must equal the size of the corresponding dimension in `A` or be `0`, beginning with the first dimension in `A`. If a dimension size does not equal the size of the corresponding dimension in `A` and is not `0`, then an `IndexError` exception must be raised.

Pull request description:
Support empty boolean indexing
Fixes #22886

Patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 3af51e305..7ef27b1ee 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -8340,7 +8340,10 @@ def _expand_bool_indices(idx, shape):
         i_shape = _shape(i)
         start = len(out) + ellipsis_offset - newaxis_offset
         expected_shape = shape[start: start + _ndim(i)]
-        if i_shape != expected_shape:
+        if len(i_shape) != len(expected_shape):
+          raise IndexError(f"too many boolean indices at index {dim_number}: got mask of shape "
+                           f"{i_shape}, but only {len(expected_shape)} dimensions remain.")
+        if not all(s1 in (0, s2) for s1, s2 in zip(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))


Test patch:
diff --git a/tests/lax_numpy_indexing_test.py b/tests/lax_numpy_indexing_test.py
index cbb8e92ed..bf2785f62 100644
--- a/tests/lax_numpy_indexing_test.py
+++ b/tests/lax_numpy_indexing_test.py
@@ -1030,6 +1030,23 @@ class IndexingTest(jtu.JaxTestCase):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @parameterized.parameters(
+      [(3,), (0,)],
+      [(3, 4), (0,)],
+      [(3, 4), (0, 4)],
+      [(3, 4), (3, 0)],
+      [(3, 4, 5), (3, 0)],
+  )
+  def testEmptyBooleanIndexing(self, x_shape, m_shape):
+    # Regression test for https://github.com/google/jax/issues/22886
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng(x_shape, np.int32), np.empty(m_shape, dtype=bool)]
+
+    np_fun = lambda x, m: np.asarray(x)[np.asarray(m)]
+    jnp_fun = lambda x, m: jnp.asarray(x)[jnp.asarray(m)]
+
+    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
+
   @jtu.sample_product(
       shape=[(2, 3, 4, 5)],
       idx=[