support complex arrays in cupy.allclose
from cupy import core
+from numpy import complex64, complex128
_is_close = core.create_ufunc(
@@ -16,6 +17,24 @@ _is_close = core.create_ufunc(
'''
)
+# Note that in cupy/core/include/cupy/complex.cuh, we already got isfinite and
+# isnan working for complex numbers, so just replace fabs above by abs (from
+# thrust) and we are ready to go
+_is_close_complex = core.create_ufunc(
+ 'cupy_is_close_complex',
+ ('FFff?->?', 'DDdd?->?'),
+ '''
+ bool equal_nan = in4;
+ if (isfinite(in0) && isfinite(in1)) {
+ out0 = abs(in0 - in1) <= in3 + in2 * abs(in1);
+ } else if (equal_nan) {
+ out0 = (in0 == in1) || (isnan(in0) && isnan(in1));
+ } else {
+ out0 = (in0 == in1);
+ }
+ '''
+)
+
def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
"""Returns True if two arrays are element-wise equal within a tolerance.
@@ -70,6 +89,10 @@ def isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
.. seealso:: :func:`numpy.isclose`
"""
+ if (a.dtype in [complex64, complex128]) or \
+ (b.dtype in [complex64, complex128]):
+ return _is_close_complex(a, b, rtol, atol, equal_nan)
+ else:
return _is_close(a, b, rtol, atol, equal_nan)
