对角矩阵的 3D 张量

Question

我有一个 m 行 n 列的矩阵 A。我想要一个维度为 m*n*n 的 3D 张量，使得张量由 A 的每一列形成的 m 个对角矩阵组成。换句话说，A 的每一列都应该转换为对角化矩阵，所有这些矩阵都应该一起形成一个 3D 张量。

使用 for 循环很容易做到这一点。但我想在不提高速度的情况下这样做。

我想出了一个糟糕且低效的方法，但我希望有人能帮助我找到更好的方法，它允许大型 A 矩阵。

# I use python
# import numpy as np
n = A.shape[0] # A is an n*k matrix
k = A.shape[1]

holding_matrix = np.repeat(np.identity(k), repeats=n, axis=1) # k rows with n*k columns
identity_stack = np.tile(np.identity(n),k) #k nxn identity matrices stacked together

B = np.array((A@holding_matrix)*identity_stack)
B = np.array(np.hsplit(B,k)) # desired result of k n*n diagonal matrices in a tensor

Answer 1

n = A.shape[0] # A.shape == (n, k)
k = A.shape[1]

B = np.zeros_like(A, shape=(k, n*n)) # to preserve dtype and order of A

B[:, ::(n+1)] = A.T
B = B.reshape(k, n, n)