极客笔记Numpy Numba 无法使用 np.mean 的问题
在本文中，我们将介绍在使用 Numpy 和 Numba 进行数值计算时遇到的一个常见问题——无法使用 Numpy 中的 np.mean 函数。
阅读更多：Numpy 教程
问题背景
在使用 Numpy 和 Numba 进行数值计算时，我们通常会使用 Numpy 的各种函数和 Numba 的 JIT 编译技术来加速计算过程。然而，在某些情况下，我们会在使用 Numba 进行 JIT 编译时遇到一些问题，例如，无法使用 Numpy 中的 np.mean 函数。
下面是一个简单的例子，展示了当我们在使用 Numba 对使用 np.mean 函数的 Numpy 代码进行 JIT 编译时会发生什么：
import numpy as np
import numba as nb
@nb.jit(nopython=True)
def mean_of_array(a):
return np.mean(a)
a = np.random.rand(1000000)
print(mean_of_array(a))
在上面的代码中，我们首先定义了一个使用 np.mean 函数的 Numpy 数组均值计算函数 mean_of_array，然后使用 Numba 对这个函数进行 JIT 编译，最后对一个由 1000000 个随机数构成的 Numpy 数组调用这个函数并打印结果。
然而，当我们运行这个代码时，会发现程序会抛出以下异常：
TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<function mean at 0x7f3ef8804390>) found for signature:
>>> mean(TensorType(float32, 1d, C))
There are 2 candidate implementations:
- Of which 2 did not match due to:
Overload of function 'mean': File: <numerous>: Line N/A.
With argument(s): '(TensorType(float32, 1d, C),)':
No match.
- Of which 2 did not match due to:
Overload of function 'mean': File: <numerous>: Line N/A.
With argument(s): '(TensorType(float32, 1d, C),)':
No match.
During: resolving callee type: Function(<function mean_of_array at 0x7f3ef1868a60>)
During: typing of call at <stdin>:9:5
Relevant code:
a = np.random.rand(1000000)
print(mean_of_array(a))
^
This error is usually caused by passing an argument of a type that is unsupported by the named function.
这个异常表明了 Numba 无法找到 Numpy np.mean 函数的有效实现。
问题探究
为了理解为什么会出现这个问题，我们需要首先了解 Numba 的 JIT 编译过程。Numba 的 JIT 编译过程分为两个阶段：首先是 Numba 的 Typing 阶段，该阶段负责确定每个函数参数和返回值的数据类型；然后是 Numba 的 Compilation 阶段，该阶段负责将 Numba 代码编译为本机代码。
在 Typing 阶段，Numba 会对每个函数参数和返回值进行类型推断。当 Numba 对使用 np.mean 函数的 Numpy 代码进行类型推断时，它会发现 np.mean 函数的实现中使用了一个名为 AxisReduce 技术的优化，该技术可以将数组的某个维度上的计算转化为一系列的迭代计算，从而提高计算效率。
然而，由于 Numba 的 Typing 阶段无法正确处理 AxisReduce 技术，因此当我们在使用 Numba 和 Numpy 进行数值计算时，如果使用了 np.mean 函数，那么在 Numba 的 Typing 阶段就会出现错误，导致 JIT 编译失败。
解决方案
为了解决这个问题，我们可以使用 Numba 的 guvectorize 装饰器来替换 np.mean 函数。gufunc 是一种可以将常规多元函数转换为适用于 ndarrays 的通用函数的装饰器。通过使用 gufunc 装饰器，我们可以将 Numpy 函数转换为支持 Numba JIT 编译的函数，从而避免了 np.mean 函数在 Numba Typing 阶段的问题。
下面是一个使用 gufunc 装饰器替换 np.mean 函数的例子：
import numpy as np
import numba as nb
@nb.guvectorize([(nb.float64[:], nb.float64[:])], '(n)->()')
def mean_of_array(a, res):
res[0] = np.sum(a) / a.shape[0]
a = np.random.rand(1000000)
print(mean_of_array(a))
在上面的代码中，我们首先定义了一个新的函数 mean_of_array，该函数使用 gufunc 装饰器将常规函数转换为适用于 ndarrays 的通用函数。在装饰器中，我们使用了一个元组和一个字符串来定义输入和输出的数据类型及形状。然后，在函数中我们使用了 Numpy 的 np.sum 函数来计算数组的和，并除以数组的大小，从而计算数组均值。
最后，我们对一个由 1000000 个随机数构成的 Numpy 数组调用这个新的均值计算函数并打印结果。这里需要注意的是，由于我们的函数返回的是一个标量值，因此需要将输出的形状定义为 ‘()’，这样才能正确地计算出数组的均值。
总结
在本文中，我们介绍了在使用 Numpy 和 Numba 进行数值计算时遇到的一个常见问题——无法使用 Numpy 中的 np.mean 函数。我们探究了这个问题的原因，并提出了使用 gufunc 装饰器来替换 np.mean 函数的解决方案。通过使用 gufunc 装饰器，我们可以将 Numpy 函数转换为支持 Numba JIT 编译的函数，从而使我们能够在使用 Numpy 和 Numba 进行数值计算时更加高效和方便。