如何使cython函数接受浮点或双数组输入？

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或{}？我目前唯一的解决方案是使用两个函数，这将创建大量重复的代码。在

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)