crusadersky's idea

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -16,6 +16,7 @@
     is_torch_array,
     is_torch_namespace,
 )
+from ._lib._utils._compat import device as get_device
 from ._lib._utils._helpers import asarrays
 from ._lib._utils._typing import Array, DType
 
@@ -169,7 +170,9 @@ def one_hot(
         msg = "x must have an integral dtype."
         raise TypeError(msg)
     if dtype is None:
-        dtype = xp.__array_namespace_info__().default_dtypes(device=get_device(x))["real floating"]
+        dtype = xp.__array_namespace_info__().default_dtypes(device=get_device(x))[
+            "real floating"
+        ]
     # Delegate where possible.
     if is_jax_namespace(xp):
         assert is_jax_array(x)
@@ -192,8 +195,6 @@ def one_hot(
             num_classes,
             dtype=dtype,
             xp=xp,
-            supports_fancy_indexing=is_numpy_namespace(xp),
-            supports_array_indexing=is_dask_namespace(xp),
         )
 
     if axis != -1:

src/array_api_extra/_lib/_funcs.py

@@ -380,8 +380,6 @@ def one_hot(
     /,
     num_classes: int,
     *,
-    supports_fancy_indexing: bool = False,
-    supports_array_indexing: bool = False,
     dtype: DType,
     xp: ModuleType,
 ) -> Array:  # numpydoc ignore=PR01,RT01
@@ -394,19 +392,16 @@ def one_hot(
         # specification.
         msg = "x must have a concrete size."
         raise TypeError(msg)
-    out = xp.zeros((x.size, num_classes), dtype=dtype, device=_compat.device(x))
-    x_flattened = xp.reshape(x, (-1,))
-    if supports_fancy_indexing:
-        out = at(out)[xp.arange(x_size), x_flattened].set(1)
-    else:
-        for i in range(x_size):
-            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
+    # TODO: Benchmark whether this is faster on the numpy backend:
+    # x_flattened = xp.reshape(x, (-1,))
+    # out = xp.zeros((x.size, num_classes), dtype=dtype, device=_compat.device(x))
+    # at(out)[xp.arange(x_size), x_flattened].set(1)
+    # if x.ndim != 1:
+    #     out = xp.reshape(out, (*x.shape, num_classes))
+    out = x[..., None] == xp.arange(
+        num_classes, dtype=x.dtype, device=_compat.device(x)
+    )
+    return xp.astype(out, dtype)
 
 
 def create_diagonal(