Use delegation file

Author: NeilGirdhar

src/array_api_extra/__init__.py

@@ -1,6 +1,6 @@
 """Extra array functions built on top of the array API standard."""
 
-from ._delegation import isclose, pad
+from ._delegation import isclose, one_hot, pad
 from ._lib._at import at
 from ._lib._funcs import (
     apply_where,
@@ -11,7 +11,6 @@
     expand_dims,
     kron,
     nunique,
-    one_hot,
     setdiff1d,
     sinc,
 )

src/array_api_extra/_delegation.py

@@ -9,15 +9,17 @@
     array_namespace,
     is_cupy_namespace,
     is_dask_namespace,
+    is_jax_array,
     is_jax_namespace,
     is_numpy_namespace,
     is_pydata_sparse_namespace,
+    is_torch_array,
     is_torch_namespace,
 )
 from ._lib._utils._helpers import asarrays
-from ._lib._utils._typing import Array
+from ._lib._utils._typing import Array, DType
 
-__all__ = ["isclose", "pad"]
+__all__ = ["isclose", "one_hot", "pad"]
 
 
 def isclose(
@@ -112,6 +114,88 @@ def isclose(
     return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp)
 
 
+def one_hot(
+    x: Array,
+    /,
+    num_classes: int,
+    *,
+    dtype: DType | None = None,
+    axis: int = -1,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    One-hot encode the given indices.
+
+    Each index in the input ``x`` is encoded as a vector of zeros of length
+    ``num_classes`` with the element at the given index set to one.
+
+    Parameters
+    ----------
+    x : array
+        An array with integral dtype having shape ``batch_dims``.
+    num_classes : int
+        Number of classes in the one-hot dimension.
+    axis : int or tuple of ints, optional
+        Position(s) in the expanded axes where the new axis is placed.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array having the same shape as `x` except for a new axis at the position
+        given by `axis` having size `num_classes`.
+
+        The dtype of the return value is the default float dtype (usually float64).
+
+        If ``x < 0`` or ``x >= num_classes``, then the result is undefined, may raise
+        an exception, or may even cause a bad state.  `x` is not checked.
+
+    Examples
+    --------
+    >>> xp.one_hot(jnp.array([1, 2, 0]), 3)
+    Array([[0., 1., 0.],
+           [0., 0., 1.],
+           [1., 0., 0.]], dtype=float64)
+    """
+    # Validate inputs.
+    if xp is None:
+        xp = array_namespace(x)
+    x_size = x.size
+    if x_size is None:
+        msg = "x must have a concrete size."
+        raise TypeError(msg)
+    if not xp.isdtype(x.dtype, "integral"):
+        msg = "x must have an integral dtype."
+        raise TypeError(msg)
+    if dtype is None:
+        dtype = xp.empty(()).dtype  # Default float dtype
+    # Delegate where possible.
+    if is_jax_namespace(xp):
+        assert is_jax_array(x)
+        from jax.nn import one_hot
+
+        return one_hot(x, num_classes, dtype=dtype, axis=axis)
+    if is_torch_namespace(xp):
+        assert is_torch_array(x)
+        from torch.nn.functional import one_hot
+
+        x = xp.astype(x, xp.int64)  # PyTorch only supports int64 here.
+        try:
+            out = 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, 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
+
+
 def pad(
     x: Array,
     pad_width: int | tuple[int, int] | Sequence[tuple[int, int]],

src/array_api_extra/_lib/_funcs.py

@@ -12,10 +12,7 @@
     array_namespace,
     is_dask_namespace,
     is_jax_array,
-    is_jax_namespace,
     is_numpy_namespace,
-    is_torch_array,
-    is_torch_namespace,
 )
 from ._utils._helpers import (
     asarrays,
@@ -388,87 +385,16 @@ def one_hot(
     /,
     num_classes: int,
     *,
-    dtype: DType | None = None,
-    axis: int = -1,
-    xp: ModuleType | None = None,
+    x_size: int,
+    dtype: DType,
+    xp: ModuleType,
 ) -> Array:
-    """
-    One-hot encode the given indices.
-
-    Each index in the input ``x`` is encoded as a vector of zeros of length
-    ``num_classes`` with the element at the given index set to one.
-
-    Parameters
-    ----------
-    x : array
-        An array with integral dtype having shape ``batch_dims``.
-    num_classes : int
-        Number of classes in the one-hot dimension.
-    axis : int or tuple of ints, optional
-        Position(s) in the expanded axes where the new axis is placed.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `x`. Default: infer.
-
-    Returns
-    -------
-    array
-        An array having the same shape as `x` except for a new axis at the position
-        given by `axis` having size `num_classes`.
-
-        The dtype of the return value is the default float dtype (usually float64).
-
-        If ``x < 0`` or ``x >= num_classes``, then the result is undefined, may raise
-        an exception, or may even cause a bad state.  `x` is not checked.
-
-    Examples
-    --------
-    >>> xp.one_hot(jnp.array([1, 2, 0]), 3)
-    Array([[0., 1., 0.],
-           [0., 0., 1.],
-           [1., 0., 0.]], dtype=float64)
-    """
-    if xp is None:
-        xp = array_namespace(x)
-    x_size = x.size
-    if x_size is None:
-        msg = "x must have a concrete size."
-        raise TypeError(msg)
-    if not xp.isdtype(x.dtype, "integral"):
-        msg = "x must have an integral dtype."
-        raise TypeError(msg)
-    if is_jax_namespace(xp):
-        assert is_jax_array(x)
-        from jax.nn import one_hot
-
-        if dtype is None:
-            dtype = xp.float_
-        return one_hot(x, num_classes, dtype=dtype, axis=axis)
-    if is_torch_namespace(xp):
-        assert is_torch_array(x)
-        from torch.nn.functional import one_hot
-
-        x = xp.astype(x, xp.int64)  # PyTorch only supports int64 here.
-        try:
-            out = one_hot(x, num_classes)
-        except RuntimeError as e:
-            raise IndexError from e
-        if dtype is None:
-            dtype = xp.float
-        out = xp.astype(out, dtype)
-    else:
-        if dtype is None:
-            dtype = xp.empty(()).dtype  # Default float dtype
-        out = xp.zeros((x.size, num_classes), dtype=dtype)
-        x_flattened = xp.reshape(x, (-1,))
-        if is_numpy_namespace(xp):
-            out = at(out)[xp.arange(x_size), x_flattened].set(1)
-        else:
-            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))
-    if axis != -1:
-        out = xp.moveaxis(out, -1, axis)
+    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)
+    for i in range(x_size):
+        out = at(out)[i, int(x_flattened[i])].set(1)
     return out