import numpy as np


def softmax(values: np.ndarray, axis: int = -1) -> np.ndarray:
    shifted = values - np.max(values, axis=axis, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=axis, keepdims=True)


def causal_mask(length: int) -> np.ndarray:
    return np.triu(np.full((length, length), -np.inf), k=1)


def scaled_dot_product_attention(q: np.ndarray, k: np.ndarray, v: np.ndarray, causal: bool = False):
    if q.ndim != 2 or k.ndim != 2 or v.ndim != 2:
        raise ValueError("q, k and v must be rank-2 arrays")
    if q.shape[1] != k.shape[1] or k.shape[0] != v.shape[0]:
        raise ValueError("incompatible attention shapes")
    scores = q @ k.T / np.sqrt(q.shape[1])
    if causal:
        if q.shape[0] != k.shape[0]:
            raise ValueError("causal self-attention requires equal sequence lengths")
        scores = scores + causal_mask(q.shape[0])
    weights = softmax(scores)
    return weights @ v, weights


def rope_2d(values: np.ndarray, positions: np.ndarray, theta: float = 10_000.0) -> np.ndarray:
    if values.shape[-1] != 2:
        raise ValueError("this teaching implementation expects head_dim=2")
    angles = positions / theta**0
    cos, sin = np.cos(angles), np.sin(angles)
    x, y = values[..., 0], values[..., 1]
    return np.stack((x * cos - y * sin, x * sin + y * cos), axis=-1)