Implement Multi-Head Self-Attention in PyTorch

Implement a Multi-Head Self-Attention module from scratch in PyTorch (do not use torch.nn.MultiheadAttention). The implementation must support batched inputs. Requirements Implement a class MultiHeadSelfAttention with: Constructor parameters: d_model: input/output feature dimension num_heads: number of attention heads; require d_model % num_heads == 0 optional dropout (if not required, you may omit) Inputs: x: a tensor of shape [batch_size, seq_len, d_model] optional mask: an attention mask to block certain tokens (padding mask or causal mask; if not required, you can skip) Output: y: a tensor of shape [batch_size, seq_len, d_model] Key Steps Use linear layers to produce Q, K, V (each with last-dim size d_model). Reshape/split into multiple heads: per-head dim d_head = d_model / num_heads transform Q, K, V into shape [batch_size, num_heads, seq_len, d_head] Compute scaled dot-product attention: scores = (Q @ K^T) / sqrt(d_head) attn = softmax(scores, dim=-1) context = attn @ V Concatenate heads back to [batch_size, seq_len, d_model], then apply an output linear layer. Constraints / Scale 1 <= batch_size <= 64 1 <= seq_len <= 512 Typical d_model: 64/128/256/512/768 Ensure tensor reshapes are correct and matmul dimensions match. Sample Checks With x = torch.randn(2, 4, 8) and num_heads=2, output shape should be [2, 4, 8]. Handle or raise error when d_model % num_heads != 0. loss = y.sum(); loss.backward() should work without errors. Example Input (Python) x = torch.randn(2, 4, 8); d_model=8; num_heads=2 Output output tensor shape: torch.Size([2, 4, 8])