Closed Zyriix closed 1 year ago
this is more general(from my perspective) q,k,v shape : (b,s,d)
q = q.contiguous() q = rearrange(q, 'b n (h d) -> (b h) n d', h = self.num_heads)
k = k.contiguous() k = rearrange(k, 'b n (h d) -> (b h) n d', h = self.num_heads)
v = v.contiguous() v = rearrange(v, 'b n (h d) -> (b h) n d', h = self.num_heads)
This is to align with the input of fairseq.
this is more general(from my perspective) q,k,v shape : (b,s,d)
q = q.contiguous() q = rearrange(q, 'b n (h d) -> (b h) n d', h = self.num_heads)
(N * h, S, d)
k = k.contiguous() k = rearrange(k, 'b n (h d) -> (b h) n d', h = self.num_heads)
(N * h, S, d)
v = v.contiguous() v = rearrange(v, 'b n (h d) -> (b h) n d', h = self.num_heads)