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

Question

我想使用Numba来矢量化一个函数，该函数将评估矩阵的每一行。这基本上会将 Numpy ufunc 应用于矩阵，而不是循环遍历行。根据文档：

你可能会问自己，“为什么我要经历这个而不是使用 @jit 装饰器编译一个简单的迭代循环？”。答案是 NumPy ufunc 会自动获得其他功能，例如减少、累积或广播。

考虑到这一点，我什至无法让一个玩具示例起作用。下面的简单示例尝试计算每行中元素的总和。

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() can't ever modify the original value of b in python syntax.

numba gets around this by requiring every param to a gufunc be an array - scalars are just length 1, that you can then assign into. So you need both params as arrays, and the assignment must use []

@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 函数（文档此处 和此处）。

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

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]