Move reshape

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -190,8 +190,6 @@ def one_hot(
     else:
         out = _funcs.one_hot(x, num_classes, x_size=x_size, dtype=dtype, xp=xp)
 
-    if x.ndim != 1:
-        out = xp.reshape(out, (*x.shape, num_classes))
     if axis != -1:
         out = xp.moveaxis(out, -1, axis)
     return out

src/array_api_extra/_lib/_funcs.py

@@ -393,9 +393,11 @@ def one_hot(
     out = xp.zeros((x.size, num_classes), dtype=dtype)
     x_flattened = xp.reshape(x, (-1,))
     if is_numpy_namespace(xp):
-        return at(out)[xp.arange(x_size), x_flattened].set(1)
+        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)
+    if x.ndim != 1:
+        out = xp.reshape(out, (*x.shape, num_classes))
     return out