Multi-Head Attention Merge

Merge attention heads back into the model dimension: the inverse of the head split. heads has shape (batch, n_heads, time, width).

Implement solve(heads).

Returns. Return an array of shape (batch, time, n_heads * width) in which the heads of each token are concatenated in head order.

Examples

Example 1

Input:

solve([[[[1, 2], [5, 6]], [[3, 4], [7, 8]]]])

Output:

[[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]

Example 2

Input:

solve([[[[1.0]], [[2.0]], [[3.0]]]])

Output:

[[[1.0, 2.0, 3.0]]]