创建 NumPy 通用函数

有两种类型的通用函数

  • 操作标量的函数，这些是“通用函数”或 ufuncs（参见下面的 @vectorize）。

  • 操作高维数组和标量的函数，这些是“广义通用函数”或 gufuncs（参见下面的 @guvectorize）。

@vectorize 装饰器

Numba 的 vectorize 允许接受标量输入参数的 Python 函数用作 NumPy ufuncs。创建一个传统的 NumPy ufunc 并非最直接的过程，它涉及编写一些 C 代码。Numba 使这变得容易。使用 vectorize() 装饰器，Numba 可以将纯 Python 函数编译成一个 ufunc，它能像用 C 编写的传统 ufuncs 一样快速地操作 NumPy 数组。

使用 vectorize()，您可以编写函数来操作输入标量，而不是数组。Numba 将生成周围的循环（或 kernel），从而实现对实际输入的有效迭代。

vectorize() 装饰器有两种操作模式

  • 即时（或装饰时）编译：如果您向装饰器传递一个或多个类型签名，您将构建一个 NumPy 通用函数 (ufunc)。本小节的其余部分描述了使用装饰时编译构建 ufuncs。

  • 延迟（或调用时）编译：当未给定任何签名时，装饰器将为您提供一个 Numba 动态通用函数（DUFunc），它在用以前不支持的输入类型调用时动态编译一个新的内核。稍后的一个子章节，“动态通用函数”，将更深入地描述这种模式。

如上所述，如果您将签名列表传递给 vectorize() 装饰器，您的函数将被编译成一个 NumPy ufunc。在基本情况下，只会传递一个签名

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_one_signature
1from numba import vectorize, float64
2
3@vectorize([float64(float64, float64)])
4def f(x, y):
5    return x + y

如果您传递多个签名，请注意必须在最不具体的签名之前传递最具体的签名（例如，单精度浮点数在双精度浮点数之前），否则基于类型的调度将无法按预期工作

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_multiple_signatures
1from numba import vectorize, int32, int64, float32, float64
2import numpy as np
3
4@vectorize([int32(int32, int32),
5            int64(int64, int64),
6            float32(float32, float32),
7            float64(float64, float64)])
8def f(x, y):
9    return x + y

该函数将按预期在指定的数组类型上工作

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_multiple_signatures
1a = np.arange(6)
2result = f(a, a)
3# result == array([ 0,  2,  4,  6,  8, 10])
来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_multiple_signatures
1a = np.linspace(0, 1, 6)
2result = f(a, a)
3# Now, result == array([0. , 0.4, 0.8, 1.2, 1.6, 2. ])

但它将在其他类型上失败

>>> 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''

您可能会问自己，“为什么我要这样做，而不是使用 @jit 装饰器编译一个简单的迭代循环？” 答案是 NumPy ufuncs 自动获得其他功能，例如归约、累加或广播。使用上面的例子

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_multiple_signatures
 1a = np.arange(12).reshape(3, 4)
 2# a == array([[ 0,  1,  2,  3],
 3#             [ 4,  5,  6,  7],
 4#             [ 8,  9, 10, 11]])
 5
 6result1 = f.reduce(a, axis=0)
 7# result1 == array([12, 15, 18, 21])
 8
 9result2 = f.reduce(a, axis=1)
10# result2 == array([ 6, 22, 38])
11
12result3 = f.accumulate(a)
13# result3 == array([[ 0,  1,  2,  3],
14#                   [ 4,  6,  8, 10],
15#                   [12, 15, 18, 21]])
16
17result4 = f.accumulate(a, axis=1)
18# result3 == array([[ 0,  1,  3,  6],
19#                   [ 4,  9, 15, 22],
20#                   [ 8, 17, 27, 38]])

另请参阅

ufuncs 的标准功能 (NumPy 文档)。

注意

在编译代码中仅支持 ufuncs 的广播和归约功能。

vectorize() 装饰器支持多种 ufunc 目标

目标

描述

cpu

单线程 CPU

parallel

多核 CPU

cuda

CUDA GPU

注意

这将创建一个 ufunc-like 对象。详情请参阅 CUDA ufunc 的文档。

一般的指导原则是为不同的数据大小和算法选择不同的目标。“cpu”目标适用于小数据量（约小于 1KB）和低计算强度算法。它的开销最小。“parallel”目标适用于中等数据量（约小于 1MB）。线程会增加少量延迟。“cuda”目标适用于大数据量（约大于 1MB）和高计算强度算法。在 GPU 之间传输内存会增加显著开销。

