Skip Dask, and use flags

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -188,7 +188,15 @@ def one_hot(
             raise IndexError from e
         out = xp.astype(out, dtype)
     else:
-        out = _funcs.one_hot(x, num_classes, x_size=x_size, dtype=dtype, xp=xp)
+        out = _funcs.one_hot(
+            x,
+            num_classes,
+            x_size=x_size,
+            dtype=dtype,
+            xp=xp,
+            supports_fancy_indexing=is_numpy_namespace(xp),
+            supports_array_indexing=is_dask_namespace(xp),
+        )
 
     if axis != -1:
         out = xp.moveaxis(out, -1, axis)

src/array_api_extra/_lib/_funcs.py

@@ -12,7 +12,6 @@
     array_namespace,
     is_dask_namespace,
     is_jax_array,
-    is_numpy_namespace,
 )
 from ._utils._helpers import (
     asarrays,
@@ -385,17 +384,22 @@ def one_hot(
     /,
     num_classes: int,
     *,
+    supports_fancy_indexing: bool = False,
+    supports_array_indexing: bool = False,
     x_size: int,
     dtype: DType,
     xp: ModuleType,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._delegation.py`."""
     out = xp.zeros((x.size, num_classes), dtype=dtype)
     x_flattened = xp.reshape(x, (-1,))
-    if is_numpy_namespace(xp):
+    if supports_fancy_indexing:
         out = at(out)[xp.arange(x_size), x_flattened].set(1)
     for i in range(x_size):
-        out = at(out)[i, int(x_flattened[i])].set(1)
+        x_i = x_flattened[i]
+        if not supports_array_indexing:
+            x_i = int(x_i)
+        out = at(out)[i, x_i].set(1)
     if x.ndim != 1:
         out = xp.reshape(out, (*x.shape, num_classes))
     return out

tests/test_funcs.py

@@ -453,6 +453,9 @@ def test_xp(self, xp: ModuleType):
 @pytest.mark.skip_xp_backend(
     Backend.SPARSE, reason="read-only backend without .at support"
 )
+@pytest.mark.skip_xp_backend(
+    Backend.DASK, reason="backend does not yet support indexed assignment"
+)
 class TestOneHot:
     @pytest.mark.parametrize("n_dim", range(4))
     @pytest.mark.parametrize("num_classes", [1, 3, 10])