Practice ML/AI coding interviews.

import math
import torch

def scaled_dot_product_attention(query, key, value, mask):
    # 1. Compute scaled query-key scores
    d_k = query.shape[-1]
    scores = query @ key.transpose(-2, -1) / math.sqrt(d_k)

    # 2. Mask and apply stable softmax
    scores = scores.masked_fill(~mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)

    # 3. Weighted sum of values
    return weights @ value