1import numpy as np
 2
 3META = {
 4    "name": "masked_candidate_low_rank",
 5    "domain": "retrieval",
 6    "description": "Masked high-dimensional candidate vectors generated from a low-rank latent action-feature subspace; regression predicts candidate utility."
 7}
 8
 9
10def get_dataset(seed, n_train=400, n_test=400):
11    rng = np.random.RandomState(seed)
12    d, rank = 128, 8
13    q, _ = np.linalg.qr(rng.normal(size=(d, rank)))
14    w = rng.normal(size=rank)
15
16    def make(n, offset):
17        rr = np.random.RandomState(seed + offset)
18        z = rr.normal(size=(n, rank)).astype(np.float32)
19        x = (z @ q.T).astype(np.float32)
20        # A fixed latent utility with modest observation-independent noise.
21        y = (z @ w + 0.15 * np.sin(z[:, 0]) + 0.08 * rr.normal(size=n)).astype(np.float32)
22        mask = (rr.rand(n, d) < 0.5).astype(np.float32)
23        return x, y, mask
24
25    xtr, ytr, mtr = make(n_train, 17)
26    xte, yte, mte = make(n_test, 50017)
27    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
28            "mtr": mtr, "mte": mte, "task": "regression", "metric": "mse",
29            "out_dim": 1, "rank": rank, "p": 0.5}