Break

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -397,9 +397,8 @@ def one_hot(
     # out = xp.zeros((x.size, num_classes), dtype=dtype, device=_compat.device(x))
     # 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)
-    )
+    range_num_classes = xp.arange(num_classes, dtype=x.dtype, device=_compat.device(x))
+    out = x[..., xp.newaxis] == range_num_classes
     return xp.astype(out, dtype)