1import numpy as np
 2
 3META = {
 4    "name": "gaussian_score_matching_1d",
 5    "domain": "loss",
 6    "description": "Scalar nonlinear Gaussian score-matching task with known N(0,1) target score for Stein discrepancy training."
 7}
 8
 9def get_dataset(seed, n_train=400, n_test=400):
10    rng = np.random.RandomState(seed)
11    ztr = rng.randn(n_train, 1).astype(np.float32)
12    zte = np.random.RandomState(seed + 5000).randn(n_test, 1).astype(np.float32)
13    def transport(z):
14        return (0.92 * z + 0.12 * np.sin(z)).astype(np.float32)
15    return {
16        "xtr": ztr, "ytr": transport(ztr),
17        "xte": zte, "yte": transport(zte),
18        "task": "regression", "metric": "mse", "out_dim": 1
19    }