Use iterative approach when non-numpy

NeilGirdhar · NeilGirdhar · commit 9b5f393d3293 · 2025-05-30T01:45:40.000-04:00
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
@@ -12,7 +12,10 @@
     array_namespace,
     is_dask_namespace,
     is_jax_array,
+    is_jax_namespace,
+    is_numpy_namespace,
     is_torch_array,
+    is_torch_namespace,
 )
 from ._utils._helpers import (
     asarrays,
@@ -391,22 +394,32 @@ def one_hot(
 ) -> Array:
     if xp is None:
         xp = array_namespace(x)
-    if is_jax_array(x):
+    if is_jax_namespace(xp):
+        assert 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):
+    if is_torch_namespace(xp):
+        assert is_torch_array(x)
         from torch.nn.functional import one_hot
         out = one_hot(x, num_classes)
         if dtype is None:
             dtype = xp.float
         out = xp.astype(out, dtype)
     else:
         if dtype is None:
-            dtype = xp.float64
+            dtype = xp.empty(()).dtype  # Default float dtype
         out = xp.zeros((x.size, num_classes), dtype=dtype)
-        at(out)[xp.arange(x.size), xp.reshape(x, (-1,))].set(1)
+        x_flattened = xp.reshape(x, (-1,))
+        x_size = x.size
+        if x_size is None:
+            raise TypeError
+        if is_numpy_namespace(xp):
+            at(out)[xp.arange(x.size), xp.reshape(x, (-1,))].set(1)
+        else:
+            for i in range(x_size):
+                at(out)[i, int(x_flattened[i])].set(1)
     if x.ndim != 1:
         out = xp.reshape(out, (*x.shape, num_classes))
     if axis != -1:
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
@@ -465,7 +465,7 @@ def test_dims_and_classes(self, xp: ModuleType, n_dim: int, num_classes: int):
         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)
+            assert float(y[*i, j]) == (int(x[i]) == j)
 
     def test_basic(self, xp: ModuleType):
         actual = one_hot(xp.asarray([0, 1, 2]), 3)
