1import numpy as np
 2
 3META = {
 4    'name': 'directed_path_holonomy',
 5    'domain': 'graph',
 6    'description': 'Directed length-3 path composition regression with noncommutative edge operators; order is essential.'
 7}
 8
 9
10def get_dataset(seed, n_train, n_test):
11    rng = np.random.default_rng(int(seed))
12
13    def make(n):
14        x = rng.normal(0.0, 0.45, size=(n, 3, 2, 2)).astype(np.float32)
15        x += np.eye(2, dtype=np.float32)[None, None, :, :] * 0.8
16        h = np.matmul(np.matmul(x[:, 0], x[:, 1]), x[:, 2])
17        y = h.reshape(n, 4).astype(np.float32)
18        return x.reshape(n, 12), y
19
20    xtr, ytr = make(n_train)
21    xte, yte = make(n_test)
22    return {
23        'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
24        'task': 'regression', 'metric': 'mse',
25        'input_shape': (12,), 'out_dim': 4,
26    }