test correction + pre commit

Signed-off-by: Pradyot Ranjan <99216956+pradyotRanjan@users.noreply.github.com>

Author: prady0t

src/array_api_extra/__init__.py

@@ -19,12 +19,12 @@
 )
 from ._lib._at import at
 from ._lib._funcs import (
+    angle,
     apply_where,
     broadcast_shapes,
     default_dtype,
     kron,
     nunique,
-    angle,
 )
 from ._lib._lazy import lazy_apply
 
@@ -33,6 +33,7 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "angle",
     "apply_where",
     "argpartition",
     "at",
@@ -55,5 +56,4 @@
     "setdiff1d",
     "sinc",
     "union1d",
-    "angle",
 ]

src/array_api_extra/_lib/_funcs.py

@@ -835,7 +835,7 @@ def angle(z: Array, deg: bool = False, /, *, xp: ModuleType | None = None) -> Ar
 
     Returns
     -------
-    angle : ndarray or scalar
+    ndarray or scalar
         The counterclockwise angle from the positive real axis on the complex
         plane in the range ``(-pi, pi]``, with dtype as float64.
     """
@@ -850,4 +850,4 @@ def angle(z: Array, deg: bool = False, /, *, xp: ModuleType | None = None) -> Ar
     a = xp.atan2(zimage, zreal)
     if deg:
         a = a * 180 / xp.pi
-    return a
+    return a

tests/test_funcs.py

@@ -13,6 +13,7 @@
 from typing_extensions import override
 
 from array_api_extra import (
+    angle,
     apply_where,
     argpartition,
     at,
@@ -33,7 +34,6 @@
     setdiff1d,
     sinc,
     union1d,
-    angle,
 )
 from array_api_extra import (
     searchsorted as xpx_searchsorted,
@@ -1883,10 +1883,11 @@ def test_device(self, xp: ModuleType, device: Device):
         b = xp.asarray([2, -2, 0], device=device)
         assert get_device(union1d(a, b)) == device
 
+
 class TestAngle:
     def test_simple(self, xp: ModuleType):
         a = xp.asarray([1, 0])
-        expected = xp.asarray([0., 0.])
+        expected = xp.asarray([0.0, 0.0], dtype=xp.float64)
         res = angle(a)
         xp_assert_equal(res, expected)
 
@@ -1898,11 +1899,13 @@ def test_complex(self, xp: ModuleType):
 
     def test_2d(self, xp: ModuleType):
         a = xp.asarray([[1 + 1j, 1 - 1j], [-1 + 1j, -1 - 1j]])
-        expected = xp.asarray([[np.pi / 4, -np.pi / 4], [3 * np.pi / 4, -3 * np.pi / 4]])
+        expected = xp.asarray(
+            [[np.pi / 4, -np.pi / 4], [3 * np.pi / 4, -3 * np.pi / 4]]
+        )
         res = angle(a)
         xp_assert_equal(res, expected)
 
     @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device")
     def test_device(self, xp: ModuleType, device: Device):
         a = xp.asarray([1 + 1j], device=device)
-        assert get_device(angle(a)) == device
+        assert get_device(angle(a)) == device