# Complete self-attention step by step
def simple_attention(X):
    Q = X  # Queries (simplified)
    K = X  # Keys 
    V = X  # Values

    scores = torch.mm(Q, K.transpose(0, 1))  # Compute similarities
    weights = F.softmax(scores, dim=1)        # Convert to probabilities
    output = torch.mm(weights, V)             # Weighted combination
    return output

# Usage
word_vectors = torch.randn(5, 8)  # 5 words, 8 dimensions each
attended_vectors = simple_attention(word_vectors)