python - 将 Cython 融合类型转换为 C++ 指针

Question

这是一个关于从 Cython 融合类型转换为 C++ 类型的一般性问题，我将用一个最小的例子来描述。考虑肤浅的 C++ 函数模板:

template <typename T>
void scale_impl(const T * x, T * y, const T a, const size_t N) {
    for (size_t n = 0; n < N; ++n) {
        y[n] = a*x[n];
    }
}

我希望能够在任何 numpy ndarray 上调用这个函数任何类型和形状。使用 Cython，我们首先声明函数模板:

cdef extern:
    void scale_impl[T](const T * x, T * y, const T a, const size_t N)

然后声明我们希望操作的有效标量类型:

ctypedef fused Scalar:
    float
    double
    ...

最后实现实际的 Cython shim:

def scale(ndarray[Scalar] x, Scalar a):
    """Scale an array x by the value a"""
    cdef ndarray[Scalar] y = np.empty_like(x)
    scale_impl(<Scalar *>x.data, <Scalar *>y.data, a, x.size)
    return y

这行不通有两个原因:

• x只能是一维的，不能是任意的(或至少是多维的)
• 转换到<Scalar *>抛出错误，因为 Scalar实际上是一个Python对象

显然可以明确地推断出特化:

    if Scalar is float:
        scale_impl(<float *>x.data, <float *>y.data, a, x.size)
    if Scalar is double:
        scale_impl(<double *>x.data, <double *>y.data, a, x.size)
    if Scalar is ...

但这会导致我必须为包含多个融合类型的函数手写组合数量的代码路径，并造成(我假设)引入融合类型以避免的情况。

有没有办法将任意维(在合理范围内)数组传递给 Cython 函数，并让它推断出标量数据的指针类型？或者，实现此类功能的最合理折衷是什么？

Answer 1

(另请参阅 Using Cython to wrap a c++ template to accept any numpy array 中给出的答案，这是一个非常相似的问题。)

使用 &x[0] 形式而不是尝试强制转换 x.data 解决了选择正确模板特化的问题。二维数组的问题有点复杂，因为不能保证数组是连续的或有序的。

我将创建一个函数来对一维数组执行实际工作，并包装在一个根据需要展平的简单函数中:

def _scale_impl(Scalar[::1] x, Scalar a):
  # the "::1" syntax ensures the array is actually continuous
  cdef np.ndarray[Scalar,ndim=1] y = np.empty_like(x)
  cdef size_t N = x.shape[0] # this seems to be necessary to avoid throwing off Cython's template deduction
  scale_impl(&x[0],&y[0],a,N)
  return y


def scale(x, a):
  """Scale an array x by the value a"""

  y = _scale_impl(np.ravel(x),a)
  return y.reshape(x.shape) # reshape needs to be kept out of Cython