Formatting

Author: NeilGirdhar

src/array_api_extra/_lib/_funcs.py

@@ -404,12 +404,14 @@ def one_hot(
     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_namespace(xp):
         assert is_torch_array(x)
         from torch.nn.functional import one_hot
+
         x = xp.astype(x, xp.int64)  # PyTorch only supports int64 here.
         try:
             out = one_hot(x, num_classes)

tests/test_funcs.py

@@ -469,15 +469,11 @@ def test_dims_and_classes(self, xp: ModuleType, n_dim: int, num_classes: int):
 
     def test_basic(self, xp: ModuleType):
         actual = one_hot(xp.asarray([0, 1, 2]), 3)
-        expected = xp.asarray([[1., 0., 0.],
-                               [0., 1., 0.],
-                               [0., 0., 1.]])
+        expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
         xp_assert_equal(actual, expected)
 
         actual = one_hot(xp.asarray([1, 2, 0]), 3)
-        expected = xp.asarray([[0., 1., 0.],
-                               [0., 0., 1.],
-                               [1., 0., 0.]])
+        expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])
         xp_assert_equal(actual, expected)
 
     @pytest.mark.skip_xp_backend(
@@ -489,32 +485,29 @@ def test_out_of_bound(self, xp: ModuleType):
             actual = one_hot(xp.asarray([-1, 3]), 3)
         except IndexError:
             return
-        expected = xp.asarray([[0., 0., 0.],
-                               [0., 0., 0.]])
+        expected = xp.asarray([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
         xp_assert_equal(actual, expected)
 
-    @pytest.mark.parametrize("int_dtype", ['int8', 'int16', 'int32', 'int64', 'uint8',
-                                           'uint16', 'uint32', 'uint64'])
+    @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.]])
+        expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
         xp_assert_equal(actual, expected)
 
     def test_custom_dtype(self, xp: ModuleType):
         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]])
+        expected = xp.asarray(
+            [[True, False, False], [False, True, False], [False, False, True]]
+        )
         xp_assert_equal(actual, expected)
 
     def test_axis(self, xp: ModuleType):
-        expected = xp.asarray([[0., 1., 0.],
-                               [0., 0., 1.],
-                               [1., 0., 0.]]).T
+        expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]).T
         actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=0)
         xp_assert_equal(actual, expected)