如何用Numpy优雅地将多向量对角化生成三维数组？

给定形状为(2,3)的数组：

import numpy as np
data = np.array([[1,2,3],[4,5,6]])

需要将每行的3元素向量转换为3×3的对角矩阵，最终得到形状为(2,3,3)的目标数组，这里提供几种高效优雅的Numpy实现方式：

方法1：使用爱因斯坦求和约定（np.einsum） #

这是最简洁的向量化实现方式，利用np.einsum的批量处理能力直接生成目标数组：

result = np.einsum('ij,jk->ijk', data, np.eye(3))

原理：通过求和约定将每个行向量data[i]与单位矩阵np.eye(3)做矩阵乘法，自动批量生成对应的对角矩阵，最终输出形状为(2,3,3)。

方法2：利用维度扩展+广播乘法 #

通过扩展原数组的维度，结合Numpy的广播机制实现批量对角化：

result = np.expand_dims(data, axis=-1) * np.eye(3)

原理：np.expand_dims(data, axis=-1)将原数组形状从(2,3)变为(2,3,1)，与形状为(3,3)的单位矩阵相乘时，广播机制会自动将每个元素映射到对角位置，生成对应的对角矩阵。

方法3：逐行应用np.diag（适合小数据量） #

如果追求代码直观性，也可以用np.apply_along_axis逐行处理：

result = np.apply_along_axis(np.diag, axis=1, arr=data)

注意：这种方法本质是循环处理每行，数据量大时效率不如前两种向量化方法，仅推荐用于小规模数据场景。

验证输出：上述三种方法得到的结果均为：

array([[[1., 0., 0.],
        [0., 2., 0.],
        [0., 0., 3.]],

       [[4., 0., 0.],
        [0., 5., 0.],
        [0., 0., 6.]]])

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