Implement Single-Head and Multi-Head Attention in NumPy

Implement single-head and multi-head attention from scratch using NumPy (no deep learning frameworks). Part A: Single-Head Scaled Dot-Product Attention Implement: single_head_attention(Q, K, V) Input Q: array of shape (Tq, d) K: array of shape (Tk, d) V: array of shape (Tk, dv) Output O: array of shape (Tq, dv) defined as: scores = (Q @ K.T) / sqrt(d) weights = softmax(scores, axis=-1) O = weights @ V Assume softmax is provided in starter code. Part B: Multi-Head Attention Implement: multi_head_attention(X, Wq, Wk, Wv, Wo, num_heads) Input X: shape (T, d_model) Wq, Wk, Wv: shape (d_model, d_model) Wo: shape (d_model, d_model) num_heads: number of heads h, with d_model % h == 0 Definition Q = X @ Wq, K = X @ Wk, V = X @ Wv (all (T, d_model)) Reshape/split into h heads with d_head = d_model / h. Run Part A attention per head to get (T, d_head) per head. Concatenate heads back to (T, d_model). Output Y = concat_heads @ Wo of shape (T, d_model). Constraints NumPy only. Must pass provided unit tests; no need for performance optimizations. Example checks (illustrative) Shape checks. If num_heads=1 and all weights are identity, multi-head should match the single-head definition (when dimensions align). If d_model % num_heads != 0, raise an error or handle explicitly. Example Input Q shape (2,4), K shape (3,4), V shape (3,5) Output O shape (2,5)