从 Numba 0.59 开始，cpu 目标在编译代码中支持以下属性和方法

  • ufunc.nin

  • ufunc.nout

  • ufunc.nargs

  • ufunc.identity

  • ufunc.signature

  • ufunc.reduce() (仅前 5 个参数 - 实验性功能)

@guvectorize 装饰器

虽然 vectorize() 允许您一次处理一个元素的 ufuncs，但 guvectorize() 装饰器将概念更进一步，允许您编写可在任意数量的输入数组元素上操作，并接受和返回不同维度的数组的 ufuncs。典型的例子是移动中值或卷积滤波器。

与 vectorize() 函数相反，guvectorize() 函数不返回它们的结果值：它们将其作为数组参数，必须由函数填充。这是因为该数组实际上是由 NumPy 的调度机制分配的，该机制调用 Numba 生成的代码。

与 vectorize() 装饰器类似，guvectorize() 也有两种操作模式：即时（或装饰时）编译和延迟（或调用时）编译。

这是一个非常简单的例子

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize
1from numba import guvectorize, int64
2import numpy as np
3
4@guvectorize([(int64[:], int64, int64[:])], '(n),()->(n)')
5def g(x, y, res):
6    for i in range(x.shape[0]):
7        res[i] = x[i] + y

底层的 Python 函数只是将给定的标量（y）添加到一维数组的所有元素中。更有趣的是声明。其中有两点

  • 输入和输出 布局 的符号形式声明：(n),()->(n) 告诉 NumPy，该函数接受一个 n 元素的一维数组，一个标量（符号上用空元组 () 表示）并返回一个 n 元素的一维数组；

  • 根据 @vectorize 支持的具体 签名 列表；在这里，如上述示例中，我们演示了 int64 数组。

注意

一维数组类型也可以接收标量参数（形状为 () 的那些）。在上面的示例中，第二个参数也可以声明为 int64[:]。在这种情况下，该值必须通过 y[0] 读取。

我们现在可以检查编译后的 ufunc 在一个简单示例中的作用

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize
1a = np.arange(5)
2result = g(a, 2)
3# result == array([2, 3, 4, 5, 6])

好的地方是 NumPy 将根据其形状自动调度更复杂的输入

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize
 1a = np.arange(6).reshape(2, 3)
 2# a == array([[0, 1, 2],
 3#             [3, 4, 5]])
 4
 5result1 = g(a, 10)
 6# result1 == array([[10, 11, 12],
 7#                   [13, 14, 15]])
 8
 9result2 = g(a, np.array([10, 20]))
10g(a, np.array([10, 20]))
11# result2 == array([[10, 11, 12],
12#                   [23, 24, 25]])

注意

无论 vectorize() 还是 guvectorize() 都支持传递 nopython=True ，如 @jit 装饰器中所示。使用它来确保生成的代码不会回退到对象模式。

标量返回值

现在假设我们想从 guvectorize() 返回一个标量值。要做到这一点，我们需要

  • 在签名中，用 [:] 声明标量返回值，就像一维数组一样（例如 int64[:]），

  • 在布局中，将其声明为 ()，

  • 在实现中，写入第一个元素（例如 res[0] = acc）。

以下示例函数计算一维数组（x）与标量（y）的和，并将其作为标量返回

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_scalar_return
1from numba import guvectorize, int64
2import numpy as np
3
4@guvectorize([(int64[:], int64, int64[:])], '(n),()->()')
5def g(x, y, res):
6    acc = 0
7    for i in range(x.shape[0]):
8        acc += x[i] + y
9    res[0] = acc

现在，如果我们将封装的函数应用于数组，我们将得到一个标量值作为输出

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_scalar_return
1a = np.arange(5)
2result = g(a, 2)
3# At this point, result == 20.

覆盖输入值

在大多数情况下，写入输入似乎也有效——然而，这种行为不可靠。考虑以下示例函数

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_overwrite
1from numba import guvectorize, float64
2import numpy as np
3
4@guvectorize([(float64[:], float64[:])], '()->()')
5def init_values(invals, outvals):
6    invals[0] = 6.5
7    outvals[0] = 4.2

