Concatenate the heads

Background

After each of the h heads produces an output of shape (n, d_v), multi-head attention concatenates them along the feature axis before the final projection:

$$\text{Concat}(\text{head}_1, \ldots, \text{head}_h) \in \mathbb{R}^{n \times h\,d_v}$$

With d_v = d_model/h, this restores the width to d_model.

Problem statement

Implement concat_heads(head_outputs) returning the concatenation of the head outputs along the last axis.

Input

• head_outputs — list of h arrays, each shape (n, d_v).

Output

Returns an np.ndarray of shape (n, h * d_v).

Examples

Example 1 — two heads, one position

Input:  head_outputs = [[[1, 2]], [[3, 4]]]
Output: [[1, 2, 3, 4]]

Example 2 — three heads, two positions

Input:  head_outputs = [[[1], [2]], [[3], [4]], [[5], [6]]]
Output: [[1, 3, 5], [2, 4, 6]]

Constraints

• Concatenate along the feature axis (axis=1).
• Return shape (n, h * d_v).

Notes

• Concatenation just lines up the heads' outputs side by side; the mixing across heads happens in the next step, the output projection W^O.