Scaled Dot-Product Attention

Attention is the core operation of the Transformer:

Implement attention(Q, K, V, causal=False) where:

• Q has shape (seq_q, d_k), K has shape (seq_k, d_k), V has shape (seq_k, d_v)
• If causal=True, position $i$ may only attend to positions $\le i$ (mask future positions with $-\infty$ before softmax). Assume seq_q == seq_k in this case.

Return a tuple (output, weights) where output has shape (seq_q, d_v) and weights has shape (seq_q, seq_k).