Masked Attention

Compute scaled dot-product attention with a mask that stops queries from attending to future positions. mask[i][j] is True when query $ii$ may attend to key $jj$; for causal (autoregressive) attention it is lower triangular. The output is $\mathrm{softmax}(Q{K}^{\mathrm{\top }}\mathrm{/}\sqrt{d}+\text{mask})V\backslash operatorname\{softmax\}(QK^\backslash top/\backslash sqrt\; d\; +\; \backslash text\{mask\})V$.

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

Returns. Return an array with one output row per query row.

Examples

Example 1

Input:

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

Output:

[[1.0, 2.0], [2.339523, 3.339523]]

Example 2

Input:

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

Output:

[[1.660477, 2.660477], [2.339523, 3.339523]]