import numpy as np
from numpy_core import attention

q = np.array([[1., 0.]])
k = np.array([[1., 0.], [0., 1.], [1., 1.]])
v = np.array([[1., 0.], [0., 2.], [2., 1.]])
output, weights = attention(q, k, v)

np.testing.assert_allclose(weights.sum(-1), 1)
np.testing.assert_allclose(output, [[1.203336, 0.796664]], atol=1e-5)
print(np.round(weights, 4))
print(np.round(output, 4))