Fix PyTorch errors and test int types

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -410,7 +410,11 @@ def one_hot(
     if is_torch_namespace(xp):
         assert is_torch_array(x)
         from torch.nn.functional import one_hot
-        out = one_hot(x, num_classes)
+        x = xp.astype(x, xp.int64)  # PyTorch only supports int64 here.
+        try:
+            out = one_hot(x, num_classes)
+        except RuntimeError as e:
+            raise IndexError from e
         if dtype is None:
             dtype = xp.float
         out = xp.astype(out, dtype)

tests/test_funcs.py

@@ -480,18 +480,32 @@ def test_basic(self, xp: ModuleType):
                                [1., 0., 0.]])
         xp_assert_equal(actual, expected)
 
+    @pytest.mark.skip_xp_backend(
+        Backend.TORCH_GPU, reason="Puts Pytorch into a bad state."
+    )
     def test_out_of_bound(self, xp: ModuleType):
         # Undefined behavior.  Either return zero, or raise.
         try:
             actual = one_hot(xp.asarray([-1, 3]), 3)
-        except (IndexError, RuntimeError):
+        except IndexError:
             return
         expected = xp.asarray([[0., 0., 0.],
                                [0., 0., 0.]])
         xp_assert_equal(actual, expected)
 
+    @pytest.mark.parametrize("int_dtype", ['int8', 'int16', 'int32', 'int64', 'uint8',
+                                           'uint16', 'uint32', 'uint64'])
+    def test_int_types(self, xp: ModuleType, int_dtype: str):
+        dtype = getattr(xp, int_dtype)
+        x = xp.asarray([0, 1, 2], dtype=dtype)
+        actual = one_hot(x, 3)
+        expected = xp.asarray([[1., 0., 0.],
+                               [0., 1., 0.],
+                               [0., 0., 1.]])
+        xp_assert_equal(actual, expected)
+
     def test_custom_dtype(self, xp: ModuleType):
-        actual = one_hot(xp.asarray([0, 1, 2]), 3, dtype=xp.bool)
+        actual = one_hot(xp.asarray([0, 1, 2], dtype=xp.int32), 3, dtype=xp.bool)
         expected = xp.asarray([[True, False, False],
                                [False, True, False],
                                [False, False, True]])