Scaled Dot-Product Attention

Compute scaled dot-product attention $\mathrm{softmax}\text{\hspace{0.17em}⁣}\left(Q{K}^{\mathrm{\top }}\mathrm{/}\sqrt{d}\right)V\backslash operatorname\{softmax\}\backslash !\backslash big(QK^\backslash top/\backslash sqrt\{d\}\backslash big)V$ with an optional boolean mask (same shape as the score matrix): positions where the mask is False are blocked, so they get (almost) zero attention weight.

Implement solve(Q,K,V,mask=None).

Returns. Return an array with one output row per query row. Rows of Q are queries (length $dd$), rows of K are keys, rows of V are values.

Examples

Example 1

Input:

solve([[1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[1.0, 2.0], [3.0, 4.0]], None)

Output:

[[1.660477, 2.660477]]

Example 2

Input:

solve([[1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[1.0, 2.0], [3.0, 4.0]], [[True, False]])

Output:

[[1.0, 2.0]]