File size: 577 Bytes
4fd79a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | 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)
|