1import numpy as np
 2
 3META = {"name": "phase_checkpoint_regression", "domain": "training-dynamics",
 4        "description": "Regression task with explicit worker/checkpoint phase metadata."}
 5
 6def get_dataset(seed, n_train, n_test):
 7    rng = np.random.RandomState(seed)
 8    def gen(n):
 9        x = rng.uniform(-1, 1, (n, 10)).astype(np.float32)
10        y = (np.sin(2*x[:, 0]) + .5*x[:, 1]**2 - .3*x[:, 2]*x[:, 3]
11             + .2*x[:, 4] + .05*rng.normal(size=n)).astype(np.float32)
12        return x, y
13    xtr, ytr = gen(n_train); xte, yte = gen(n_test)
14    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
15            "task": "regression", "metric": "mse", "out_dim": 1}