Add one_hot

Author: NeilGirdhar

docs/api-reference.md

@@ -15,6 +15,7 @@
     expand_dims
     isclose
     kron
+    one_hot
     nunique
     pad
     setdiff1d

src/array_api_extra/__init__.py

@@ -10,6 +10,7 @@
     create_diagonal,
     expand_dims,
     kron,
+    one_hot,
     nunique,
     setdiff1d,
     sinc,
@@ -31,6 +32,7 @@
     "isclose",
     "kron",
     "lazy_apply",
+    "one_hot",
     "nunique",
     "pad",
     "setdiff1d",

src/array_api_extra/_lib/_funcs.py

@@ -8,15 +8,20 @@
 
 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,
+    is_torch_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 +380,39 @@ def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     return xp.squeeze(c, axis=axes)
 
 
+def one_hot(
+    x: Array,
+    /,
+    num_classes: int,
+    *,
+    dtype: DType | None = None,
+    axis: int = -1,
+    xp: ModuleType | None = None,
+) -> Array:
+    if xp is None:
+        xp = array_namespace(x)
+    if 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_array(x):
+        from torch.nn.functional import one_hot
+        out = one_hot(x, num_classes)
+        if dtype is not None:
+            out = xp.astype(out, dtype)
+    else:
+        if dtype is None:
+            dtype = xp.float64
+        out = xp.zeros((x.size, num_classes), dtype=dtype)
+        at(out)[xp.arange(x.size), xp.reshape(x, -1)].set(1)
+    if x.ndim != 1:
+        out = xp.reshape(out, (*x.shape, num_classes))
+    if axis != -1:
+        out = xp.moveaxis(out, -1, axis)
+    return out
+
+
 def create_diagonal(
     x: Array, /, *, offset: int = 0, xp: ModuleType | None = None
 ) -> Array:

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,24 @@ def test_xp(self, xp: ModuleType):
         )
 
 
+@pytest.mark.skip_xp_backend(
+    Backend.SPARSE, reason="read-only backend without .at support"
+)
+class TestOneHot:
+    @pytest.mark.parametrize("n_dim", range(4))
+    @pytest.mark.parametrize("num_classes", range(1, 5, 2))
+    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 y[*i, j] == (x[i] == j)
+
+
 @pytest.mark.skip_xp_backend(
     Backend.SPARSE, reason="read-only backend without .at support"
 )