python - numba np.diff 轴=0

Question

在使用 numba 时，axis=0 是 np.sum() 可接受的参数，但不适用于 np.diff()。为什么会这样？我正在使用 2D，因此需要指定轴。

@jit(nopython=True)
def jitsum(y):
    np.sum(y, axis=0)

@jit(nopython=True)
def jitdiff(y): #this one will cause error
    np.diff(y, axis=0)

错误:np_diff_impl() 得到了一个意外的关键字参数“axis”

2D 中的解决方法是:

@jit(nopython=True)
def jitdiff(y):
    np.diff(y.T).T

Answer 1

np.diff 在 n=1 的二维数组上，axis=1 就是

a[:, 1:] - a[:, :-1]

对于axis=0:

a[1:, :] - a[:-1, :]

我怀疑上面的行可以用 numba 编译得很好。