Comment

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -392,12 +392,11 @@ def one_hot(
         # specification.
         msg = "x must have a concrete size."
         raise TypeError(msg)
-    # TODO: Benchmark whether this is faster on the numpy backend:
+    # 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 = at(out)[xp.arange(x_size), x_flattened].set(1)
+    # out = xp.reshape(out, (*x.shape, num_classes))
     out = x[..., None] == xp.arange(
         num_classes, dtype=x.dtype, device=_compat.device(x)
     )