Compute scaled query-key dot products independently for multiple attention heads, without a softmax or mask. Queries have shape (B, S_q, D), keys have shape (B, S_k, D), and each token's model-width axis contains H consecutive head blocks.
dh=D/H Ab,h,i,j=dh∑r=1dhQb,h,i,rKb,h,j,rHere B is batch size, S_q and S_k are the query and key lengths, D is model width, H is num_heads, and d_h is head width. Q and K denote the query and key values within those head blocks. Indices b and h select a batch and head, i and j select query and key positions, r selects a coordinate within a head, and A contains the scores.
Return a PyTorch tensor of shape (B, H, S_q, S_k), preserving the shared input dtype and device.
Input: q = [[[1,0,0,1],[0,1,1,0]]], k = [[[1,0,1,0],[0,1,0,1]]], num_heads = 2
Output: [[[[0.707107,0],[0,0.707107]],[[0,0.707107],[0.707107,0]]]]
Explanation: Each 4-value token becomes two heads of width 2. Every head computes its own query-key dot products and divides them by sqrt(2).
Input: q = [[[1,2,3,4]]], k = [[[1,0,1,0],[0,1,0,1]]], num_heads = 2
Output: [[[[0.707107,1.414214]],[[2.12132,2.828427]]]]
Input: q = [[[2,0]]], k = [[[1,3]]], num_heads = 1
Output: [[[[1.414214]]]]
torch.reshape can separate a dimension even when the input is noncontiguous.
torch.transpose exchanges two tensor axes.
In torch.einsum, a label omitted from the output is summed over.
Accepts: array
Accepts: array
Accepts: number
Compute scaled query-key dot products independently for multiple attention heads, without a softmax or mask. Queries have shape (B, S_q, D), keys have shape (B, S_k, D), and each token's model-width axis contains H consecutive head blocks.
dh=D/H Ab,h,i,j=dh∑r=1dhQb,h,i,rKb,h,j,rHere B is batch size, S_q and S_k are the query and key lengths, D is model width, H is num_heads, and d_h is head width. Q and K denote the query and key values within those head blocks. Indices b and h select a batch and head, i and j select query and key positions, r selects a coordinate within a head, and A contains the scores.
Return a PyTorch tensor of shape (B, H, S_q, S_k), preserving the shared input dtype and device.
Input: q = [[[1,0,0,1],[0,1,1,0]]], k = [[[1,0,1,0],[0,1,0,1]]], num_heads = 2
Output: [[[[0.707107,0],[0,0.707107]],[[0,0.707107],[0.707107,0]]]]
Explanation: Each 4-value token becomes two heads of width 2. Every head computes its own query-key dot products and divides them by sqrt(2).
Input: q = [[[1,2,3,4]]], k = [[[1,0,1,0],[0,1,0,1]]], num_heads = 2
Output: [[[[0.707107,1.414214]],[[2.12132,2.828427]]]]
Input: q = [[[2,0]]], k = [[[1,3]]], num_heads = 1
Output: [[[[1.414214]]]]
torch.reshape can separate a dimension even when the input is noncontiguous.
torch.transpose exchanges two tensor axes.
In torch.einsum, a label omitted from the output is summed over.
Accepts: array
Accepts: array
Accepts: number
