在numpy中沿一个轴对角化二维矩阵

Diagonalize 2d matrix along one axis in numpy

给定维度为 MxN 的矩阵 X。我想为 X 的每一行创建一个对角矩阵。结果应该在 MxNxN 中。如何有效地做到这一点?谢谢!

out = np.zeros((m, n, n))
out[:, np.arange(n), np.arange(n)] = X