Use thrust::complex
@@ -9,7 +9,7 @@ from cupyx.scipy.sparse import csr_matrix
import numpy as np
-TYPES = ['double', 'complex<double>']
+TYPES = ['double', 'thrust::complex<double>']
INT_TYPES = ['int', 'long long']
INTERVAL_KERNEL = r'''
@@ -206,7 +206,12 @@ DESIGN_MAT_MODULE = cupy.RawModule(
def _get_module_func(module, func_name, *template_args):
- args_dtypes = [get_typename(arg.dtype) for arg in template_args]
+ def _get_typename(dtype):
+ typename = get_typename(dtype)
+ if dtype.kind == 'c':
+ typename = 'thrust::' + typename
+ return typename
+ args_dtypes = [_get_typename(arg.dtype) for arg in template_args]
template = ', '.join(args_dtypes)
kernel_name = f'{func_name}<{template}>' if template_args else func_name
kernel = module.get_function(kernel_name)
