KV Cache Append

Append the key/value vector of the newest token to a KV cache. The cache is a 2-D array of shape (seq_len, dim) and new_value is a vector of length dim; the new row is added at the end of the sequence axis (axis $-2-2$).

Implement solve(cache, new_value).

Returns. Return an array of shape (seq_len + 1, dim).

Examples

Example 1

Input:

solve([[1.0, 2.0], [3.0, 4.0]], np.array([5.0, 6.0]))

Output:

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

Example 2

Input:

solve(np.zeros((0, 2)), np.array([1.0, 2.0]))

Output:

[[1.0, 2.0]]