Always check types in one-hot

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -394,6 +394,13 @@ def one_hot(
 ) -> Array:
     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)
     if is_jax_namespace(xp):
         assert is_jax_array(x)
         from jax.nn import one_hot
@@ -412,10 +419,6 @@ def one_hot(
             dtype = xp.empty(()).dtype  # Default float dtype
         out = xp.zeros((x.size, num_classes), dtype=dtype)
         x_flattened = xp.reshape(x, (-1,))
-        x_size = x.size
-        if x_size is None:
-            msg = "x must have a concrete size."
-            raise TypeError(msg)
         if is_numpy_namespace(xp):
             at(out)[xp.arange(x_size), x_flattened].set(1)
         else:

tests/test_funcs.py

@@ -455,7 +455,7 @@ def test_xp(self, xp: ModuleType):
 )
 class TestOneHot:
     @pytest.mark.parametrize("n_dim", range(4))
-    @pytest.mark.parametrize("num_classes", range(1, 5, 2))
+    @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)
@@ -508,7 +508,7 @@ def test_axis(self, xp: ModuleType):
         xp_assert_equal(actual, expected)
 
     def test_non_integer(self, xp: ModuleType):
-        with pytest.raises((TypeError, RuntimeError, IndexError, DeprecationWarning)):
+        with pytest.raises(TypeError):
             one_hot(xp.asarray([1.0]), 3)