调用 init_values 函数并传入 float64 类型的数组会导致输入发生可见变化

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_overwrite
1invals = np.zeros(shape=(3, 3), dtype=np.float64)
2# invals == array([[6.5, 6.5, 6.5],
3#                  [6.5, 6.5, 6.5],
4#                  [6.5, 6.5, 6.5]])
5
6outvals = init_values(invals)
7# outvals == array([[4.2, 4.2, 4.2],
8#                   [4.2, 4.2, 4.2],
9#                   [4.2, 4.2, 4.2]])

这之所以有效，是因为 NumPy 可以将输入数据直接传递给 init_values 函数，因为数据 dtype 与声明的参数类型匹配。然而，它也可能创建并传入一个临时数组，在这种情况下，对输入的更改将丢失。例如，当需要类型转换时，可能会发生这种情况。为了演示，我们可以对 init_values 函数使用 float32 类型的数组

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_overwrite
 1invals = np.zeros(shape=(3, 3), dtype=np.float32)
 2# invals == array([[0., 0., 0.],
 3#                  [0., 0., 0.],
 4#                  [0., 0., 0.]], dtype=float32)
 5outvals = init_values(invals)
 6# outvals == array([[4.2, 4.2, 4.2],
 7#                   [4.2, 4.2, 4.2],
 8#                   [4.2, 4.2, 4.2]])
 9print(invals)
10# invals == array([[0., 0., 0.],
11#                  [0., 0., 0.],
12#                  [0., 0., 0.]], dtype=float32)

在这种情况下，invals 数组没有变化，因为被修改的是临时类型转换后的数组。

为了解决这个问题，需要告知 GUFunc 引擎 invals 参数是可写的。这可以通过向 @guvectorize 传递 writable_args=('invals',)（按名称指定）或 writable_args=(0,)（按位置指定）来实现。现在，上面的代码可以按预期工作了

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_overwrite
 1@guvectorize(
 2    [(float64[:], float64[:])],
 3    '()->()',
 4    writable_args=('invals',)
 5)
 6def init_values(invals, outvals):
 7    invals[0] = 6.5
 8    outvals[0] = 4.2
 9
10invals = np.zeros(shape=(3, 3), dtype=np.float32)
11# invals == array([[0., 0., 0.],
12#                  [0., 0., 0.],
13#                  [0., 0., 0.]], dtype=float32)
14outvals = init_values(invals)
15# outvals == array([[4.2, 4.2, 4.2],
16#                   [4.2, 4.2, 4.2],
17#                   [4.2, 4.2, 4.2]])
18print(invals)
19# invals == array([[6.5, 6.5, 6.5],
20#                  [6.5, 6.5, 6.5],
21#                  [6.5, 6.5, 6.5]], dtype=float32)

动态通用函数

如上所述，如果您没有向 vectorize() 装饰器传递任何签名，您的 Python 函数将用于构建一个动态通用函数，即 DUFunc。例如

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_dynamic
1from numba import vectorize
2
3@vectorize
4def f(x, y):
5    return x * y

结果的 f() 是一个 DUFunc 实例，它最初不支持任何输入类型。当您调用 f() 时，每当您传递以前不支持的输入类型时，Numba 都会生成新的内核。鉴于上述示例，以下解释器交互集说明了动态编译的工作原理

>>> f
<numba._DUFunc 'f'>
>>> f.ufunc
<ufunc 'f'>
>>> f.ufunc.types
[]

上面的例子表明 DUFunc 实例不是 ufuncs。相反，DUFunc 实例通过维护一个 ufunc 成员来工作，然后将 ufunc 属性读取和方法调用委托给此成员（也称为类型聚合）。当我们查看 ufunc 最初支持的类型时，我们可以验证它们不存在。

我们来尝试调用 f()

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_dynamic
1result = f(3,4)
2# result == 12
3
4print(f.types)
5# ['ll->l']

如果这是一个普通的 NumPy ufunc，我们就会看到一个异常，抱怨 ufunc 无法处理输入类型。当我们用整数参数调用 f() 时，我们不仅得到了答案，而且可以验证 Numba 创建了一个支持 C long 整数的循环。

我们可以通过使用不同的输入调用 f() 来添加额外的循环

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_dynamic
1result = f(1.,2.)
2# result == 2.0
3
4print(f.types)
5# ['ll->l', 'dd->d']

现在我们可以验证 Numba 已经为处理浮点输入添加了第二个循环，即 "dd->d"。

如果我们将混合输入类型传递给 f()，我们可以验证 NumPy ufunc 转换规则仍然有效

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_dynamic
1result = f(1,2.)
2# result == 2.0
3
4print(f.types)
5# ['ll->l', 'dd->d']

