1import numpy as np
 2
 3META = {
 4    'name': 'conditional_sequence_diffusion',
 5    'domain': 'sequence-level diffusion sampling',
 6    'description': 'Conditional generation of multi-token sinusoidal sequences from a context window; the metric is held-out conditional reconstruction MSE.'
 7}
 8
 9def get_dataset(seed, n_train=400, n_test=100):
10    rng = np.random.default_rng(int(seed))
11    n = n_train + n_test
12    ctx = rng.uniform(-1, 1, (n, 4)).astype('float32')
13    phase = rng.uniform(0, 2*np.pi, n)
14    freq = rng.uniform(.7, 1.8, n)
15    amp = .7 + .5 * (ctx[:, 0] + 1) / 2
16    tt = np.arange(8, dtype='float32')[None, :]
17    y = (amp[:, None] * np.sin(freq[:, None] * tt + phase[:, None]) +
18         .25 * ctx[:, 1, None] * np.cos(.55 * tt + ctx[:, 2, None]) +
19         .08 * rng.normal(size=(n, 8))).astype('float32')
20    # Context contains informative phase/frequency proxies, making this a conditional task.
21    ctx[:, 2] = np.sin(phase); ctx[:, 3] = np.cos(phase)
22    return {'xtr': ctx[:n_train], 'ytr': y[:n_train],
23            'xte': ctx[n_train:], 'yte': y[n_train:], 'task': 'regression', 'metric': 'mse',
24            'input_shape': (4,), 'out_dim': 8, 'track': META['name']}