numpy einsum 用於沿軸循環的乘法

Question

我有一個 2 x 2 矩陣yy

yy = np.array([[0.5, 0], [0, 2]])
print(yy)

array([[0.5, 0. ],
       [0. , 2. ]])

和 n=3 x 4 x 2 矩陣xy

xy = np.array([
    [[1, 0.1], [2, 0.2], [3, 0.3], [4, 0.4]],
    [[2, 0.2], [4, 0.3], [6, 0.4], [8, 0.5]],
    [[1, 0.3], [3, 0.4], [5, 0.5], [7, 0.6]],
])
print(xy)

array([[[1. , 0.1],
        [2. , 0.2],
        [3. , 0.3],
        [4. , 0.4]],

       [[2. , 0.2],
        [4. , 0.3],
        [6. , 0.4],
        [8. , 0.5]],

       [[1. , 0.3],
        [3. , 0.4],
        [5. , 0.5],
        [7. , 0.6]]])

我想通過numpy.einsum在xy沿 0 軸的轉置元素上乘以yy而無需循環得到

result = np.concatenate((np.expand_dims((yy @ xy[0].T).T, 0), np.expand_dims((yy @ xy[1].T).T, 0), np.expand_dims((yy @ xy[2].T).T, 0)), axis=0)
print(result)

array([[[0.5, 0.2],
        [1. , 0.4],
        [1.5, 0.6],
        [2. , 0.8]],

       [[1. , 0.4],
        [2. , 0.6],
        [3. , 0.8],
        [4. , 1. ]],

       [[0.5, 0.6],
        [1.5, 0.8],
        [2.5, 1. ],
        [3.5, 1.2]]])

也就是說，每個第 i 個元素都應該是

i = 0
print((yy @ xy[i,:,:].T).T)

[[0.5 0.2]
 [1.  0.4]
 [1.5 0.6]
 [2.  0.8]]

如何使用numpy.einsum實現這一目標？

Answer 1

或重寫result ：

 result1 = np.stack([(yy @ xyi.T).T for xyi in xy])

等效的einsum ：

In [32]: np.allclose(np.einsum('ij,kmj->kmi', yy, xy),result)
Out[32]: True

由於xy的k維度是一個“批量”維度，我們可以堅持使用matmul ：

In [34]: np.allclose(xy@yy.T,result)
Out[34]: True

如果我們用更少的轉置編寫迭代版本，這會更明顯：

result2 = np.stack([xyi@yy.T for xyi in xy])