1from __future__ import annotations
 2import numpy as np
 3
 4
 5def tetrahedron_design() -> np.ndarray:
 6    """Four unit vertices of a regular tetrahedron, a spherical 2-design."""
 7    X = np.array([[1, 1, 1], [1, -1, -1], [-1, 1, -1], [-1, -1, 1]], dtype=float).T
 8    return X / np.linalg.norm(X, axis=0, keepdims=True)
 9
10
11def slack_matrix(X: np.ndarray, h: np.ndarray, c: float, U: np.ndarray,
12                 clamp: bool = False) -> np.ndarray:
13    A = c * np.outer(h, h) - X.T @ U.T @ X
14    return np.maximum(A, 0.0) if clamp else A
15
16
17def polar_attention_logits(qk_logits: np.ndarray, A: np.ndarray,
18                           alpha: float = 1.0, eps: float = 1e-8) -> np.ndarray:
19    return qk_logits + alpha * np.log(np.maximum(A, 0.0) + eps)
20
21
22def row_softmax(z: np.ndarray) -> np.ndarray:
23    z = z - np.max(z, axis=-1, keepdims=True)
24    p = np.exp(z)
25    return p / np.sum(p, axis=-1, keepdims=True)