CosineSimCodebook疑问求解答

[guoswang]

在您开源的代码中，CosineSimCodebook实现部分。
我看首先会用下面代码对emb进行初始化。
raw_flatten = rearrange(x, 'h ... d -> h (...) d') flatten = l2norm(raw_flatten)

然后再通过余弦相似度得到embed_ind
embed = l2norm(embed) dist = einsum('h n d, h c d -> h n c', flatten, embed) embed_ind = gumbel_sample(dist, dim=-1, temperature=self.sample_codebook_temp)
上述操作，是不是可以理解在计算相似度约等价为
dist = einsum('h n d, h c d -> h n c', l2norm(raw_flatten), l2norm(l2norm(raw_flatten)))
但感觉这个计算逻辑怪怪的，还是我理解有偏差，请大佬帮忙解答一下。