Simplify

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -188,15 +188,9 @@ def one_hot(
             out = torch_one_hot(x, num_classes)
         except RuntimeError as e:
             raise IndexError from e
-        out = xp.astype(out, dtype)
     else:
-        out = _funcs.one_hot(
-            x,
-            num_classes,
-            dtype=dtype,
-            xp=xp,
-        )
-
+        out = _funcs.one_hot(x, num_classes, xp=xp)
+    out = xp.astype(out, dtype)
     if axis != -1:
         out = xp.moveaxis(out, -1, axis)
     return out

src/array_api_extra/_lib/_funcs.py

@@ -16,7 +16,7 @@
     meta_namespace,
     ndindex,
 )
-from ._utils._typing import Array, DType
+from ._utils._typing import Array
 
 __all__ = [
     "apply_where",
@@ -380,7 +380,6 @@ def one_hot(
     /,
     num_classes: int,
     *,
-    dtype: DType,
     xp: ModuleType,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._delegation.py`."""
@@ -398,8 +397,7 @@ def one_hot(
     # out = at(out)[xp.arange(x_size), x_flattened].set(1)
     # out = xp.reshape(out, (*x.shape, num_classes))
     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)
+    return x[..., xp.newaxis] == range_num_classes
 
 
 def create_diagonal(