1import numpy as np
 2META = {"name":"gaussian_score_matching","domain":"loss","description":"Maps latent Gaussian inputs to 2D target samples; known N(0,I) score permits Langevin Stein KSD training."}
 3
 4def get_dataset(seed, n_train=400, n_test=400):
 5    rng=np.random.RandomState(seed)
 6    # A fixed nonlinear transport makes ordinary supervised fitting nontrivial,
 7    # while the output distribution remains close to a known Gaussian target.
 8    ztr=rng.randn(n_train,2).astype(np.float32)
 9    zte=np.random.RandomState(seed+5000).randn(n_test,2).astype(np.float32)
10    def transport(z):
11        y=z.copy(); y[:,0]=0.92*z[:,0]+0.12*np.sin(z[:,1]); y[:,1]=0.92*z[:,1]+0.12*np.sin(z[:,0]); return y.astype(np.float32)
12    return {"xtr":ztr,"ytr":transport(ztr),"xte":zte,"yte":transport(zte),"task":"regression","metric":"mse","out_dim":2}