当函数将数组作为参数时，如何使用 Numba 对函数进行矢量化？

Question

我想使用 Numba 来矢量化一个函数，该函数将计算矩阵的每一行。这实际上会将 Numpy ufunc 应用于矩阵，而不是循环遍历行。根据 docs:

You might ask yourself, “why would I go through this instead of compiling a simple iteration loop using the @jit decorator?”. The answer is that NumPy ufuncs automatically get other features such as reduction, accumulation or broadcasting.

考虑到这一点，我连一个玩具示例都无法运行。下面的简单示例尝试计算每行中元素的总和。

import numba, numpy as np

# Define the row-wise function to be vectorized:
@numba.guvectorize(["void(float64[:],float64)"],"(n)->()")
def f(a,b):
    b = a.sum() 

# Apply the function to an array with five rows:
a = np.arange(10).reshape(5,2)
b = f(a)

我使用了 @guvectorize 装饰器，因为我希望装饰函数将参数 a 作为矩阵的每一行，这是一个数组； @vectorize 只接受标量输入。我还编写了签名以获取数组参数并修改标量输出。根据 docs，装饰函数不使用 return 语句。

结果应该是b = [1,5,9,13,17]，但我却得到了b=[0.,1.,2.,3.,4.]。显然，我错过了一些东西。我会很感激一些方向，记住总和只是一个例子。

Answer 1

b = a.sum() 永远无法修改 python 语法中 b 的原始值。

numba 通过要求 gufunc 的每个参数都是一个数组来解决这个问题 - 标量的长度仅为 1，然后您可以将其分配给它。所以你需要两个参数都作为数组，并且赋值必须使用 []

@numba.guvectorize(["void(float64[:],float64[:])"],"(n)->()")
def f(a,b):
    b[:] = a.sum()
    # or b[0] = a.sum()

f(a)
Out[246]: array([ 1.,  5.,  9., 13., 17.])

Answer 2

@chrisb 在上面有一个很好的答案。这个答案应该为那些刚接触矢量化的人增加一些说明。

就向量化而言（在 numpy 和 numba 中），您传递输入向量。

例如：

import numpy as np

a=[1,2]
b=[3,4]

@np.vectorize
def f(x_1,x_2):
    return x_1+x_2

print(f(a,b))
#-> [4,6]

在 numba 中，传统上您需要将输入类型传递给向量化装饰器。在较新版本的 numba 中，如果将 numpy 数组作为输入传递给通用矢量化函数，则无需指定矢量输入类型。

例如：

import numpy as np
import numba as nb

a=np.array([1,2])
b=np.array([3,4])

# Note a generic vectorize decorator with input types not specified
@nb.vectorize
def f(x_1,x_2):
    return x_1+x_2

print(f(a,b))
#-> [4,6]

到目前为止，变量是从输入数组传递到函数中的简单单个对象。这使得 numba 可以将 python 代码转换为可以在 numpy 数组上运行的简单 ufunc。

在求和向量的示例中，您需要将数据作为向量的单个向量传递。为此，您需要创建对向量本身进行操作的 ufunc。这需要更多的工作和规范来确定如何创建任意输出输入 guvectorize 函数 (docs here and here)。

因为您提供的是向量的向量。您的外部向量的处理方式类似于您在上面使用 vectorize 的方式。现在您需要为您的输入值指定每个内部向量的外观。

EG 添加任意整数向量。（由于下面解释的一些原因，这将不起作用）

@nb.guvectorize([(nb.int64[:])])
def f(x):
    return x.sum()

现在您还需要向您的函数和装饰器添加额外的输入。这允许您指定任意类型来存储函数的输出。您现在将更新此输入变量，而不是返回输出。将此最终变量视为 numba 在为 numpy 评估创建 ufunc 时用于生成任意输出向量的自定义变量。

此输入也需要在装饰器中指定，您的函数应如下所示：

@nb.guvectorize([(nb.int64[:],nb.int64[:])])
def f(x, out):
    out[:]=x.sum()

最后你需要在装饰器中指定输入和输出格式。这些按输入向量的顺序作为矩阵形状给出，并使用箭头指示输出向量形状（实际上是您的最终输入）。在这种情况下，您将采用大小为 n 的向量并将结果输出为值而不是向量。您的格式应为 (n)->().

作为一个更复杂的示例，假设您有两个输入向量用于大小为 (m,n) 和 (n,o) 的矩阵乘法，并且您希望输出向量的大小为 (m,o) 您的装饰器格式看起来像 (m,n),(n,o)->(m,o).

当前问题的完整函数如下所示：

@nb.guvectorize([(nb.int64[:],nb.int64[:])], '(n)->()')
def f(x, out):
    out[:]=x.sum()

您的结束代码应该类似于：

import numpy as np
import numba as nb

a=np.arange(10).reshape(5,2)
# Equivalent to
# a=np.array([
#   [0,1],
#   [2,3],
#   [4,5],
#   [6,7],
#   [8,9]
# ])

@nb.guvectorize([(nb.int64[:],nb.int64[:])], '(n)->()')
def f(x, out):
    out[:]=x.sum()

print(f(a))
#-> [ 1  5  9 13 17]

当函数将数组作为参数时，如何使用 Numba 对函数进行矢量化？

How do I vectorize a function with Numba, when the function takes arrays as arguments?

python

numpy

scipy

numba