时间戳

首页 > 解决方案 > numba 和 numpy.expand_dims

python - numba 和 numpy.expand_dims

问题描述

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

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

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

TypingError: Cannot unify array(float64, 2d, A) and array(float64, 3d, A) for 'x', defined at C:/Users/DNP_Student_3/Documents/Python Scripts/GCFuncsTests.py (332)

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

标签： pythonnumba

解决方案

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

Numba 有这个案例的@generated_jit-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

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

@nb.njit
def FormHistMatrix2(x, Whc, Lm):
    x = nb_expander(x)
    [N, Ncells] = x.shape

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

推荐阅读