Add conj, imag, and real functions to numpy.array_api

Original NumPy Commit: 103bca57407ba69632c005609c58538c3a765123

Author: asmeurer

array_api_strict/__init__.py

@@ -234,6 +234,7 @@
     bitwise_right_shift,
     bitwise_xor,
     ceil,
+    conj,
     cos,
     cosh,
     divide,
@@ -244,6 +245,7 @@
     floor_divide,
     greater,
     greater_equal,
+    imag,
     isfinite,
     isinf,
     isnan,
@@ -263,6 +265,7 @@
     not_equal,
     positive,
     pow,
+    real,
     remainder,
     round,
     sign,

array_api_strict/_dtypes.py

@@ -83,6 +83,7 @@
     "integer or boolean": _integer_or_boolean_dtypes,
     "boolean": _boolean_dtypes,
     "real floating-point": _floating_dtypes,
+    "complex floating-point": _complex_floating_dtypes,
     "floating-point": _floating_dtypes,
 }
 

array_api_strict/_elementwise_functions.py

@@ -3,6 +3,7 @@
 from ._dtypes import (
     _boolean_dtypes,
     _floating_dtypes,
+    _complex_floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
     _numeric_dtypes,
@@ -238,6 +239,17 @@ def ceil(x: Array, /) -> Array:
     return Array._new(np.ceil(x._array))
 
 
+def conj(x: Array, /) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.conj <numpy.conj>`.
+
+    See its docstring for more information.
+    """
+    if x.dtype not in _complex_floating_dtypes:
+        raise TypeError("Only complex floating-point dtypes are allowed in conj")
+    return Array._new(np.conj(x))
+
+
 def cos(x: Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.cos <numpy.cos>`.
@@ -364,6 +376,17 @@ def greater_equal(x1: Array, x2: Array, /) -> Array:
     return Array._new(np.greater_equal(x1._array, x2._array))
 
 
+def imag(x: Array, /) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.imag <numpy.imag>`.
+
+    See its docstring for more information.
+    """
+    if x.dtype not in _complex_floating_dtypes:
+        raise TypeError("Only complex floating-point dtypes are allowed in imag")
+    return Array._new(np.imag(x))
+
+
 def isfinite(x: Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.isfinite <numpy.isfinite>`.
@@ -599,6 +622,17 @@ def pow(x1: Array, x2: Array, /) -> Array:
     return Array._new(np.power(x1._array, x2._array))
 
 
+def real(x: Array, /) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.real <numpy.real>`.
+
+    See its docstring for more information.
+    """
+    if x.dtype not in _complex_floating_dtypes:
+        raise TypeError("Only complex floating-point dtypes are allowed in real")
+    return Array._new(np.real(x))
+
+
 def remainder(x1: Array, x2: Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.remainder <numpy.remainder>`.

array_api_strict/tests/test_elementwise_functions.py

@@ -38,6 +38,7 @@ def test_function_types():
         "bitwise_right_shift": "integer",
         "bitwise_xor": "integer or boolean",
         "ceil": "numeric",
+        "conj": "complex floating-point",
         "cos": "floating-point",
         "cosh": "floating-point",
         "divide": "floating-point",
@@ -48,6 +49,7 @@ def test_function_types():
         "floor_divide": "numeric",
         "greater": "numeric",
         "greater_equal": "numeric",
+        "imag": "complex floating-point",
         "isfinite": "numeric",
         "isinf": "numeric",
         "isnan": "numeric",
@@ -67,6 +69,7 @@ def test_function_types():
         "not_equal": "all",
         "positive": "numeric",
         "pow": "numeric",
+        "real": "complex floating-point",
         "remainder": "numeric",
         "round": "numeric",
         "sign": "numeric",