1"""Small deterministic attributed relational-graph benchmark for stage-2."""
 2import numpy as np
 3
 4META = {
 5    "name": "relational_graph_classification",
 6    "domain": "graph-nn",
 7    "description": "Permutation-augmented attributed graph classification with dense relational kernels.",
 8}
 9
10
11def _make(rng, n):
12    out, ys = [], []
13    for _ in range(int(n)):
14        y = int(rng.randint(0, 2))
15        p = 0.22 if y == 0 else 0.58
16        a = (rng.rand(8, 8) < p).astype(np.float32)
17        a = np.triu(a, 1)
18        a = a + a.T
19        np.fill_diagonal(a, 0.0)
20        node = rng.normal(0.0, 0.12, size=(8, 1)).astype(np.float32)
21        node += 0.08 if y else -0.08
22        out.append(np.concatenate([a, node], axis=1))
23        ys.append(y)
24    return np.asarray(out, dtype=np.float32), np.asarray(ys, dtype=np.int64)
25
26
27def get_dataset(seed, n_train, n_test):
28    xtr, ytr = _make(np.random.RandomState(int(seed)), n_train)
29    xte, yte = _make(np.random.RandomState(int(seed) + 100003), n_test)
30    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
31            "task": "classification", "metric": "err", "out_dim": 2}