Use cmath where obvious

Author: honno

array_api_tests/pytest_helpers.py

@@ -1,3 +1,4 @@
+import cmath
 import math
 from inspect import getfullargspec
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
@@ -345,12 +346,12 @@ def assert_scalar_equals(
     if type_ in [bool, int]:
         msg = f"{repr_name}={out}, but should be {expected} [{f_func}]"
         assert out == expected, msg
-    elif math.isnan(expected):
+    elif cmath.isnan(expected):
         msg = f"{repr_name}={out}, but should be {expected} [{f_func}]"
-        assert math.isnan(out), msg
+        assert cmath.isnan(out), msg
     else:
         msg = f"{repr_name}={out}, but should be roughly {expected} [{f_func}]"
-        assert math.isclose(out, expected, rel_tol=0.25, abs_tol=1), msg
+        assert cmath.isclose(out, expected, rel_tol=0.25, abs_tol=1), msg
 
 
 def assert_fill(
@@ -368,7 +369,7 @@ def assert_fill(
 
     """
     msg = f"out not filled with {fill_value} [{func_name}({fmt_kw(kw)})]\n{out=}"
-    if math.isnan(fill_value):
+    if cmath.isnan(fill_value):
         assert xp.all(xp.isnan(out)), msg
     else:
         assert xp.all(xp.equal(out, xp.asarray(fill_value, dtype=dtype))), msg

array_api_tests/test_array_object.py

@@ -1,3 +1,4 @@
+import cmath
 import math
 from itertools import product
 from typing import List, Sequence, Tuple, Union, get_args
@@ -135,7 +136,7 @@ def test_setitem(shape, dtypes, data):
     f_res = sh.fmt_idx("x", key)
     if isinstance(value, get_args(Scalar)):
         msg = f"{f_res}={res[key]!r}, but should be {value=} [__setitem__()]"
-        if math.isnan(value):
+        if cmath.isnan(value):
             assert xp.isnan(res[key]), msg
         else:
             assert res[key] == value, msg

array_api_tests/test_creation_functions.py

@@ -1,3 +1,4 @@
+import cmath
 import math
 from itertools import count
 from typing import Iterator, NamedTuple, Union
@@ -247,8 +248,8 @@ def test_asarray_scalars(shape, data):
 
 
 def scalar_eq(s1: Scalar, s2: Scalar) -> bool:
-    if math.isnan(s1):
-        return math.isnan(s2)
+    if cmath.isnan(s1):
+        return cmath.isnan(s2)
     else:
         return s1 == s2
 

array_api_tests/test_set_functions.py

@@ -1,4 +1,5 @@
 # TODO: disable if opted out, refactor things
+import cmath
 import math
 from collections import Counter, defaultdict
 
@@ -61,7 +62,7 @@ def test_unique_all(x):
 
     for idx in sh.ndindex(out.indices.shape):
         val = scalar_type(out.values[idx])
-        if math.isnan(val):
+        if cmath.isnan(val):
             break
         i = int(out.indices[idx])
         expected = firsts[val]
@@ -88,7 +89,7 @@ def test_unique_all(x):
     for idx in sh.ndindex(out.values.shape):
         val = scalar_type(out.values[idx])
         count = int(out.counts[idx])
-        if math.isnan(val):
+        if cmath.isnan(val):
             nans += 1
             assert count == 1, (
                 f"out.counts[{idx}]={count} for out.values[{idx}]={val}, "
@@ -225,7 +226,7 @@ def test_unique_values(x):
     nans = 0
     for idx in sh.ndindex(out.shape):
         val = scalar_type(out[idx])
-        if math.isnan(val):
+        if cmath.isnan(val):
             nans += 1
         else:
             assert val in distinct, f"out[{idx}]={val}, but {val} not in input array"

array_api_tests/test_sorting_functions.py

@@ -1,4 +1,4 @@
-import math
+import cmath
 from typing import Set
 
 import pytest
@@ -26,7 +26,7 @@ def assert_scalar_in_set(
     **kw,
 ):
     out_repr = "out" if idx == () else f"out[{idx}]"
-    if math.isnan(out):
+    if cmath.isnan(out):
         raise NotImplementedError()
     msg = f"{out_repr}={out}, but should be in {set_} [{func_name}({ph.fmt_kw(kw)})]"
     assert out in set_, msg

array_api_tests/test_statistical_functions.py

@@ -1,3 +1,4 @@
+import cmath
 import math
 from typing import Optional
 
@@ -162,7 +163,7 @@ def test_prod(x, data):
     scalar_type = dh.get_scalar_type(out.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
         prod = scalar_type(out[out_idx])
-        assume(math.isfinite(prod))
+        assume(cmath.isfinite(prod))
         elements = []
         for idx in indices:
             s = scalar_type(x[idx])
@@ -267,7 +268,7 @@ def test_sum(x, data):
     scalar_type = dh.get_scalar_type(out.dtype)
     for indices, out_idx in zip(sh.axes_ndindex(x.shape, _axes), sh.ndindex(out.shape)):
         sum_ = scalar_type(out[out_idx])
-        assume(math.isfinite(sum_))
+        assume(cmath.isfinite(sum_))
         elements = []
         for idx in indices:
             s = scalar_type(x[idx])