Rudimentary complex testing for unary elwise functions

Author: honno

array_api_tests/dtype_helpers.py

@@ -15,6 +15,7 @@
     "uint_dtypes",
     "all_int_dtypes",
     "float_dtypes",
+    "real_dtypes",
     "numeric_dtypes",
     "all_dtypes",
     "dtype_to_name",
@@ -30,6 +31,7 @@
     "promotion_table",
     "dtype_nbits",
     "dtype_signed",
+    "dtype_components",
     "func_in_dtypes",
     "func_returns_bool",
     "binary_op_to_symbol",
@@ -86,14 +88,19 @@ def __repr__(self):
 _uint_names = ("uint8", "uint16", "uint32", "uint64")
 _int_names = ("int8", "int16", "int32", "int64")
 _float_names = ("float32", "float64")
-_dtype_names = ("bool",) + _uint_names + _int_names + _float_names
+_real_names = _uint_names + _int_names + _float_names
+_complex_names = ("complex64", "complex128")
+_numeric_names = _real_names + _complex_names
+_dtype_names = ("bool",) + _numeric_names
 
 
 uint_dtypes = tuple(getattr(xp, name) for name in _uint_names)
 int_dtypes = tuple(getattr(xp, name) for name in _int_names)
 float_dtypes = tuple(getattr(xp, name) for name in _float_names)
 all_int_dtypes = uint_dtypes + int_dtypes
-numeric_dtypes = all_int_dtypes + float_dtypes
+real_dtypes = all_int_dtypes + float_dtypes
+complex_dtypes = tuple(getattr(xp, name) for name in _complex_names)
+numeric_dtypes = real_dtypes + complex_dtypes
 all_dtypes = (xp.bool,) + numeric_dtypes
 bool_and_all_int_dtypes = (xp.bool,) + all_int_dtypes
 
@@ -129,6 +136,8 @@ def get_scalar_type(dtype: DataType) -> ScalarType:
         return int
     elif is_float_dtype(dtype):
         return float
+    elif dtype in complex_dtypes:
+        return complex
     else:
         return bool
 
@@ -157,7 +166,8 @@ class MinMax(NamedTuple):
     [(d, 8) for d in [xp.int8, xp.uint8]]
     + [(d, 16) for d in [xp.int16, xp.uint16]]
     + [(d, 32) for d in [xp.int32, xp.uint32, xp.float32]]
-    + [(d, 64) for d in [xp.int64, xp.uint64, xp.float64]]
+    + [(d, 64) for d in [xp.int64, xp.uint64, xp.float64, xp.complex64]]
+    + [(xp.complex128, 128)]
 )
 
 
@@ -166,6 +176,11 @@ class MinMax(NamedTuple):
 )
 
 
+dtype_components = EqualityMapping(
+    [(xp.complex64, xp.float32), (xp.complex128, xp.float64)]
+)
+
+
 if isinstance(xp.asarray, _UndefinedStub):
     default_int = xp.int32
     default_float = xp.float32
@@ -226,6 +241,11 @@ class MinMax(NamedTuple):
     ((xp.float32, xp.float32), xp.float32),
     ((xp.float32, xp.float64), xp.float64),
     ((xp.float64, xp.float64), xp.float64),
+    # complex
+    ((xp.complex64, xp.complex64), xp.complex64),
+    ((xp.complex64, xp.complex128), xp.complex128),
+    ((xp.complex128, xp.complex128), xp.complex128),
+
 ]
 _numeric_promotions += [((d2, d1), res) for (d1, d2), res in _numeric_promotions]
 _promotion_table = list(set(_numeric_promotions))

array_api_tests/hypothesis_helpers.py

@@ -46,7 +46,7 @@
 shared_dtypes = shared(dtypes, key="dtype")
 shared_floating_dtypes = shared(floating_dtypes, key="dtype")
 
-_dtype_categories = [(xp.bool,), dh.uint_dtypes, dh.int_dtypes, dh.float_dtypes]
+_dtype_categories = [(xp.bool,), dh.uint_dtypes, dh.int_dtypes, dh.float_dtypes, dh.complex_dtypes]
 _sorted_dtypes = [d for category in _dtype_categories for d in category]
 
 def _dtypes_sorter(dtype_pair: Tuple[DataType, DataType]):

