1import numpy as np
 2
 3META = {
 4    "name": "router_regime_regression",
 5    "domain": "moe-routing",
 6    "description": "Regression with three observable nonlinear regimes, requiring a learned fixed-budget expert router."
 7}
 8
 9
10def get_dataset(seed, n_train, n_test):
11    rng = np.random.default_rng(seed)
12    def make(n):
13        x = rng.uniform(-2.0, 2.0, size=(n, 4)).astype(np.float32)
14        # Three regimes are observable from x but have distinct functions.
15        z = np.argmax(np.stack([x[:, 0], -x[:, 0], x[:, 1]], axis=1), axis=1)
16        y = np.where(z == 0, np.sin(2*x[:, 1]) + 0.25*x[:, 2],
17            np.where(z == 1, 0.5*x[:, 0]**2 - x[:, 2],
18                     x[:, 1]*x[:, 2] + 0.4*x[:, 3]))
19        y += 0.04 * rng.normal(size=n)
20        return x, y.astype(np.float32)[:, None]
21    xtr, ytr = make(n_train); xte, yte = make(n_test)
22    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
23            "task": "regression", "metric": "mse", "out_dim": 1}