NumPy 中的自定义矩阵产品，具有用户定义的逐元素操作

Question

我正在寻找一种在 NumPy 中执行自定义矩阵乘积的方法，其中我可以用自定义函数（例如 lambda 表达式）替换逐元素乘法和加法运算。具体来说，我想保留矩阵乘法签名

(n,k),(k,m) -> (n,m)

，但在矩阵乘积计算中使用自定义运算符进行标量乘法和加法。

示例：

给定两个矩阵

A

和

B

：

A = [[a_11, a_12],
     [a_21, a_22]]

B = [[b_11, b_12],
     [b_21, b_22]]

标准矩阵乘积

C = AB

计算如下：

C_ij = sum(A_ik * B_kj for k in range(n))

我想用自定义函数替换乘法 (

*

) 和加法 (

sum

) 运算，例如

custom_multiply

和

custom_add

，这样：

C_ij = custom_add(custom_multiply(A_i1, B_1j), custom_multiply(A_i2, B_2j), ..., custom_multiply(A_ik, B_kj))

自定义运算符：

例如，我可能想定义：

custom_multiply = lambda x, y: (x&0xff) * y  # Custom multiplication function
custom_add = lambda x, y: max(x, y)          # Custom addition function

用例：

我相信为各种应用提供定制矩阵产品是有用的。例如，在网络中，当给定邻接矩阵时，路由的计算可以遵循矩阵乘法的一般过程。但是，必须自定义逐元素操作以符合某些路由协议所需的特定限制。

问题：

我不想使用普通循环来实现矩阵乘积。我的自定义按元素运算可以分解为 NumPy 支持的一系列基本算术运算，并且可以扩展/广播为按矩阵运算。我想知道：

• 有内置的NumPy函数支持自定义矩阵产品吗？
• 如果没有，是否有一种通用方法可以有效地将矩阵乘积过程分解为一系列矩阵运算？

任何指导或示例将不胜感激。

Answer 1

尽管评论给出了一些关于如何实现这一目标的提示，但这里有一个充实的版本。

from typing import Callable

import numpy as np


def munge(A: np.ndarray, B: np.ndarray, mult: Callable, sum_: Callable) -> np.ndarray:
    n, m = A.shape
    m_, k = B.shape
    assert m == m_

    return np.array([[sum_(mult(A[i, :], B[:, j])) for j in range(k)] for i in range(n)])


if __name__ == '__main__':
    A = np.random.randn(3, 4)
    B = np.random.randn(4, 5)

    C = A @ B
    C_ = munge(A, B, np.multiply, np.sum)

    print(np.all(np.isclose(C, C_)))

这会在共享维度上进行广播，但会在（可能）不同的维度上进行迭代，因此它可能不会像标准 matmul 运算那么高效。我没有找到解决这个问题的简单方法。