import numpy as np

A = np.array(
    [
        [1, 1, 2],
        [2, 2, 2],
        [5, 5, 5]
    ]
)
B = np.array(
    [
        [0, 1, 0],
        [1, 1, 0],
        [1, 1, 1]
    ]
)

A[..., np.newaxis]
'''
array([[[1],
        [1],
        [2]],

       [[2],
        [2],
        [2]],

       [[5],
        [5],
        [5]]])
'''

# 向量乘法不求积，对A[:, :, np.newaxis]或A[..., np.newaxis]来说，之前的每行(1,3)，变成三行->(3,1)，重复三次
# 那么现在的每份(1,3,3)，对应的就是之前(1,3),转换为列后再重复三次，就可以实现列和列的对应位置的乘积
result_1 = A[:, :, np.newaxis] * B[np.newaxis, :, :]
result_2 = np.einsum('ij,jk->ijk', A, B) # 爱因斯坦求和
# 实现矩阵乘法不求和，以及求和后和矩阵乘法一致
np.array_equal(result_1, result_2), np.array_equal(A @ B, result_2.sum(axis=1))
# (True, True)

A = np.array(
    [
        [1, 1, 2],
        [2, 2, 2],
        [5, 5, 5]
    ]
)
B = np.array(
    [
        [0, 1, 0],
        [1, 1, 0],
        [1, 1, 1]
    ]
)

score1 = A @ B.T  # A和B都是(token,dim)的矩阵，转置则有(dim,token),计算内积来获得注意力分数
score2 = np.einsum("ij,dj->id", A, B) # 爱因斯坦积的表示
np.array_equal(score1, score2) # True

def forward(self, x):
    b, c, h, w = x.shape
    # 划分q,k,v
    qkv = self.to_qkv(x).chunk(3, dim=1)
    # rearrange
    q, k, v = map(
        lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
    )
    q = q * self.scale
  # 相似度计算
    sim = einsum("b h d i, b h d j -> b h i j", q, k)
    sim = sim - sim.amax(dim=-1, keepdim=True).detach()
    attn = sim.softmax(dim=-1)
  # 相似度计算
    out = einsum("b h i j, b h d j -> b h i d", attn, v)
    # rearrange
    out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
    return self.to_out(out)