1import numpy as np
 2
 3META = {
 4    'name': 'fisher_multitoken_denoising',
 5    'domain': 'diffusion-sampling',
 6    'description': 'Correlated multi-token Gaussian sequences for neural denoising under an annealed noise path.'
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    L = 12
11    def make(n, rs):
12        t = np.linspace(0.0, 1.0, L, dtype=np.float32)
13        out = np.empty((n, L), dtype=np.float32)
14        for i in range(n):
15            a = rs.uniform(0.6, 1.4)
16            f = rs.uniform(0.7, 1.5)
17            ph = rs.uniform(0.0, 2.0*np.pi)
18            z = a*np.sin(2.0*np.pi*f*t + ph) + 0.35*np.cos(np.pi*f*t - 0.4*ph)
19            out[i] = z + rs.normal(0.0, 0.06, L)
20        return out.astype(np.float32)
21    xtr = make(n_train, np.random.RandomState(seed))
22    xte = make(n_test, np.random.RandomState(seed + 10007))
23    return {'xtr': xtr, 'ytr': xtr.copy(), 'xte': xte, 'yte': xte.copy(),
24            'task': 'regression', 'metric': 'mse', 'out_dim': L}