Implement Multi-Head Attention (NumPy or PyTorch)

Implement a simplified Multi-Head Attention forward pass using NumPy or PyTorch, and clearly state tensor shapes. Inputs Q, K, V: each of shape (B, T, D_model) number of heads H with D_model % H == 0 optional attention mask M broadcastable to (B, H, T, T), where masked positions are -inf (or a very large negative number) Output O: shape (B, T, D_model) Requirements Split D_model into H heads with D_head = D_model / H Compute scaled dot-product attention: A = softmax((Q K^T) / sqrt(D_head) + M) head_out = A V Concatenate heads back to D_model Explain how you ensure numerical stability (e.g., stable softmax) Scale B <= 8, T <= 2048, D_model <= 4096, H <= 64 No need to implement backward pass or learnable projections (Wq/Wk/Wv/Wo); only the core attention computation. Example Input # PyTorch snippet\n# B=1,T=2,D=4,H=2\n# Q=K=V all ones -> uniform attention, output ones\n Output O shape = (1,2,4); all elements ~= 1.0