具有任意维度的 Cython ndarray

Question

我正在编写一些需要能够处理具有任意维数的 NumPy ndarray 的 Cython 代码。目前，我只有不同的函数来接受不同大小的 ndarrays，有点像：

def func1(np.ndarray[DTYPE_float64_t, ndim=1] arr):
    # Do something with the 1-D ndarray.

def func2(np.ndarray[DTYPE_float64_t, ndim=2] arr):
    # Do something with the 2-D ndarray.

def func3(np.ndarray[DTYPE_float64_t, ndim=3] arr):
    # Do something with the 3-D ndarray.

但是我想编写一个以任意维度的 ndarray 作为参数的通用函数。我试着简单地忽略“ndim”参数，但是 Cython 假设 ndim=1，这不好。

有没有办法做到这一点，还是我只需要为每个维数编写一个函数？

Answer 1

如果你只是想做一些元素方面的事情，诀窍就是获取数组的一维视图并对其进行操作

def func(arr):
   shape = arr.shape
   output = _func_impl(arr.ravel())
   return output.reshape(shape) # ensure that the output is the same shape
       # as the input. Skip this if it doesn't make sense!

def _func_impl(np.ndarray[DTYPE_float64_t, ndim=1] arr):
   # do something useful