Scaled_dot_product_attention(q, k, v, mask) function:

Maybe you use the tensorflow tutorial: