python - Python numba：如何找到数组中最大元素的位置

Question

我正在编写用于实时处理来自相机的图像的代码。我使用 Python 3.5 和 Anaconda Accelerate/Numba 包在 GPU 上执行大部分计算。我在实现一个函数时遇到问题，该函数将在 float32 二维数组中找到最大元素的位置。该数组已经在 GPU 内存中。问题是：它非常慢。这是我整个代码的瓶颈。编码：

@n_cuda.jit('void(float32[:,:], float32, float32, float32)')
def d_findcarpeak(temp_mat, height, width, peak_flat):
    row, col = cuda.grid(2)
    if row < height and col < width:
        peak_flat = temp_mat.argmax()

这就是我所说的：

d_findcarpeak[number_of_blocks, threads_per_block](
            d_temp_mat, height, width, d_peak_flat)

我怎样才能重写这段代码？