numba vectorize 的问题：数组不能表示为 numpy 类型

Question

我遇到的基本问题如下：

In [39]: import numba as nb

In [40]: import numpy as np

In [41]: mat = np.array([[1.2, 1.6, 0.5, 0  ],
    ...:                 [1.2, 1.6, 0,   0.5]])

In [42]: @nb.vectorize
    ...: def f(a, b, c, d):
    ...:     load = np.array([a, b, c, d])
    ...:     return mat @ load
    ...:

In [43]: f(1.1, 2.1, 3, 4)
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<SNIP>
NotImplementedError: array(float64, 1d, C) cannot be represented as a Numpy dtype

旁注：在此示例中，mat 表示负载因子矩阵（要乘以一组负载）并且永远不会改变；在我的实际问题中，我动态生成了许多这样的mats 并动态生成了许多这些分解的负载组合函数（每个函数都有自己的常量mat 值）。

现在，显然这个函数更容易实现为基本的矩阵乘法运算，而无需使用 numba：

In [44]: def f(a, b, c, d):
    ...:     load = np.array([a, b, c, d])
    ...:     return mat @ load
    ...:

In [45]: f(1.1, 2.1, 3, 4)
Out[45]: array([6.18, 6.68])

但是，我使用 numba 的目标是允许一起广播各种形状的数组的参数，而不仅仅是原子值。示例：

>>> a = np.array([1, 2, 3])
>>> b = np.array([[1, 2, 3], [4, 5, 6]])
>>> c = 1
>>> d = 2
>>> f(a, b, c, d)
# expected output:
array([[[ 3.3,  3.8],
        [ 6.1,  6.6],
        [ 8.9,  9.4]],

       [[ 8.1,  8.6],
        [10.9, 11.4],
        [13.7, 14.2]]])

如何对每个操作生成一个 numpy 数组的操作进行向量化，并返回与输入的广播形状匹配的结果嵌套矩阵？

编辑：

顺便说一句，我对输出形状不是很固执。输出这样的形状可能会更方便？

>>> f(a, b, c, d)
# expected output:
array([[[ 3.3,  6.1,  8.9],
        [ 8.1, 10.9, 13.7]],

       [[ 3.8,  6.6,  9.4],
        [ 8.6, 11.4, 14.2]]])

编辑：

这是我使用nb.guvectorize 所能做出的最好尝试。它会给出性能警告和错误的结果。

In [97]: @nb.guvectorize(["f8, f8, f8, f8, f8[:,:], f8[::1]"], "(), (), (), (), (n,m)-> (n)")
    ...: def f(a, b, c, d, mat, res):
    ...:     for i in range(mat.shape[0]):
    ...:         res[i] = mat[i] @ np.array([a, b, c, d])
    ...:
<ipython-input-97-9dea8a9aed86>:4: NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (array(float64, 1d, A), array(float64, 1d, C))
  res[i] = mat[i] @ np.array([a, b, c, d])
<SNIP>

In [98]: f(a, b, c, d, mat)
Out[98]:
array([[[ 3.3,  3.3],
        [ 6.1,  6.1],
        [ 8.9,  8.9]],

       [[ 8.1,  8.1],
        [10.9, 10.9],
        [13.7, 13.7]]])

Answer 1

一种方法是一次将mat 的每一行向量化。这将需要为每一行创建一个单独的向量化函数，并且还需要取消矩阵乘法运算以支持基本算术。它看起来像这样：

# note that `mat` is of this form:
"""
        a    b    c    d
array([[1.2, 1.6, 0.5, 0. ],
       [1.2, 1.6, 0. , 0.5]])
"""

# So we create two functions to perform the matrix multiplication operation of each row separately, f_c and f_d:

@nb.vectorize
def f_c(a, b, c):
    """Load factor for d is zero."""
    return 1.2*a + 1.6*b + 0.5*c


@nb.vectorize
def f_d(a, b, d):
    """Load factor for c is zero."""
    return 1.2*a + 1.6*b + 0.5*d


# now call both functions and pack the results into a matrix:


def f(a, b, c, d):
    return np.array([f_c(a, b, c), f_d(a, b, d)])

这行得通。时机似乎也很不错（如果用@nb.njit jit-ed 会更好，但这会破坏广播）：

In [94]: %timeit f(1.1, 2.1, 3.1, 4.1)
2.7 µs ± 9.58 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

将此与简单的矩阵乘法版本进行比较：

def g(a, b, c, d):
    return mat @ np.array([a, b, c, d])

...速度几乎是原来的两倍：

In [98]: %timeit g(1.1, 2.1, 3.1, 4.1)
1.34 µs ± 1.84 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

...但是这个版本不适用于广播数组。

为了支持广播，我认为支付 2 倍减速的成本并不算太糟糕；至少不是 10 倍。但是，我不喜欢既要创建一堆编译函数又要取消快速矩阵乘法的想法。这种“多功能”方法也没有尝试利用并行处理。

这是另一个使用 guvectorize 的实现，它看起来比上面的方法更简单，但速度较慢并且会产生性能警告：

In [69]: @nb.guvectorize(["f8, f8, f8, f8, f8[:,:], f8[::1]"], "(), (), (), (), (m,n) -> (m)")
    ...: def f(a, b, c, d, m, res):
    ...:     res[:] = m @ np.array([a, b, c, d])
    ...: 
<ipython-input-69-fe24a895fcf3>:3: NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (array(float64, 2d, A), array(float64, 1d, C))
  res[:] = m @ np.array([a, b, c, d])
<SNIP>

In [70]: %timeit f(1.1, 2.1, 3.1, 4.1, mat)
14.5 µs ± 93.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

必须有更好的方法。