使用 Numba 沿 O 维度将 (N,N) 矩阵乘以 (N,M,O) 矩阵答案

【问题标题】：Multiply a (N,N) matrix by a (N,M,O) matrix along the O dimension with Numba使用 Numba 沿 O 维度将 (N,N) 矩阵乘以 (N,M,O) 矩阵
【发布时间】：2023-01-19 02:12:39
【问题描述】：

我正在尝试将大小为 $(N,N)$ 的矩阵 A 乘以大小为 $(N,M,O)$ 的矩阵 B 沿着 O 维度（即左乘所有“页面” B 沿 O 维度乘以 A)，使用 jitted numba 函数。

我想出了这个解决方案：

@njit
def fast_expectation(Pi, X):
    
    res = np.empty_like(X)
    
    for i in range(Pi.shape[0]):
        for j in range(X.shape[1]):
            for k in range(X.shape[2]):
                res[i,j,k] = np.dot(Pi[i,:], X[:,j,k])
                            
    return res

但是，这会返回警告 NumbaPerformanceWarning: np.dot() is faster on contiguous arrays, called on (array(float64, 1d, C), array(float64, 1d, A))。您知道我如何使用 numba 兼容功能快速执行此操作吗？

我尝试运行之前的代码，并交换矩阵 B 的数组（将其转换为 (N,M,O) 矩阵）。没用。

编辑：

我还尝试了以下代码：

@njit
def multiply_ith_dimension(Pi, i, X):
    """If Pi is a matrix, multiply Pi times the ith dimension of X and return"""
    X = np.swapaxes(X, 0, i)
    shape = X.shape
    X = X.reshape(shape[0], -1)

    # iterate forward using Pi
    X = Pi @ X

    # reverse steps
    X = X.reshape(Pi.shape[0], *shape[1:])
    return np.swapaxes(X, 0, i)

这也给了我一个错误

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
reshape() supports contiguous array only
...
    <source elided>
    shape = X.shape
    X = X.reshape(shape[0], -1)
    ^

【问题讨论】：

您可以尝试对变量 Pi 使用列优先排序吗？例如在循环之前，插入行 Pi = np.asfortranarray(Pi) 以将行优先更改为列优先。不要更改索引。
谢谢！如果我将它应用于 X，它就会起作用！但是，当我通过 Pi (Pi.T) 的转置时，我又得到了之前的错误... NumbaPerformanceWarning: np.dot() is faster on contiguous arrays, called on (array(float64, 1d, A), array(float64, 1d, F)) 你知道如何让它在这两种情况下都起作用吗？

标签： python matrix-multiplication numba

【解决方案1】：

您可以使用 np.einsum（或者更快，opt_einsum 库）轻松完成此操作。

result = np.einsum('ab,bcd->acd', A, B, optimize=True)

请注意，einsum 在后台使用 np.dot（基于 BLAS），因此可能是最快的方法。

【讨论】：