| 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 | |