麻木和numpy.expand U dims公司

Question

我正在重写一些函数以适合Numba。现在我有了一个函数，我在脚本中用不同维的输入数组调用了多次。在

def FormHistMatrix2(x,Whc,Lm):
    if x.ndim == 1:
       x = np.expand_dims(x,axis=1)
    [N,Ncells] = x.shape

这是我函数的开始，Numba抛出以下错误：

^{pr2}$

在本例中，“x”是二维数组，但在其他情况下，它可以是一维数组。 那么Numba不喜欢if循环吗？或者这里发生了什么？在

Answer 1

JoshAdel所说的在一般情况下是正确的，但是在这种情况下，您需要根据输入类型对函数进行不同的实现/专门化。在

Numba有这个例子的^{}-decorator。在

在您的情况下，需要编写一个专门的expand dims函数，该函数取决于输入数组的维数：

import numba as nb
@nb.generated_jit(nopython=True)
def nb_expander(x):
    if x.ndim == 1:
        return lambda x: np.expand_dims(x, axis=1)
    else:
        return lambda x: x

此函数需要从其他函数中调用：

^{pr2}$

这将适用于维度1和2的x。对于x.ndim==3，还需要为形状实现一个类似的方法。在

Answer 2

在Numba中，与标准python不同，变量在函数执行期间不能更改其类型。您应该能够将对np.expand_dims的调用结果分配给另一个变量，这样它就可以工作了。如果x有时是1d，有时是2d，只要函数执行过程中所有变量的类型都保持一致，就可以了。在