如何沿高維數組的對角線分布 Numpy 數組？

Question

我有三個二維 Numpy arrays x ， w ， d並想創建第四個稱為a 。 w和d僅用d.shape + w.shape定義a的形狀。 我想在 a 的條目中a x在其他地方有一個零。 具體來說，我想要這個代碼的無循環版本：

a = np.zeros(d.shape + w.shape)

for j in range(d.shape[1]):
    a[:,j,:,j] = x

例如，給定：

x = np.array([
    [2, 3],
    [1, 1],
    [8,10],
    [0, 1]
])

w = np.array([
    [ 0, 1, 1],
    [-1,-2, 1]
])

d = np.matmul(x,w)

我想a

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

        [[ 0.,  2.,  0.],
         [ 0.,  3.,  0.]],

        [[ 0.,  0.,  2.],
         [ 0.,  0.,  3.]]],


       [[[ 1.,  0.,  0.],
         [ 1.,  0.,  0.]],

        [[ 0.,  1.,  0.],
         [ 0.,  1.,  0.]],

        [[ 0.,  0.,  1.],
         [ 0.,  0.,  1.]]],


       [[[ 8.,  0.,  0.],
         [10.,  0.,  0.]],

        [[ 0.,  8.,  0.],
         [ 0., 10.,  0.]],

        [[ 0.,  0.,  8.],
         [ 0.,  0., 10.]]],


       [[[ 0.,  0.,  0.],
         [ 1.,  0.,  0.]],

        [[ 0.,  0.,  0.],
         [ 0.,  1.,  0.]],

        [[ 0.,  0.,  0.],
         [ 0.,  0.,  1.]]]])

Answer 1

這個答案啟發了以下解決方案：

# shape a: (4, 3, 2, 3)
# shape x: (4, 2)

a = np.zeros(d.shape + w.shape)
a[:, np.arange(a.shape[1]), :, np.arange(a.shape[3])] = x

它使用 Numpy 的廣播（參見此處或 此處）與Advanced Indexing結合來放大x以適應切片。

Answer 2

我碰巧有一個更簡單的解決方案： a = np.tensordot(x, np.identity(3), axes = 0).swapaxes(1,2)
單位矩陣的大小將取決於您希望重復x元素的次數。