1import numpy as np
 2
 3META = {
 4    "name": "sparse_cycle_compatibility",
 5    "domain": "masked_categorical_compatibility",
 6    "description": "Paired conditional categorical prediction on a sparse bipartite support with a compatible-joint generator."
 7}
 8N = 8
 9EDGES = np.array([(0,0),(0,1),(1,1),(1,2),(2,2),(2,3),(3,3),(3,4),
10                  (4,4),(4,5),(5,5),(5,6),(6,6),(6,7),(7,7),(7,0),
11                  (0,4),(2,6),(4,0),(6,2)], dtype=np.int64)
12
13def get_dataset(seed, n_train, n_test):
14    rng = np.random.default_rng(int(seed))
15    ux, vy = rng.normal(0, .45, N), rng.normal(0, .45, N)
16    joint_logits = np.array([ux[x] + vy[y] for x, y in EDGES])
17    p = np.exp(joint_logits - joint_logits.max()); p /= p.sum()
18    def sample(k):
19        z = rng.choice(len(EDGES), k, p=p)
20        return EDGES[z, 1].astype(np.int64), EDGES[z, 0].astype(np.int64)
21    ytr, xtr_target = sample(n_train)
22    yte, xte_target = sample(n_test)
23    return {
24        "xtr": np.eye(N, dtype=np.float32)[ytr], "ytr": xtr_target,
25        "xte": np.eye(N, dtype=np.float32)[yte], "yte": xte_target,
26        "task": "classification", "metric": "err", "out_dim": N,
27        "n": N, "edges": EDGES.copy(), "train_y": ytr, "test_y": yte,
28        "train_x": xtr_target, "test_x": xte_target,
29    }