1import numpy as np
 2META = {"name":"correlated_multitask_regression","domain":"multi_task_learning","description":"Six regression tasks with latent task groups and correlated observation residuals; used to test residual-redundancy routing."}
 3
 4def get_dataset(seed, n_train, n_test):
 5    rng=np.random.default_rng(seed); p=12; t=6
 6    w=rng.normal(size=(3,p)).astype(np.float32)
 7    nuisance=rng.normal(size=(t,p)).astype(np.float32)
 8    cov=np.eye(t,dtype=np.float32)*.20
 9    for a,b in [(0,1),(2,3),(4,5)]: cov[a,b]=cov[b,a]=.16
10    def make(n):
11        x=rng.normal(size=(n,p)).astype(np.float32)
12        y=np.zeros((n,t),dtype=np.float32)
13        for j,g in enumerate([[0,1],[2,3],[4,5]]):
14            latent=np.tanh(x@w[j]) + .35*np.sin(x[:,j%4])
15            for q in g: y[:,q]=latent + .12*(x@nuisance[q])
16        y += rng.multivariate_normal(np.zeros(t),cov,size=n).astype(np.float32)
17        return x,y
18    xtr,ytr=make(n_train); xte,yte=make(n_test)
19    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":6}