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