时间戳

首页 > 解决方案 > Numba 矢量化与常数的参与

python - Numba 矢量化与常数的参与

问题描述

我在 GPU 设备上有一个 Numba 矩阵，想用一个常数替换 < 0 的值。如何做到这一点？就像是

import numba
from numba import cuda, void, int32, int64, float32, float64
import numpy as np

@numba.vectorize([int32[:, :](int32[:, :], int32)], '(i, j), () -> (i, j)', target='cuda')
def replace(a, c):  
    if a[i, j] < 0:
        return c
    else:
        return a[i, j]

b = np.array([-1, 1, 2, 3, -1], dtype=np.int32)
b = cuda.to_device(b)
replace(b, 100)

?

这产生

TypeError: vectorize() takes from 0 to 1 positional arguments but 2 were given

另一个问题，是否可以就地实施这样的事情？

标签： pythonvectorizationnumba

解决方案

嗯。它通过外壳工作

cc = self.c    # I am declaring and running it from class' method now

@numba.vectorize([int32(int32)], target='cuda')
def replace(b):  
    if b < 0:
        return cc   # Btw, direct self.c here causes error
    else:
        return b

...

但这意味着对于典型的调用，我们需要额外的包装类或类似的东西......

推荐阅读