| import torch | |
| from spectral_world_models.models import SelectiveSpectralTransition | |
| def test_selective_caps_only_expansive_modes(): | |
| m=SelectiveSpectralTransition(16,5,n_ops=2,rank=4,max_gain=1.05) | |
| ids=torch.arange(5) | |
| _, raw, clipped=m.selective_matrix(ids) | |
| assert float(clipped.max()) <= 1.05001 | |
| mask=raw <= 1.05 | |
| assert torch.allclose(clipped[mask], raw[mask], atol=1e-6) | |
| def test_selective_forward_shape(): | |
| m=SelectiveSpectralTransition(16,5,n_ops=2,rank=4,max_gain=1.05) | |
| y=m(torch.randn(3,16),torch.tensor([0,1,4])) | |
| assert y.shape==(3,16) | |