python - Python numba:配列内の最大要素の位置を見つける方法

Question

カメラからの画像をリアルタイムで処理するためのコードを書いています。Python 3.5 と Anaconda Accelerate/Numba パッケージを使用して、ほとんどの計算を GPU で実行しています。float32 2d 配列の最大要素の位置を見つける関数の実装に問題があります。配列は既に 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)

このコードを書き直すにはどうすればよいですか?