1"""NumPy reference implementation of Laplace-heterogeneous MoE routing."""
 2import numpy as np
 3
 4
 5class LaplaceHeterogeneousRouter:
 6    def __init__(self, n_experts, lambdas=(0.25, 1.0, 4.0, 16.0),
 7                 weights=None, alpha=1.0, beta=0.9, epsilon=1e-12):
 8        self.n_experts = int(n_experts)
 9        self.lambdas = np.asarray(lambdas, dtype=float)
10        if self.lambdas.ndim != 1 or len(self.lambdas) == 0 or np.any(self.lambdas < 0):
11            raise ValueError("lambdas must be nonempty and nonnegative")
12        if weights is None:
13            weights = np.ones(len(self.lambdas)) / len(self.lambdas)
14        self.weights = np.asarray(weights, dtype=float)
15        if self.weights.shape != self.lambdas.shape or np.any(self.weights < 0) or self.weights.sum() <= 0:
16            raise ValueError("weights must be nonnegative and match lambdas")
17        self.weights /= self.weights.sum()
18        self.alpha, self.beta, self.epsilon = float(alpha), float(beta), float(epsilon)
19        self.pressure = np.zeros(self.n_experts)
20
21    @staticmethod
22    def softmax(logits):
23        z = logits - logits.max(axis=1, keepdims=True)
24        e = np.exp(z)
25        return e / e.sum(axis=1, keepdims=True)
26
27    def availability(self):
28        return (self.weights[:, None] * np.exp(-self.lambdas[:, None] * self.pressure[None, :])).sum(axis=0)
29
30    def effective_hazard(self):
31        terms = self.weights[:, None] * np.exp(-self.lambdas[:, None] * self.pressure[None, :])
32        q = terms.sum(axis=0)
33        return (terms * self.lambdas[:, None]).sum(axis=0) / np.maximum(q, self.epsilon)
34
35    def route(self, logits, top_k=2):
36        logits = np.asarray(logits, dtype=float)
37        if logits.ndim != 2 or logits.shape[1] != self.n_experts:
38            raise ValueError("logits must have shape [batch, n_experts]")
39        if not 1 <= top_k <= self.n_experts:
40            raise ValueError("invalid top_k")
41        mass = self.softmax(logits).mean(axis=0)
42        self.pressure = self.beta * self.pressure + (1.0 - self.beta) * mass
43        q = self.availability()
44        adjusted = logits + self.alpha * np.log(q + self.epsilon)[None, :]
45        chosen = np.argpartition(-adjusted, top_k - 1, axis=1)[:, :top_k]
46        return chosen, adjusted, q, self.pressure.copy()