Compute scaled dot-product attention for a batch of examples with several heads.
For one head of one example, the queries Q have shape (Tq, D), the keys K have shape (Tk, D) and the values V have shape (Tk, Dv), and the output is
Attention(Q,K,V)=softmax(DQK⊤)V,where the softmax runs over the Tk keys separately for each query row. The input q has shape (B, H, Tq, D), k has shape (B, H, Tk, D) and v has shape (B, H, Tk, Dv), where B is the batch size, H is the number of heads, Tq is the number of queries, Tk is the number of keys, D is the shared query and key width and Dv is the value width. Example b and head h use the blocks q[b, h], k[b, h] and v[b, h].
Return a JAX array of shape (B, H, Tq, Dv) with the dtype of q.
Input: q = [[[[1,0],[0,1]],[[1,1],[0,0]]]], k = [[[[1,0],[0,1]],[[1,0],[0,1]]]], v = [[[[1,2],[3,4]],[[0,1],[1,0]]]]
Output: [[[[1.660477,2.660477],[2.339523,3.339523]],[[0.5,0.5],[0.5,0.5]]]]
Explanation: In head 0, the query (1, 0) scores 0.7071 against the first key and 0 against the second, so its weights are 0.6698 and 0.3302 and its output is 0.6698 times (1, 2) plus 0.3302 times (3, 4). In head 1 every query scores the two keys equally, so every output is the plain average (0.5, 0.5) of the two value rows.
Input: q = [[[[2,0]]],[[[0,-1]]]], k = [[[[1,0],[0,1],[-1,0]]],[[[1,0],[0,1],[-1,0]]]], v = [[[[1,0],[0,1],[1,1]]],[[[2,0],[0,2],[2,2]]]]
Output: [[[[0.813306,0.232082]]],[[[1.604448,1.197776]]]]
Input: q = [[[[1,1,0,0],[0,0,1,-1]]]], k = [[[[1,0,1,0],[0,1,0,1],[1,1,1,1]]]], v = [[[[3],[-1],[0.5]]]]
Output: [[[[0.774069],[1.486715]]]]
jax.vmap(f) maps over axis 0, and jax.vmap(jax.vmap(f)) maps over two leading axes.
jax.nn.softmax(scores, axis=-1) normalizes over the keys.
Divide by jnp.sqrt(d).astype(q.dtype) to keep bfloat16 inputs in bfloat16.
Accepts: array
Accepts: array
Accepts: array
Compute scaled dot-product attention for a batch of examples with several heads.
For one head of one example, the queries Q have shape (Tq, D), the keys K have shape (Tk, D) and the values V have shape (Tk, Dv), and the output is
Attention(Q,K,V)=softmax(DQK⊤)V,where the softmax runs over the Tk keys separately for each query row. The input q has shape (B, H, Tq, D), k has shape (B, H, Tk, D) and v has shape (B, H, Tk, Dv), where B is the batch size, H is the number of heads, Tq is the number of queries, Tk is the number of keys, D is the shared query and key width and Dv is the value width. Example b and head h use the blocks q[b, h], k[b, h] and v[b, h].
Return a JAX array of shape (B, H, Tq, Dv) with the dtype of q.
Input: q = [[[[1,0],[0,1]],[[1,1],[0,0]]]], k = [[[[1,0],[0,1]],[[1,0],[0,1]]]], v = [[[[1,2],[3,4]],[[0,1],[1,0]]]]
Output: [[[[1.660477,2.660477],[2.339523,3.339523]],[[0.5,0.5],[0.5,0.5]]]]
Explanation: In head 0, the query (1, 0) scores 0.7071 against the first key and 0 against the second, so its weights are 0.6698 and 0.3302 and its output is 0.6698 times (1, 2) plus 0.3302 times (3, 4). In head 1 every query scores the two keys equally, so every output is the plain average (0.5, 0.5) of the two value rows.
Input: q = [[[[2,0]]],[[[0,-1]]]], k = [[[[1,0],[0,1],[-1,0]]],[[[1,0],[0,1],[-1,0]]]], v = [[[[1,0],[0,1],[1,1]]],[[[2,0],[0,2],[2,2]]]]
Output: [[[[0.813306,0.232082]]],[[[1.604448,1.197776]]]]
Input: q = [[[[1,1,0,0],[0,0,1,-1]]]], k = [[[[1,0,1,0],[0,1,0,1],[1,1,1,1]]]], v = [[[[3],[-1],[0.5]]]]
Output: [[[[0.774069],[1.486715]]]]
jax.vmap(f) maps over axis 0, and jax.vmap(jax.vmap(f)) maps over two leading axes.
jax.nn.softmax(scores, axis=-1) normalizes over the keys.
Divide by jnp.sqrt(d).astype(q.dtype) to keep bfloat16 inputs in bfloat16.
Accepts: array
Accepts: array
Accepts: array
