Review response

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -126,13 +126,13 @@ def one_hot(
     """
     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.
+    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``.
+        An array with integral dtype and concrete size (``x.size`` cannot be `None`).
     num_classes : int
         Number of classes in the one-hot dimension.
     dtype : DType, optional
@@ -147,17 +147,20 @@ def one_hot(
     -------
     array
         An array having the same shape as `x` except for a new axis at the position
-        given by `axis` having size `num_classes`.
+        given by `axis` having size `num_classes`.  If `axis` is unspecified, it
+        defaults to -1, which appends a new axis.
 
         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)
+    >>> import array_api_extra as xpx
+    >>> import array-api-strict as xp
+    >>> xpx.one_hot(xp.asarray([1, 2, 0]), 3)
     Array([[0., 1., 0.],
-           [0., 0., 1.],
-           [1., 0., 0.]], dtype=float64)
+          [0., 0., 1.],
+          [1., 0., 0.]], dtype=array_api_strict.float64)
     """
     # Validate inputs.
     if xp is None:

src/array_api_extra/_lib/_funcs.py

@@ -8,11 +8,7 @@
 
 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,
@@ -392,6 +388,10 @@ def one_hot(
     """See docstring in `array_api_extra._delegation.py`."""
     x_size = x.size
     if x_size is None:  # pragma: no cover
+        # This cannot be tested because there is no way to create an array with abstract
+        # size today.  However, it must be blocked for the sake of type-checking and
+        # future-proofing since x.size is allowed to None according to the
+        # specification.
         msg = "x must have a concrete size."
         raise TypeError(msg)
     out = xp.zeros((x.size, num_classes), dtype=dtype)