@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 中,您通常需要将输入类型传递给 vectorize 装饰器。在较新版本的 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 hereand 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]