ENH: add angle delegation wrapper

Author: UshnahAbbasi

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/_elementwise.py

@@ -4,7 +4,41 @@
 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)
+
+    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)