why input is [s b dim] but not [b s dim]?

[Zyriix]

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)

[Doraemonzzz]

This is to align with the input of fairseq.