- @vectorize 装饰器
Numba 的 @vectorize 装饰器可以将以标量为输入的的python函数编译为类似Numpy的 ufuncs。创建一个传统的NumPy ufunc并不是最简单的过程,它可能需要编写一些C代码。Numba让这很容易。使用@vectorize装饰器 ,Numba可以将纯Python函数编译成ufunc,该ufunc在NumPy数组上运行的速度与用C编写的传统ufunc一样快。
一个简单的例子:
代码语言:javascript复制import numpy as np
from numba import vectorize, float64
@vectorize(nopython=True)
def f(x, y):
return x y
a = np.arange(1,11)
b = np.arange(10,0,-1)
print(f(a,b))
以上是惰性编译,不指定参数类型。也可以采用积极编译,给出函数签名,即指定返回值和输入参数的类型。注意函数签名需写在列表中。
代码语言:javascript复制@vectorize([float64(float64, float64)], nopython=True)
def f(x, y):
return x y
还可以指定多个函数签名,需注意越通用的类型越要排在后面。
代码语言:javascript复制import numpy as np
from numba import vectorize, int32, int64, float32, float64
@vectorize([int32(int32, int32),
int64(int64, int64),
float32(float32, float32),
float64(float64, float64)])
def f(x, y):
return x y
此时如果传入其它类型,程序就会报错。
代码语言:javascript复制>>> a = np.linspace(0, 1 1j, 6)
>>> f(a, a)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: ufunc 'ufunc' not supported for the input types, and the
inputs could not be safely coerced to any supported types according
to the casting rule ''safe''
代码语言:javascript复制你可能会问,为什么不之间用@jit装饰器写一个简单的循环来替代呢?
答案是NumPy UFUNC会自动获得其他功能,如 reduce,accumulate或
广播。
a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
>>> f.reduce(a, axis=0) # 按行求和合并
array([12, 15, 18, 21])
>>> f.reduce(a, axis=1) # 按列求和合并
array([ 6, 22, 38])
>>> f.accumulate(a) # 按行积累
array([[ 0, 1, 2, 3],
[ 4, 6, 8, 10],
[12, 15, 18, 21]])
>>> f.accumulate(a, axis=1) # 按列积累
array([[ 0, 1, 3, 6],
[ 4, 9, 15, 22],
[ 8, 17, 27, 38]])