np.newaxis с Numba nopython
Есть ли способ использовать np.newaxis
с Нумбой nopython
? Для того, чтобы применить функцию вещания без отступления на питоне?
например
@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
Спасибо
2 ответа
В моей кассе (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
используя трансляцию.
Вы можете сделать это с помощью Reshape, похоже, [:, 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
Это можно сделать с помощью новейшей версии Numba (0.27) и numpy. stride_tricks
, Вы должны быть осторожны с этим, и это немного уродливо. Прочитайте документацию для as_strided
чтобы убедиться, что вы понимаете, что происходит, поскольку это не "безопасно", так как не проверяет форму или шаги.
import numpy as np
import numba as nb
a = np.random.randn(20, 10)
b = np.random.randn(20)
c = np.random.randn(10)
def toto(a, b, c):
d = a - b[:, np.newaxis] * c[np.newaxis, :]
return d
@nb.jit(nopython=True)
def toto2(a, b, c):
_b = np.lib.stride_tricks.as_strided(b, shape=(b.shape[0], 1), strides=(b.strides[0], 0))
_c = np.lib.stride_tricks.as_strided(c, shape=(1, c.shape[0]), strides=(0, c.strides[0]))
d = a - _b * _c
return d
x = toto(a,b,c)
y = toto2(a,b,c)
print np.allclose(x, y) # True