numpy - Np.newaxis与Numba nopython

Question

有没有办法将Numba np.newaxis与nopython一起使用？为了应用广播功能而不回退python？

例如

@jit(nopython=True)
def toto():
    a = np.random.randn(20, 10)
    b = np.random.randn(20) 
    c = np.random.randn(10)
    d = a - b[:, np.newaxis] * c[np.newaxis, :]
    return d

谢谢

Answer 1

您可以使用整形来完成此操作，看起来当前不支持[:, None]索引。注意，这可能不会比用python快得多，因为它已经被矢量化了。

@jit(nopython=True)
def toto():
    a = np.random.randn(20, 10)
    b = np.random.randn(20) 
    c = np.random.randn(10)
    d = a - b.reshape((-1, 1)) * c.reshape((1,-1))
    return d