这个例子表明，用混合类型调用 f() 导致 NumPy 选择浮点循环，并将整数参数转换为浮点值。因此，Numba 没有创建特殊的 "dl->d" 内核。

这种 DUFunc 行为导致我们回到与上面“@vectorize 装饰器”小节中给出的警告类似的问题，但这里重要的不是装饰器中的签名声明顺序，而是调用顺序。如果我们先传递浮点参数，那么任何带有整数参数的调用都将被转换为双精度浮点值。例如

来自 numba/tests/doc_examples/test_examples.py 的 test_vectorize_dynamic
 1@vectorize
 2def g(a, b):
 3    return a / b
 4
 5print(g(2.,3.))
 6# 0.66666666666666663
 7
 8print(g(2,3))
 9# 0.66666666666666663
10
11print(g.types)
12# ['dd->d']

如果您需要精确支持各种类型签名，您应该在 vectorize() 装饰器中指定它们，而不是依赖动态编译。

动态广义通用函数

类似于动态通用函数，如果您没有向 guvectorize() 装饰器指定任何类型，您的 Python 函数将用于构建一个动态广义通用函数，即 GUFunc。例如

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_dynamic
1from numba import guvectorize
2import numpy as np
3
4@guvectorize('(n),()->(n)')
5def g(x, y, res):
6    for i in range(x.shape[0]):
7        res[i] = x[i] + y

我们可以验证结果函数 g() 是一个 GUFunc 实例，它最初不支持任何输入类型。例如

>>> g
<numba._GUFunc 'g'>
>>> g.ufunc
<ufunc 'g'>
>>> g.ufunc.types
[]

类似于 DUFunc，当调用 g() 时，numba 会为以前不支持的输入类型生成新的内核。以下一系列解释器交互将说明 GUFunc 的动态编译如何工作

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_dynamic
1x = np.arange(5, dtype=np.int64)
2y = 10
3res = np.zeros_like(x)
4g(x, y, res)
5# res == array([10, 11, 12, 13, 14])
6print(g.types)
7# ['ll->l']

如果这是一个正常的 guvectorize() 函数，我们就会看到一个异常，抱怨 ufunc 无法处理给定的输入类型。当我们用输入参数调用 g() 时，numba 会为输入类型创建一个新的循环。

我们可以通过使用新参数调用 g() 来添加额外的循环

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_dynamic
1x = np.arange(5, dtype=np.double)
2y = 2.2
3res = np.zeros_like(x)
4g(x, y, res)
5# res == array([2.2, 3.2, 4.2, 5.2, 6.2])

现在我们可以验证 Numba 已经为处理浮点输入添加了第二个循环，即 "dd->d"。

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_dynamic
1print(g.types) # shorthand for g.ufunc.types
2# ['ll->l', 'dd->d']

还可以验证 NumPy ufunc 类型转换规则是否按预期工作

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_dynamic
1x = np.arange(5, dtype=np.int64)
2y = 2
3res = np.zeros_like(x)
4g(x, y, res)
5print(res)
6# res == array([2, 3, 4, 5, 6])

如果您需要精确支持各种类型签名，则不应依赖动态编译，而应将类型作为第一个参数在 guvectorize() 装饰器中指定。

@guvectorize 函数也可以从 JIT 编译的函数中调用。例如

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_jit
 1import numpy as np
 2
 3from numba import jit, guvectorize
 4
 5@guvectorize('(n)->(n)')
 6def copy(x, res):
 7    for i in range(x.shape[0]):
 8        res[i] = x[i]
 9
10@jit(nopython=True)
11def jit_fn(x, res):
12    copy(x, res)

警告

目前尚不支持广播。在需要广播的情况下调用 guvectorize 函数可能会导致不正确的行为。Numba 将尝试检测这些情况并引发异常。

来自 numba/tests/doc_examples/test_examples.py 的 test_guvectorize_jit
 1import numpy as np
 2from numba import jit, guvectorize
 3
 4@guvectorize('(n)->(n)')
 5def copy(x, res):
 6    for i in range(x.shape[0]):
 7        res[i] = x[i]
 8
 9@jit(nopython=True)
10def jit_fn(x, res):
11    copy(x, res)
12
13x = np.ones((1, 5))
14res = np.empty((5,))
15with self.assertRaises(ValueError) as raises:
16    jit_fn(x, res)