np.newaxis 与 Numba nopython答案

【问题标题】：np.newaxis with Numba nopythonnp.newaxis 与 Numba nopython
【发布时间】：2016-12-10 04:43:08
【问题描述】：

有没有办法将 np.newaxis 与 Numba 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

谢谢

【问题讨论】：

标签： numpy numba array-broadcasting

【解决方案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。

【讨论】：

【解决方案2】：

这可以使用最新版本的 Numba (0.27) 和 numpy stride_tricks 来完成。你需要小心这个，它有点难看。阅读docstring 的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

【讨论】：

【解决方案3】：

您可以使用 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

【讨论】：

我已经尝试过了，但我得到了：reshape() supports contiguous array only。当然，toto() 只是一个例子，不是我的实际功能
你可以做b.copy().reshape((-1,1))。如果您的数组不连续，我相信无论如何都会复制，尽管不是 100% 肯定。