Np.newaxis与Numba nopython

Question

np.newaxisNumba 有使用方法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

在我的情况下 ( numba: 0.35, numpy: 1.14.0) expand_dims工作正常：

import numpy as np
from numba import jit

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

当然，我们可以expand_dims使用广播省略第二个。

Answer 2

您可以使用整形来完成此操作，好像[:, 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