ENH: angle: add delegation

src/array_api_extra/__init__.py

@@ -1,11 +1,11 @@
 """Extra array functions built on top of the array API standard."""
 
 from . import testing
-from ._agnostic._elementwise import angle, apply_where
+from ._agnostic._elementwise import apply_where
 from ._agnostic._inspection import default_dtype
 from ._at import at
 from ._creation import create_diagonal, one_hot
-from ._elementwise import deg2rad, isclose, nan_to_num, rad2deg, sinc
+from ._elementwise import angle, deg2rad, isclose, nan_to_num, rad2deg, sinc
 from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index
 from ._lazy import lazy_apply
 from ._linalg import kron

src/array_api_extra/_agnostic/_elementwise.py

@@ -7,7 +7,6 @@
 from .._at import at
 from .._lib import _compat, _helpers
 from .._lib._typing import Array, ArrayNamespace
-from . import _inspection
 
 __all__ = [
     "angle",
@@ -291,46 +290,13 @@ def sinc(x: Array, /, *, xp: ArrayNamespace) -> Array:
     return xp.sin(y) / y
 
 
-def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array:
-    """
-    Return the angle of the complex argument.
-
-    Parameters
-    ----------
-    z : Array
-        Input array.
-    deg : bool, optional
-        Return angle in degrees if True, radians if False (default).
-    xp : array_namespace, optional
-        The standard-compatible namespace for `z`. Default: infer.
-
-    Returns
-    -------
-    array
-        The counterclockwise angle from the positive real axis on the complex
-        plane in the range ``(-pi, pi]``.
-
-    Notes
-    -----
-    Real input ``x`` is interpreted as ``x + 0j``.
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), xp=xp)
-    Array([0.        , 1.57079633, 0.78539816], dtype=array_api_strict.float64)
-    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp)
-    Array([ 0., 90., 45.], dtype=array_api_strict.float64)
-    """
-    if xp is None:
-        xp = _compat.array_namespace(z)
+def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace) -> Array:
+    # numpydoc ignore=PR01,RT01
+    """See docstring in `array_api_extra._elementwise`."""
     if xp.isdtype(z.dtype, "complex floating"):
         zimag = xp.imag(z)
         zreal = xp.real(z)
     else:
-        if not xp.isdtype(z.dtype, "real floating"):
-            z = xp.astype(z, _inspection.default_dtype(xp, device=_compat.device(z)))
         zimag = xp.zeros_like(z)
         zreal = z
     a = xp.atan2(zimag, zreal)

src/array_api_extra/_elementwise.py

@@ -4,7 +4,59 @@
 from ._lib import _compat, _helpers
 from ._lib._typing import Array, ArrayNamespace
 
-__all__ = ["deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"]
+__all__ = ["angle", "deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"]
+
+
+def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array:
+    """
+    Return the angle of the complex argument.
+
+    Parameters
+    ----------
+    z : array
+        Input array. Real input is interpreted as having zero imaginary part.
+    deg : bool, optional
+        Return angle in degrees if True, radians if False (default).
+    xp : array_namespace, optional
+        The standard-compatible namespace for `z`. Default: infer.
+
+    Returns
+    -------
+    array
+        The counterclockwise angle from the positive real axis on the complex
+        plane in the range ``(-pi, pi]``.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), xp=xp)
+    Array([0.        , 1.57079633, 0.78539816], dtype=array_api_strict.float64)
+    >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp)
+    Array([ 0., 90., 45.], dtype=array_api_strict.float64)
+    """
+    if xp is None:
+        xp = _compat.array_namespace(z)
+
+    if not xp.isdtype(z.dtype, ("real floating", "complex floating")):
+        z = xp.astype(
+            z, _agnostic._inspection.default_dtype(xp, device=_compat.device(z))
+        )
+
+    if (
+        _compat.is_numpy_namespace(xp)
+        or _compat.is_cupy_namespace(xp)
+        or _compat.is_dask_namespace(xp)
+        or _compat.is_jax_namespace(xp)
+    ):
+        return xp.angle(z, deg=deg)
+
+    # Torch treats real negative zero as positive zero, unlike atan2(0, z).
+    if _compat.is_torch_namespace(xp) and xp.isdtype(z.dtype, "complex floating"):
+        result = xp.angle(z)
+        return result * 180 / xp.pi if deg else result
+
+    return _agnostic._elementwise.angle(z, deg=deg, xp=xp)
 
 
 def deg2rad(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array:

tests/main/test_elementwise.py

@@ -23,6 +23,7 @@
 from array_api_extra._lib._typing import Array, ArrayNamespace, Device
 from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function
 
+lazy_xp_function(angle)
 lazy_xp_function(apply_where)
 lazy_xp_function(deg2rad)
 lazy_xp_function(isclose)