1import numpy as np
 2META={"name":"unordered_set_regression","domain":"set_models","description":"Invariant regression on six exchangeable 2D tokens with unary and pairwise interactions."}
 3def get_dataset(seed,n_train,n_test):
 4 def make(rng,n):
 5  x=rng.normal(size=(n,6,2)).astype(np.float32)
 6  unary=(x*x).sum(2).mean(1)
 7  pair=np.tanh((x[:,:,None,:]*x[:,None,:,:]).sum(3).mean((1,2)))
 8  return x,(unary+0.35*pair).astype(np.float32)
 9 xtr,ytr=make(np.random.RandomState(seed),n_train); xte,yte=make(np.random.RandomState(seed+5000),n_test)
10 return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","out_dim":1}