Implement Multi-Head Attention from Tensor Operations

Implement multi-head attention from basic tensor operations and explain the dimensions throughout the computation. Do not substitute a call to a complete built-in multi-head-attention layer for the implementation. The reported exercise does not specify a framework, tensor layout, or attention variant. For this practice version, use batched self-attention with input shape `[batch, sequence, model_dimension]`, an equal head dimension, and learned query, key, value, and output projections. State any additional choices. ### Constraints & Assumptions - `model_dimension` is divisible by the number of heads. - The baseline is the attention module, not a full Transformer block; residual connections and layer normalization are outside it unless requested. - Explain whether a padding or causal mask is supported and what its values mean. - A fully masked query row needs an explicit policy; silently producing undefined softmax values is not acceptable. ### Clarifying Questions to Ask - Is self-attention sufficient, or must query and key/value sequences differ? - Are projection biases, dropout, and causal masking required? - What layout and mask convention does the chosen tensor library use? ```hint Write the shape after each projection, reshape, transpose, and matrix multiplication. Softmax must normalize over the key positions for each query and head. ``` ### What a Strong Answer Covers The implementation should include learned projections, splitting into heads, scaled dot-product scores, correctly applied masks, softmax over keys, weighted values, head recombination, and output projection. Explain numerical stability, shape checks, and what is excluded from the module. ### Follow-up Questions - Why scale scores by the square root of the per-head key dimension? - How would cross-attention change the sequence dimensions? - Which tensor grows quadratically with sequence length in a straightforward implementation?