Allow boolean indexing

Author: ricardoV94

pytensor/xtensor/indexing.py

@@ -44,20 +44,27 @@ def as_idx_variable(idx):
             dims = (dim,)
         else:
             dims = tuple(dim)
-        idx = xtensor_from_tensor(as_tensor(idx), dims=dims)
+        idx = as_xtensor(as_tensor(idx), dims=dims)
     else:
         # Must be integer indices, we already counted for None and slices
         try:
             idx = as_xtensor(idx)
         except TypeError:
             idx = as_tensor(idx)
         if idx.type.dtype == "bool":
-            raise NotImplementedError("Boolean indexing not yet supported")
-        if idx.type.dtype not in discrete_dtypes:
+            if idx.type.ndim != 1:
+                # xarray allaws `x[True]`, but I think it is a bug: https://github.com/pydata/xarray/issues/10379
+                # Otherwise, it is always restricted to 1d boolean indexing arrays
+                raise NotImplementedError(
+                    "Only 1d boolean indexing arrays are supported"
+                )
+            # Convert to nonzero indices
+            if isinstance(idx.type, XTensorType):
+                idx = as_xtensor(idx.values.nonzero()[0], dims=idx.type.dims)
+            else:
+                idx = idx.nonzero()[0]
+        elif idx.type.dtype not in discrete_dtypes:
             raise TypeError("Numerical indices must be integers or boolean")
-        if idx.type.dtype == "bool" and idx.type.ndim == 0:
-            # This can't be triggered right now, but will once we lift the boolean restriction
-            raise NotImplementedError("Scalar boolean indices not supported")
     return idx
 
 

tests/xtensor/test_indexing.py

@@ -295,3 +295,39 @@ def test_scalar_integer_indexing(dims_order):
     expected_res2 = x_test[tuple(idxs)]
     xr_assert_allclose(res1, expected_res1)
     xr_assert_allclose(res2, expected_res2)
+
+
+def test_unsupported_boolean_indexing():
+    x = xtensor(dims=("a", "b"), shape=(3, 5))
+
+    mat_idx = xtensor("idx", dtype=bool, shape=(4, 2), dims=("a", "b"))
+    scalar_idx = mat_idx.isel(a=0, b=1)
+
+    for idx in (mat_idx, mat_idx.values, scalar_idx, scalar_idx.values):
+        with pytest.raises(
+            NotImplementedError,
+            match="Only 1d boolean indexing arrays are supported",
+        ):
+            x[idx]
+
+
+def test_boolean_indexing():
+    x = xtensor("x", shape=(8, 7), dims=("a", "b"))
+    bool_idx = xtensor("bool_idx", dtype=bool, shape=(8,), dims=("a",))
+    int_idx = xtensor("int_idx", dtype=int, shape=(4, 3), dims=("a", "new_dim"))
+
+    out_vectorized = x[bool_idx, int_idx]
+    out_orthogonal = x[bool_idx, int_idx.rename(a="b")]
+    fn = xr_function([x, bool_idx, int_idx], [out_vectorized, out_orthogonal])
+
+    x_test = xr_arange_like(x)
+    bool_idx_test = DataArray(np.array([True, False] * 4, dtype=bool), dims=("a",))
+    int_idx_test = DataArray(
+        np.random.binomial(n=4, p=0.5, size=(4, 3)),
+        dims=("a", "new_dim"),
+    )
+    res1, res2 = fn(x_test, bool_idx_test, int_idx_test)
+    expected_res1 = x_test[bool_idx_test, int_idx_test]
+    expected_res2 = x_test[bool_idx_test, int_idx_test.rename(a="b")]
+    xr_assert_allclose(res1, expected_res1)
+    xr_assert_allclose(res2, expected_res2)