1/* gpt2.rs's attention_of_one_key_is_its_value. (Golden data for attention over many keys is in
2 * @ajax/mesh-ref-gpt2's tests.) */
3#include "mesh/ref/gpt2_attend.h"
4
5int main(void) {
6    float q[64], k[64], v[64], out[64];
7    for (int i = 0; i < 64; i++) {
8        q[i] = (float)i / 64.0f;
9        k[i] = 1.0f - (float)i / 64.0f;
10        v[i] = (float)i * 0.25f - 3.0f;
11    }
12    if (!mesh_gpt2_attend_row(q, 1, k, 0, v, 0, out)) return 1;
13    for (int i = 0; i < 64; i++)
14        if (out[i] != v[i]) return 1;
15    return !mesh_gpt2_attend_row(q, 0, k, 0, v, 0, out) || mesh_gpt2_attend_row(q, 1025, k, 0, v, 0, out);
16}