Julia 矩阵按行外积的线性组合实现方法咨询

你要实现的运算本质是加权逐行外积的和，公式可以写为：$\sum_{i=1}^n x_i \cdot \boldsymbol{a}_i \boldsymbol{b}_i^T$，其中$\boldsymbol{a}_i$是A的第i行转成的列向量，$\boldsymbol{b}_i$是B的第i行转成的列向量，最终输出是m行p列的矩阵。

你贴的循环写法逻辑完全正确，但当n的数值比较大时，显式循环的性能会很差，下面给出更高效的实现：

Julia 优化实现（和参考代码同技术栈） #

直接用矩阵乘法替换循环，底层调用高度优化的BLAS运算，性能比手写循环高几个数量级，且逻辑等价：

using LinearAlgebra

# 原参数定义
n, m, p = 100, 10, 3
A, B, x = randn(n, m), randn(n, p), randn(n)

# 优化实现，和循环结果完全一致
result_opt = A' * Diagonal(x) * B

# 可以用下面代码验证结果一致性
# result_loop = zeros(m, p)
# for i in 1:n
#     result_loop += x[i] * (A[i, :] * B[i, :]')
# end
# @assert isapprox(result_opt, result_loop)

Python 等价实现（Numpy） #

如果用Python的numpy库，写法逻辑完全一致：

import numpy as np

n, m, p = 100, 10, 3
A = np.random.randn(n, m)
B = np.random.randn(n, p)
x = np.random.randn(n)

# 优化实现
result_opt = A.T @ np.diag(x) @ B

补充说明 #

如果不想构造对角矩阵占用额外内存，也可以用爱因斯坦求和语法实现：

• Julia 可以用Einsum包：@einsum res[i,j] = x[k] * A[k,i] * B[k,j]
• Python Numpy 可以用：np.einsum("ki,k,kj->ij", A, x, B)

内容的提问来源于stack exchange，提问作者Physics_Student