Masked Multihead Attention

I made a masked Multihead Attention mechanism for LLMs.

This is the code:

They are combined the context vectors in the last dimension each learning something different from the input x.