如何向量化扁平化外积函数，使其适用于矩阵的每对行？

你的原函数是针对两个一维数组计算外积并扁平化，要扩展到矩阵的每一对行，完全可以用NumPy的广播机制或einsum实现纯向量化的高效处理，不需要显式循环。

先优化原一维函数（可选） #

原函数里的np.resize可以替换为np.reshape，因为我们只是改变数组形状，resize在形状不匹配时会重复元素，用reshape更安全：

import numpy as np

def cross(a, b):
    return (a.reshape(-1, 1) * b.reshape(1, -1)).flatten()

这个版本和原函数效果完全一致，但更健壮。

矩阵行对的向量化实现 #

假设你有两个矩阵：

• A：形状为(N, D1)，共N行，每行D1个元素
• B：形状为(M, D2)，共M行，每行D2个元素

我们需要对A的每一行和B的每一行计算扁平化外积，最终得到形状为(N, M, D1*D2)的结果（每个位置对应一对行的计算结果）。

方法1：广播+Reshape（直观易懂）

def cross_matrix_rows(A, B):
    N, D1 = A.shape
    M, D2 = B.shape
    # 扩展维度实现广播：A的每一行对应(N,1,D1,1)，B的每一行对应(1,M,1,D2)
    # 相乘后得到所有行对的外积矩阵，形状为(N,M,D1,D2)
    outer_pairs = A[:, None, :, None] * B[None, :, None, :]
    # 扁平化最后两个维度，得到每个行对的扁平化外积
    return outer_pairs.reshape(N, M, D1 * D2)

方法2：用np.einsum（代码更简洁）

einsum可以清晰描述维度间的运算关系，适合这种多维度的行对操作：

def cross_matrix_rows(A, B):
    # 计算所有行对的外积，得到形状为(N,M,D1,D2)的结果
    outer_pairs = np.einsum('nd,me->nmed', A, B)
    # 扁平化最后两个维度
    return outer_pairs.reshape(A.shape[0], B.shape[0], -1)

测试验证 #

# 测试用例
A = np.array([[1,2], [3,4]])  # 2行2列
B = np.array([[5,6], [7,8], [9,10]])  # 3行2列

result = cross_matrix_rows(A, B)
print(result.shape)  # 输出 (2, 3, 4)，对应2×3个行对，每个行对的结果是4个元素

# 验证第一个行对的结果是否正确
expected = np.outer(A[0], B[0]).flatten()
print(np.allclose(result[0, 0], expected))  # 输出 True

为什么这是向量化？ #

整个过程没有Python层面的循环，完全依赖NumPy底层的C实现处理维度运算，在矩阵规模较大时，效率比显式循环遍历行对高几个数量级。

内容的提问来源于stack exchange，提问作者W.314