Add docstring

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -392,6 +392,41 @@ def one_hot(
     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)
+    """
     if xp is None:
         xp = array_namespace(x)
     x_size = x.size