如何在Python（cuPy/Numpy）中实现行与对应列的相乘？

由于NumPy和CuPy的API几乎完全一致，以下代码可直接在两个库中切换使用（将np替换为cp即可），全程无需for循环。

场景1：计算行与对应列的点积（得到标量数组） #

即对每个索引i，计算矩阵第i行与转置矩阵第i列（即原矩阵第i行的列向量形式）的点积（元素相乘后求和），推荐两种高效实现方式：

• 逐元素相乘后按行求和
  这是性能最优的方式，仅需一次逐元素运算和行维度求和：

  import numpy as np  # 替换为import cupy as cp即可适配CuPy
  
  mat = np.array([[1,2,3], [4,5,6], [7,8,9]])
  result = (mat * mat).sum(axis=1)
  # 输出：array([14, 77, 194])，对应各行元素的平方和
  

• 使用Einstein求和
  通过索引表达式清晰描述运算逻辑，性能与上一方法持平：

  result = np.einsum('ij,ij->i', mat, mat)
  # 或明确标注转置关系：np.einsum('ij,ji->i', mat, mat.T)
  

注意：若用np.diag(mat @ mat.T)实现，会额外计算大量冗余元素，大矩阵场景下效率较低，不推荐。

场景2：计算行与对应列的逐元素相乘（得到同维度矩阵） #

若需求是将每行与对应列（即自身）逐元素相乘，直接用逐元素乘法或平方函数即可：

# 逐元素相乘
result = mat * mat
# 或用平方函数，运算效率更高
result = np.square(mat)
# 输出：
# array([[ 1,  4,  9],
#        [16, 25, 36],
#        [49, 64, 81]])

场景3：计算行与对应列的外积（得到三维张量） #

若需对每个i，计算第i行与对应列的外积（得到二维矩阵），最终组成三维张量，可通过广播实现：

result = mat[:, :, np.newaxis] * mat[:, np.newaxis, :]
# 3x3矩阵的输出为3x3x3张量，其中result[i]即为第i行与对应列的外积

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