try to fix casting error for complex type
@@ -702,6 +702,13 @@ cpdef _ndarray_base _nanvar(_ndarray_base a, axis, dtype, out, ddof, keepdims):
if out is None:
out = _nanvar_core(
a, arrsum, _count, ddof, axis=axis, keepdims=keepdims)
+ else:
+ if a.dtype == cupy.complex64 or dtype == cupy.complex64:
+ out = _nanvar_core_complex64_out(
+ a, arrsum, _count, ddof, out, axis=axis, keepdims=keepdims)
+ elif a.dtype == cupy.complex128 or dtype == cupy.complex128:
+ out = _nanvar_core_complex128_out(
+ a, arrsum, _count, ddof, out, axis=axis, keepdims=keepdims)
else:
out = _nanvar_core_out(
a, arrsum, _count, ddof, out, axis=axis, keepdims=keepdims)
@@ -736,6 +743,12 @@ template <typename S, typename T>
__device__ T nanvar_impl(S x, T mean, long long alpha) {
return (isnan(x) ? T(0) : T(nanvar_complex_cast(x - mean))) / alpha;
}
+
+template <typename S, typename T>
+__device__ T nanvar_impl_complex(S x, complex<T> mean, long long alpha) {
+ return (isnan(x) ? T(0) : T(nanvar_complex_cast(x - mean))) / alpha;
+}
+
'''
@@ -751,6 +764,18 @@ cdef _nanvar_core_out = ReductionKernel(
'a + b', 'out = a', '0', '_nanvar_core', preamble=_nanvar_preamble)
+cdef _nanvar_core_complex64_out = ReductionKernel(
+ 'complex64 x, T sum, int64 _count, int64 ddof', 'float64 out',
+ 'nanvar_impl_complex<S, T>(x, sum / _count, max(_count - ddof, 0LL))',
+ 'a + b', 'out = a', '0', '_nanvar_core', preamble=_nanvar_preamble)
+
+
+cdef _nanvar_core_complex128_out = ReductionKernel(
+ 'complex128 x, T sum, int64 _count, int64 ddof', 'float64 out',
+ 'nanvar_impl_complex<S, T>(x, sum / _count, max(_count - ddof, 0LL))',
+ 'a + b', 'out = a', '0', '_nanvar_core', preamble=_nanvar_preamble)
+
+
# Variables to expose to Python
# (cythonized data cannot be exposed to Python, even with cpdef.)
