import math


def attend(query, keys, values, visible):
    scores = [
        sum(q * k for q, k in zip(query, key)) / math.sqrt(len(query))
        for key, allowed in zip(keys, visible)
        if allowed
    ]
    peak = max(scores)
    weights = [math.exp(score - peak) for score in scores]
    scale = sum(weights)
    allowed_weights = iter(weight / scale for weight in weights)
    weights = [next(allowed_weights) if allowed else 0.0 for allowed in visible]
    output = [sum(weight * value[index] for weight, value in zip(weights, values))
              for index in range(len(values[0]))]
    return weights, output


keys = [[1, 0], [0, 1], [1, 1]]
values = [[1, 0], [0, 2], [3, 3]]
visible = [True, True, False]

for query in ([2, 0], [0, 2]):
    weights, mixed = attend(query, keys, values, visible)
    print(f"Query {query}: weights {[round(w, 3) for w in weights]}, "
          f"mixed value {[round(value, 3) for value in mixed]}")