python - 如何使 cython 函数接受 float 或 double 数组输入？

Question

假设我有以下 (MCVE...) cython 函数

cimport cython

from scipy.linalg.cython_blas cimport dnrm2


cpdef double func(int n, double[:] x):
   cdef int inc = 1
   return dnrm2(&n, &x[0], &inc)

然后，我无法在 np.float32 数组 x 上调用它。

如何让 func 接受 double[:] 或 float[:]，并调用 dnrm2 或 snrm2 或者？我目前唯一的解决方案是拥有两个函数，这会产生大量重复代码。

Answer 1

您可以使用融合类型。请注意，以下代码无法在我的系统上编译，因为 ddot 和 sdot 显然需要 5 个参数:

# cython: infer_types=True
cimport cython

from scipy.linalg.cython_blas cimport ddot, sdot

ctypedef fused anyfloat:
   double
   float

cpdef anyfloat func(int n, anyfloat[:] x):
   cdef int inc = 1
   if anyfloat is double:
      return ddot(&n, &x[0], &inc)
   else:
      return sdot(&n, &x[0], &inc)