import torch
from rsil import RSILConfig, RSILForMaskedLM, ResonantSpectralMixer

def test_forward_backward():
 c=RSILConfig(vocab_size=101,dim=32,depth=2,n_bands=3,groups=4,max_seq_len=64)
 m=RSILForMaskedLM(c); x=torch.randint(0,101,(2,64)); y=x.clone(); o=m(x,y); assert o['logits'].shape==(2,64,101); o['loss'].backward()
def test_translation_equivariance():
 c=RSILConfig(dim=32,n_bands=3,groups=4,max_degree=3); m=ResonantSpectralMixer(c).eval(); x=torch.randn(2,64,32)
 with torch.no_grad():
  for p in m.gates.values(): p.normal_(0,.1)
  a=torch.roll(m(x),7,1); b=m(torch.roll(x,7,1))
 assert (a-b).abs().max()<2e-4