1import numpy as np
 2
 3META = {
 4    'name': 'bidirectional_conditional_joint',
 5    'domain': 'conditional_compatibility',
 6    'description': 'Strictly positive categorical joint samples for two neural conditional directions.'
 7}
 8
 9def get_dataset(seed, n_train=400, n_test=400):
10    rng = np.random.default_rng(int(seed))
11    raw = np.array([
12        [1.0, 2.0, .7, 1.4, .8, 1.7],
13        [2.1, .8, 1.8, .5, 1.3, .9],
14        [.6, 1.7, 2.4, 1.1, .9, 1.5],
15        [1.3, .9, 1.5, 2.2, 1.1, .7],
16        [1.8, 1.2, .5, 1.6, 2.0, 1.0],
17        [.7, 1.4, 1.1, .8, 1.6, 2.3]], dtype=np.float64)
18    joint = raw / raw.sum()
19    z = rng.choice(36, size=n_train + n_test, p=joint.reshape(-1))
20    pairs = np.stack([z // 6, z % 6], axis=1).astype(np.int64)
21    return {
22        'xtr': pairs[:n_train], 'ytr': pairs[:n_train, 0],
23        'xte': pairs[n_train:], 'yte': pairs[n_train:, 0],
24        'task': 'classification', 'metric': 'bidirectional_nll',
25        'input_shape': (2,), 'out_dim': 6, 'nx': 6, 'ny': 6,
26        'joint': joint
27    }