array_api_tests/pytest_helpers.py

@@ -374,6 +374,21 @@ def assert_fill(
         assert xp.all(xp.equal(out, xp.asarray(fill_value, dtype=dtype))), msg
 
 
+def _assert_float_element(at_out: Array, at_expected: Array, msg: str):
+    if xp.isnan(at_expected):
+        assert xp.isnan(at_out), msg
+    elif at_expected == 0.0 or at_expected == -0.0:
+        scalar_at_expected = float(at_expected)
+        scalar_at_out = float(at_out)
+        if is_pos_zero(scalar_at_expected):
+            assert is_pos_zero(scalar_at_out), msg
+        else:
+            assert is_neg_zero(scalar_at_expected)  # sanity check
+            assert is_neg_zero(scalar_at_out), msg
+    else:
+        assert at_out == at_expected, msg
+
+
 def assert_array_elements(
     func_name: str, out: Array, expected: Array, /, *, out_repr: str = "out", **kw
 ):
@@ -392,26 +407,26 @@ def assert_array_elements(
     dh.result_type(out.dtype, expected.dtype)  # sanity check
     assert_shape(func_name, out.shape, expected.shape, **kw)  # sanity check
     f_func = f"[{func_name}({fmt_kw(kw)})]"
-    if dh.is_float_dtype(out.dtype):
+    if out.dtype in dh.float_dtypes:
+        for idx in sh.ndindex(out.shape):
+            at_out = out[idx]
+            at_expected = expected[idx]
+            msg = (
+                f"{sh.fmt_idx(out_repr, idx)}={at_out}, should be {at_expected} "
+                f"{f_func}"
+            )
+            _assert_float_element(at_out, at_expected, msg)
+    elif out.dtype in dh.complex_dtypes:
+        assert (out.dtype in dh.complex_dtypes) == (expected.dtype in dh.complex_dtypes)
         for idx in sh.ndindex(out.shape):
             at_out = out[idx]
             at_expected = expected[idx]
             msg = (
                 f"{sh.fmt_idx(out_repr, idx)}={at_out}, should be {at_expected} "
                 f"{f_func}"
             )
-            if xp.isnan(at_expected):
-                assert xp.isnan(at_out), msg
-            elif at_expected == 0.0 or at_expected == -0.0:
-                scalar_at_expected = float(at_expected)
-                scalar_at_out = float(at_out)
-                if is_pos_zero(scalar_at_expected):
-                    assert is_pos_zero(scalar_at_out), msg
-                else:
-                    assert is_neg_zero(scalar_at_expected)  # sanity check
-                    assert is_neg_zero(scalar_at_out), msg
-            else:
-                assert at_out == at_expected, msg
+            _assert_float_element(at_out.real, at_expected.real, msg)
+            _assert_float_element(at_out.imag, at_expected.imag, msg)
     else:
         assert xp.all(
             out == expected

array_api_tests/test_operators_and_elementwise_functions.py

@@ -103,6 +103,8 @@ def default_filter(s: Scalar) -> bool:
     """
     if isinstance(s, int):  # note bools are ints
         return True
+    elif isinstance(s, complex):
+        return default_filter(s.real) and default_filter(s.imag)
     else:
         return math.isfinite(s) and s != 0
 
@@ -247,7 +249,12 @@ def unary_assert_against_refimpl(
     in_stype = dh.get_scalar_type(in_.dtype)
     if res_stype is None:
         res_stype = in_stype
-    m, M = dh.dtype_ranges.get(res.dtype, (None, None))
+    if res.dtype == xp.bool:
+        m, M = (None, None)
+    if res.dtype in dh.complex_dtypes:
+        m, M = dh.dtype_ranges[dh.dtype_components[res.dtype]]
+    else:
+        m, M = dh.dtype_ranges[res.dtype]
     for idx in sh.ndindex(in_.shape):
         scalar_i = in_stype(in_[idx])
         if not filter_(scalar_i):
@@ -257,9 +264,13 @@ def unary_assert_against_refimpl(
         except Exception:
             continue
         if res.dtype != xp.bool:
-            assert m is not None and M is not None  # for mypy
-            if expected <= m or expected >= M:
-                continue
+            if res.dtype in dh.complex_dtypes:
+                for component in [expected.real, expected.imag]:
+                    if component <= m or expected >= M:
+                        continue
+            else:
+                if expected <= m or expected >= M:
+                    continue
         scalar_o = res_stype(res[idx])
         f_i = sh.fmt_idx("x", idx)
         f_o = sh.fmt_idx("out", idx)
@@ -418,8 +429,11 @@ def __repr__(self):
 
 
 def make_unary_params(
-    elwise_func_name: str, dtypes_strat: st.SearchStrategy[DataType]
+    elwise_func_name: str, dtypes: Sequence[DataType]
 ) -> List[Param[UnaryParamContext]]:
+    if hh.FILTER_UNDEFINED_DTYPES:
+        dtypes = [d for d in dtypes if not isinstance(d, xp._UndefinedStub)]
+    dtypes_strat = st.sampled_from(dtypes)
     strat = xps.arrays(dtype=dtypes_strat, shape=hh.shapes())
     func_ctx = UnaryParamContext(
         func_name=elwise_func_name, func=getattr(xp, elwise_func_name), strat=strat
@@ -633,7 +647,7 @@ def binary_param_assert_against_refimpl(
         )
 
 
-@pytest.mark.parametrize("ctx", make_unary_params("abs", xps.numeric_dtypes()))
+@pytest.mark.parametrize("ctx", make_unary_params("abs", dh.numeric_dtypes))
 @given(data=st.data())
 def test_abs(ctx, data):
     x = data.draw(ctx.strat, label="x")
@@ -643,7 +657,10 @@ def test_abs(ctx, data):
 
     out = ctx.func(x)
 
-    ph.assert_dtype(ctx.func_name, x.dtype, out.dtype)
+    if x.dtype in dh.complex_dtypes:
+        assert out.dtype == dh.complex_components[x.dtype]
+    else:
+        ph.assert_dtype(ctx.func_name, x.dtype, out.dtype)
     ph.assert_shape(ctx.func_name, out.shape, x.shape)
     unary_assert_against_refimpl(
         ctx.func_name,
@@ -783,7 +800,7 @@ def test_bitwise_left_shift(ctx, data):
 
 
 @pytest.mark.parametrize(
-    "ctx", make_unary_params("bitwise_invert", boolean_and_all_integer_dtypes())
+    "ctx", make_unary_params("bitwise_invert", dh.bool_and_all_int_dtypes)
 )
 @given(data=st.data())
 def test_bitwise_invert(ctx, data):
@@ -1187,9 +1204,7 @@ def test_multiply(ctx, data):
 
 
 # TODO: clarify if uints are acceptable, adjust accordingly
-@pytest.mark.parametrize(
-    "ctx", make_unary_params("negative", xps.integer_dtypes() | xps.floating_dtypes())
-)
+@pytest.mark.parametrize("ctx", make_unary_params("negative", dh.numeric_dtypes))
 @given(data=st.data())
 def test_negative(ctx, data):
     x = data.draw(ctx.strat, label="x")
@@ -1226,7 +1241,7 @@ def test_not_equal(ctx, data):
     )
 
 
-@pytest.mark.parametrize("ctx", make_unary_params("positive", xps.numeric_dtypes()))
+@pytest.mark.parametrize("ctx", make_unary_params("positive", dh.numeric_dtypes))
 @given(data=st.data())
 def test_positive(ctx, data):
     x = data.draw(ctx.strat, label="x")
@@ -1317,7 +1332,7 @@ def test_square(x):
     ph.assert_dtype("square", x.dtype, out.dtype)
     ph.assert_shape("square", out.shape, x.shape)
     unary_assert_against_refimpl(
-        "square", x, out, lambda s: s ** 2, expr_template="{}²={}"
+        "square", x, out, lambda s: s**2, expr_template="{}²={}"
     )