Pragma no cover

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -162,10 +162,6 @@ def one_hot(
     # 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)
@@ -191,7 +187,6 @@ def one_hot(
         out = _funcs.one_hot(
             x,
             num_classes,
-            x_size=x_size,
             dtype=dtype,
             xp=xp,
             supports_fancy_indexing=is_numpy_namespace(xp),

src/array_api_extra/_lib/_funcs.py

@@ -386,20 +386,24 @@ def one_hot(
     *,
     supports_fancy_indexing: bool = False,
     supports_array_indexing: bool = False,
-    x_size: int,
     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)
-    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)
+    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