1import numpy as np
 2
 3META = {"name": "relational_block_graph", "domain": "graph-nn", "description": "Multi-relational stochastic-block edge prediction with node features; targets are sampled held-out relation edges."}
 4
 5def get_dataset(seed, n_train=400, n_test=120):
 6    rng = np.random.default_rng(seed)
 7    n_nodes, d, m, r = 36, 8, 3, 2
 8    z = rng.integers(m, size=n_nodes)
 9    centers = rng.normal(0, 1, (m, d))
10    X = centers[z] + 0.35 * rng.normal(size=(n_nodes, d))
11    A = np.array([[[.78,.16,.32],[.16,.72,.20],[.32,.20,.70]],
12                  [[.65,.25,.15],[.25,.72,.28],[.15,.28,.62]]], dtype=np.float32)
13    pairs, ys = [], []
14    for _ in range(n_train + n_test):
15        u, v = rng.integers(n_nodes, size=2)
16        k = int(rng.integers(r))
17        y = float(rng.random() < A[k, z[u], z[v]])
18        pairs.append(np.concatenate([X[u], X[v], [k]]))
19        ys.append(y)
20    pairs = np.asarray(pairs, dtype=np.float32)
21    ys = np.asarray(ys, dtype=np.float32)[:, None]
22    mu = pairs[:n_train, :2*d].mean(0)
23    sd = pairs[:n_train, :2*d].std(0) + 1e-6
24    pairs[:, :2*d] = (pairs[:, :2*d] - mu) / sd
25    return {"xtr": pairs[:n_train], "ytr": ys[:n_train],
26            "xte": pairs[n_train:], "yte": ys[n_train:],
27            "task": "regression", "metric": "mse", "input_shape": (2*d+1,),
28            "out_dim": 1}