1import numpy as np
 2
 3META = {"name": "congestion_moe_sequence", "domain": "moe-routing", "description": "Multi-token regime-conditioned regression where token groups are processed by shared experts."}
 4
 5def get_dataset(seed, n_train, n_test):
 6    def make(n, s):
 7        rng = np.random.default_rng(s)
 8        length, dim = 12, 4
 9        regime = rng.integers(0, 4, size=n)
10        x = rng.normal(0, 0.7, size=(n, length, dim)).astype(np.float32)
11        x[:, :, 0] += (regime[:, None] - 1.5) * 0.8
12        x[:, :, 1] += np.sin(regime[:, None] * 1.7)
13        y = (0.8 * np.sin(x[:, 0, 0] + regime) + 0.5 * x[:, 3, 1] ** 2
14             - 0.35 * x[:, 7, 2] + 0.25 * (regime - 1.5)
15             + rng.normal(0, 0.08, n))
16        return x, y.astype(np.float32)[:, None]
17    xtr, ytr = make(n_train, seed)
18    xte, yte = make(n_test, seed + 5000)
19    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
20            "task": "regression", "metric": "mse", "out_dim": 1}