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