[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-1)def forward(self, x):
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-2)    x = self.qkv(self.norm(x))
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-3)    x = rearrange(x, 'n s (h d) -> (n h) s d', h=self.nheads)
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-4)    q,k,v = torch.chunk(x, 3, dim=-1)
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-5)    s = (q@k.transpose(1,2))/self.scale
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-6)    x = s.softmax(dim=-1)@v
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-7)    x = rearrange(x, '(n h) s d -> n s (h d)', h=self.nheads)
[](https://www.fast.ai/posts/2023-05-03-mojo-launch.html#cb2-8)    return self.proj(x)