如何在NumPy中将二维数组行平方转换为三维数组？

云原生阿岚

2026-8-5

高效实现矩阵行的外积运算（用于Softmax导数批量计算）

核心需求

给定形状为(n_samples, d)的二维NumPy矩阵，对每一行计算外积（即row.reshape(-1, 1) @ row.T），得到形状为(n_samples, d, d)的三维数组，避免低效的Python循环。

解决方案1：利用NumPy广播机制

通过维度扩展实现广播相乘，直接得到批量外积结果：

import numpy as np

a = np.array([[1,2,3],
              [4,5,6]])

# 扩展维度后广播相乘
outer_product = a[:, :, np.newaxis] * a[:, np.newaxis, :]
print(outer_product)

输出结果：

array([[[ 1,  2,  3],
        [ 2,  4,  6],
        [ 3,  6,  9]],

       [[16, 20, 24],
        [20, 25, 30],
        [24, 30, 36]]])

原理说明：

a[:, :, np.newaxis]将原数组从(2,3)转换为(2,3,1)
a[:, np.newaxis, :]将原数组转换为(2,1,3)
NumPy广播机制会自动扩展维度，对应位置元素相乘，最终得到每行的外积结果，与row.reshape(-1,1) @ row.T完全等价。

解决方案2：使用`np.einsum`（更简洁）

利用爱因斯坦求和约定，代码更直观：

outer_product = np.einsum('ni,nj->nij', a, a)

该方式同样会生成(n_samples, d, d)的三维数组，结果与广播方法一致，且底层优化效率很高。

结合Softmax导数批量实现

结合你已实现的对角三维数组构造方法，可直接完成批量Softmax导数计算：

def batch_softmax_derivative(s_matrix):
    d = s_matrix.shape[-1]
    identity = np.eye(d)
    # 构造批量对角矩阵
    diag_part = s_matrix[..., np.newaxis] * identity
    # 计算批量外积
    outer_product = s_matrix[:, :, np.newaxis] * s_matrix[:, np.newaxis, :]
    # 返回Softmax导数：diag(s) - s@s.T
    return diag_part - outer_product

其中s_matrix是批量的Softmax输出（形状(n_samples, d)），函数返回每个样本的Softmax导数矩阵（形状(n_samples, d, d)）。