Add one_hot

NeilGirdhar · NeilGirdhar · commit ecc8b40dd02e · 2025-05-30T23:46:19.000-04:00
diff --git a/docs/api-reference.md b/docs/api-reference.md
@@ -15,6 +15,7 @@
     expand_dims
     isclose
     kron
+    one_hot
     nunique
     pad
     setdiff1d
diff --git a/src/array_api_extra/__init__.py b/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,
@@ -32,6 +32,7 @@
     "kron",
     "lazy_apply",
     "nunique",
+    "one_hot",
     "pad",
     "setdiff1d",
     "sinc",
diff --git a/src/array_api_extra/_delegation.py b/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,90 @@ 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.
+    dtype : DType, optional
+        The dtype of the return value.  Defaults to the default float dtype (usually
+        float64).
+    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`.
+
+        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)
+    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 as jax_one_hot
+
+        return jax_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 as torch_one_hot
+
+        x = xp.astype(x, xp.int64)  # PyTorch only supports int64 here.
+        try:
+            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,
+            supports_fancy_indexing=is_numpy_namespace(xp),
+            supports_array_indexing=is_dask_namespace(xp),
+        )
+
+    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]],
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
@@ -8,15 +8,19 @@
 
 from ._at import at
 from ._utils import _compat, _helpers
-from ._utils._compat import array_namespace, is_dask_namespace, is_jax_array
+from ._utils._compat import (
+    array_namespace,
+    is_dask_namespace,
+    is_jax_array,
+)
 from ._utils._helpers import (
     asarrays,
     capabilities,
     eager_shape,
     meta_namespace,
     ndindex,
 )
-from ._utils._typing import Array
+from ._utils._typing import Array, DType
 
 __all__ = [
     "apply_where",
@@ -375,6 +379,36 @@ def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     return xp.squeeze(c, axis=axes)
 
 
+def one_hot(
+    x: Array,
+    /,
+    num_classes: int,
+    *,
+    supports_fancy_indexing: bool = False,
+    supports_array_indexing: bool = False,
+    dtype: DType,
+    xp: ModuleType,
+) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in `array_api_extra._delegation.py`."""
+    x_size = x.size
+    if x_size is None:  # pragma: no cover
+        msg = "x must have a concrete size."
+        raise TypeError(msg)
+    out = xp.zeros((x.size, num_classes), dtype=dtype)
+    x_flattened = xp.reshape(x, (-1,))
+    if supports_fancy_indexing:
+        out = at(out)[xp.arange(x_size), x_flattened].set(1)
+    else:
+        for i in range(x_size):
+            x_i = x_flattened[i]
+            if not supports_array_indexing:
+                x_i = int(x_i)
+            out = at(out)[i, x_i].set(1)
+    if x.ndim != 1:
+        out = xp.reshape(out, (*x.shape, num_classes))
+    return out
+
+
 def create_diagonal(
     x: Array, /, *, offset: int = 0, xp: ModuleType | None = None
 ) -> Array:
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
@@ -21,6 +21,7 @@
     isclose,
     kron,
     nunique,
+    one_hot,
     pad,
     setdiff1d,
     sinc,
@@ -44,6 +45,7 @@
 lazy_xp_function(expand_dims)
 lazy_xp_function(kron)
 lazy_xp_function(nunique)
+lazy_xp_function(one_hot)
 lazy_xp_function(pad)
 # FIXME calls in1d which calls xp.unique_values without size
 lazy_xp_function(setdiff1d, jax_jit=False)
@@ -448,6 +450,78 @@ def test_xp(self, xp: ModuleType):
         )
 
 
+@pytest.mark.skip_xp_backend(
+    Backend.SPARSE, reason="read-only backend without .at support"
+)
+@pytest.mark.skip_xp_backend(
+    Backend.DASK, reason="backend does not yet support indexed assignment"
+)
+class TestOneHot:
+    @pytest.mark.parametrize("n_dim", range(4))
+    @pytest.mark.parametrize("num_classes", [1, 3, 10])
+    def test_dims_and_classes(self, xp: ModuleType, n_dim: int, num_classes: int):
+        shape = tuple(range(2, 2 + n_dim))
+        rng = np.random.default_rng(2347823)
+        np_x = rng.integers(num_classes, size=shape)
+        x = xp.asarray(np_x)
+        y = one_hot(x, num_classes)
+        assert y.shape == (*x.shape, num_classes)
+        for *i_list, j in ndindex(*shape, num_classes):
+            i = tuple(i_list)
+            assert float(y[(*i, j)]) == (int(x[i]) == j)
+
+    def test_basic(self, xp: ModuleType):
+        actual = one_hot(xp.asarray([0, 1, 2]), 3)
+        expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
+        xp_assert_equal(actual, expected)
+
+        actual = one_hot(xp.asarray([1, 2, 0]), 3)
+        expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
+        xp_assert_equal(actual, expected)
+
+    @pytest.mark.skip_xp_backend(
+        Backend.TORCH_GPU, reason="Puts Pytorch into a bad state."
+    )
+    def test_out_of_bound(self, xp: ModuleType):
+        # Undefined behavior.  Either return zero, or raise.
+        try:
+            actual = one_hot(xp.asarray([-1, 3]), 3)
+        except IndexError:
+            return
+        expected = xp.asarray([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
+        xp_assert_equal(actual, expected)
+
+    @pytest.mark.parametrize(
+        "int_dtype",
+        ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"],
+    )
+    def test_int_types(self, xp: ModuleType, int_dtype: str):
+        dtype = getattr(xp, int_dtype)
+        x = xp.asarray([0, 1, 2], dtype=dtype)
+        actual = one_hot(x, 3)
+        expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
+        xp_assert_equal(actual, expected)
+
+    def test_custom_dtype(self, xp: ModuleType):
+        actual = one_hot(xp.asarray([0, 1, 2], dtype=xp.int32), 3, dtype=xp.bool)
+        expected = xp.asarray(
+            [[True, False, False], [False, True, False], [False, False, True]]
+        )
+        xp_assert_equal(actual, expected)
+
+    def test_axis(self, xp: ModuleType):
+        expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]).T
+        actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=0)
+        xp_assert_equal(actual, expected)
+
+        actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=-2)
+        xp_assert_equal(actual, expected)
+
+    def test_non_integer(self, xp: ModuleType):
+        with pytest.raises(TypeError):
+            _ = one_hot(xp.asarray([1.0]), 3)
+
+
 @pytest.mark.skip_xp_backend(
     Backend.SPARSE, reason="read-only backend without .at support"
 )
