1import numpy as np
 2META={"name":"categorical_persistent_sequences","domain":"diffusion-sampling","description":"Correlated binary multi-token sequences conditioned on continuous context; evaluates categorical persistent sampling."}
 3def get_dataset(seed,n_train,n_test):
 4    rng=np.random.RandomState(seed); L=8; K=2
 5    def make(n):
 6        z=rng.choice(2,size=n).astype(np.float32)
 7        phase=rng.uniform(-1,1,size=n).astype(np.float32)
 8        x=np.stack([z*2-1,phase,np.sin(phase),np.cos(phase),rng.normal(size=n),rng.uniform(-1,1,size=n)],1).astype(np.float32)
 9        y=np.empty((n,L),np.int64); y[:,0]=rng.binomial(1,0.5,size=n)
10        for j in range(1,L):
11            # persistent categorical transitions, with context-dependent bias
12            p=np.clip(.88-.12*z + .04*np.sin(phase+j),.55,.93)
13            y[:,j]=np.where(rng.rand(n)<p,y[:,j-1],1-y[:,j-1])
14        y ^= (z[:,None].astype(np.int64) if False else 0)
15        return x,y
16    xtr,ytr=make(n_train); xte,yte=make(n_test)
17    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"classification","metric":"err","out_dim":L*K}