1import numpy as np
 2
 3META = {
 4    'name': 'quadratic_context_multitoken',
 5    'domain': 'diffusion-sampling',
 6    'description': 'Multitoken denoising with transferable analytic quadratic Gaussian context.'
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    rng = np.random.RandomState(seed)
11    L = 8
12    def make(n):
13        t = np.linspace(0, 1, L, dtype=np.float32)
14        z = rng.choice([-1., 1.], size=n).astype(np.float32)
15        phase = rng.uniform(-np.pi, np.pi, n).astype(np.float32)
16        amp = rng.uniform(.8, 1.2, n).astype(np.float32)
17        x = np.empty((n, L), dtype=np.float32)
18        for i in range(n):
19            x[i] = amp[i] * (np.sin(2*np.pi*1.25*t + phase[i]) + .35*z[i]*np.cos(2*np.pi*2.0*t))
20        return x
21    rtr, rte = make(n_train), make(n_test)
22    k = 1.5
23    return {'xtr': rtr, 'ytr': rtr/(1+k), 'xte': rte, 'yte': rte/(1+k),
24            'task': 'regression', 'metric': 'mse', 'out_dim': L, 'k': k}