← 返回 meta 的题目列表Implement Attention Mechanism
类型：online_judge
Implement a basic Attention mechanism. You need to encode and implement according to the given formula. The specific requirements are as follows:
Input:
Queries (Q): a 2D array representing query vectors.
Keys (K): a 2D array representing key vectors.
Values (V): a 2D array representing value vectors.
num_heads: an integer indicating the number of attention heads.
Output:
A 2D array representing the weighted sum obtained by the Attention mechanism.
Assumptions:
The shapes of Queries, Keys, and Values are (batch_size, seq_length, depth).
depth is divisible by num_heads.
Test Cases:
Input:
Queries: [[1, 0], [0, 1]], Keys: [[1, 0], [0, 1]], Values: [[1, 0], [0, 1]], num_heads: 1
Output: [[1, 0], [0, 1]]
Input:
Queries: [[1, 1]], Keys: [[1, 0], [0, 1]], Values: [[1, 2], [3, 4]], num_heads: 1
Output: [[2, 3]]
Example
Input
Queries: [[1, 0], [0, 1]], Keys: [[1, 0], [0, 1]], Values: [[1, 0], [0, 1]], num_heads: 1