np.newaxis with Numba nopython

Question

Is there a way to use np.newaxis with Numba nopython ? In order to apply broadcasting function without fallbacking on python ?

for example

@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

Thanks

Answer 1

You can accomplish this using reshape, it looks like the [:, None] indexing isn't currently supported. Note that this probably won't be much faster than doing it python, since it was already vectorized.

@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