由于 Numpy ufunc 签名导致的 Numba guvectorize 异常答案

【问题标题】：Numba guvectorize exception due to Numpy ufunc signature由于 Numpy ufunc 签名导致的 Numba guvectorize 异常
【发布时间】：2021-04-10 00:45:10
【问题描述】：

我正在尝试使用 @guvectorize 装饰器，它将从开始和停止向量返回数组的 axis=0 方法，例如，如果：

values = array([[82, 3],
                [76, 7],
                [23, 8],
                [46, 5],
                [38, 6]])
starts = array([ 0,  1,  1,  1,  4])
stops  = array([ 1,  4,  4,  4,  5])

我的函数会返回：

array([[82.        ,  3.        ],
       [48.33333333,  6.66666667],
       [48.33333333,  6.66666667],
       [48.33333333,  6.66666667],
       [38.        ,  6.        ]])

按照simple example in the numba documentation，这是我想出的（函数对一列的向量进行操作）：

@guvectorize([(int64[:], int64[:], int64[:], float64[:])], '(n),(n),(n)->(n)', nopython=True)
def variable_window_avg(values, starts, stops, means):
    for i in range(values.shape[0]):
        means[i] = np.nanmean(values[starts[i]:stops[i]])

当values 是向量时，这很有效，但当值是所需的 ndarray 时，它不会按比例放大：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-128-8a213e6ace9c> in <module>
----> 1 variable_window_avg(arr[:,:2], arr[:,1], arr[:,2])

ValueError: variable_window_avg: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n),(n),(n)->(n) (size 100 is different from 2)

错误表明形状没有通过 Numpy 通用通用函数 API 正确广播在一起，确实

res = variable_window_avg(values.T, starts, stops)

res = res.T

有效。但是修改输入和输出布局签名似乎是避免两个转置的更好解决方案。我也不明白为什么 numba 文档中的示例但我的代码失败了。

【问题讨论】：

标签： python arrays numpy numba

【解决方案1】：

我了解该算法的目的是获取向量切片的平均值。例如，有向量 [1,2,3,4,5] 和切片 1:4，结果将是 3。如果我有错误的想法，请纠正我。

矢量化背后的想法不是在知道您将使用的维度的情况下实现函数。相反，您只需要实现“内核”或最简单的情况。对于这个算法，它将是：

@nb.guvectorize([(nb.int64[:], nb.int64, nb.int64, nb.float64[:])],
                '(n),(),()->()', nopython=True)
def variable_window_avg(values, start, stop, means):
    means[:] = np.nanmean(values[start:stop])

注意签名'(n),(),()->()'表示输入是一个向量和两个标量，输出是另一个标量。矢量化机制将根据每种情况下的输入参数扩展维度。

例如，最简单的情况不需要任何向量化：

>>> variable_window_avg([3,7,8,5,6], 1, 4)
6.666666666666667

当使用多个索引对时，应用矢量化：

>>> variable_window_avg([3,7,8,5,6], [0,1,4], [1,4,5])
array([3.        , 6.66666667, 6.        ])

当使用二维数组和单个索引对时：

>>> variable_window_avg([[82,76,23,46,38], [3,7,8,5,6]], 1, 4)
array([48.33333333,  6.66666667])

但是当使用数组和多个索引时，会出现广播问题：

>>>> variable_window_avg([[82,76,23,46,38], [3,7,8,5,6]], [0,1,4], [1,4,5])
Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File ".../numba/np/ufunc/gufunc.py", line 151, in __call__
    return self.ufunc(*args, **kwargs)
ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (2,6)->(2,) (3,)->(3,) (3,)->(3,)  and requested shape ()

所以你需要转置索引：

>>> variable_window_avg([[82,76,23,46,38], [3,7,8,5,6]], [[0],[1],[4]], [[1],[4],[5]])
array([[82.        ,  3.        ],
       [48.33333333,  6.66666667],
       [38.        ,  6.        ]])

转置输入数组而不是索引将不起作用，因为该函数在数组的最内层执行切片。

【讨论】：