Multi-Head Attention Split

Split the feature dimension of a batch into attention heads. X has shape (batch, time, features); reshape the features into n_heads groups of features // n_heads and move the head axis forward. features must be divisible by a positive n_heads, otherwise ValueError.

Implement solve(X, n_heads).

Returns. Return an array of shape (batch, n_heads, time, features // n_heads).

Examples

Example 1

Input:

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

Output:

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

Example 2

Input:

solve([[[1, 2, 3]]], 2)

Output: Raises ValueError.