Spaces:
Sleeping
Sleeping
| import pytest | |
| import torch | |
| from plastic.config import ModelConfig | |
| from plastic.model.memory import FastWeightMemory | |
| from plastic.model.state import SessionState | |
| def test_chunk_and_recurrent_modes_agree(device): | |
| cfg = ModelConfig(d_model=32, n_heads=2, chunk=16) | |
| mem = FastWeightMemory(cfg).to(device) | |
| u = torch.randn(2, 40, 32, device=device) | |
| S0 = torch.randn(2, 2, 16, 16, device=device) * 0.1 | |
| m1, S1, _, _, sig1 = mem(u, S0, None, mode="chunk") | |
| m2, S2, _, _, sig2 = mem(u, S0, None, mode="recurrent") | |
| assert torch.allclose(m1, m2, atol=1e-4, rtol=1e-4) | |
| assert torch.allclose(S1, S2, atol=1e-4, rtol=1e-4) | |
| assert torch.allclose(sig1.write_norm, sig2.write_norm, atol=1e-4) | |
| def test_beta_scale_zero_freezes_memory(device): | |
| cfg = ModelConfig(d_model=32, n_heads=2, chunk=16) | |
| mem = FastWeightMemory(cfg).to(device) | |
| u = torch.randn(1, 20, 32, device=device) | |
| S0 = torch.randn(1, 2, 16, 16, device=device) * 0.1 | |
| _, S1, _, _, sig = mem(u, S0, None, mode="chunk", beta_scale=0.0) | |
| assert sig.write_norm.abs().max() == 0 | |
| assert (S1.abs() <= S0.abs() + 1e-6).all() | |
| def test_signals_shapes(device): | |
| cfg = ModelConfig(d_model=32, n_heads=2, chunk=16) | |
| mem = FastWeightMemory(cfg).to(device) | |
| u = torch.randn(3, 17, 32, device=device) | |
| S0 = torch.zeros(3, 2, 16, 16, device=device) | |
| m, S, M, _, sig = mem(u, S0, None, mode="chunk") | |
| assert m.shape == (3, 17, 32) and S.shape == (3, 2, 16, 16) and M is None | |
| assert sig.err.shape == sig.beta.shape == sig.alpha.shape == sig.write_norm.shape == (3, 2, 17) | |
| assert (sig.beta > 0).all() and (sig.beta < 1).all() and (sig.alpha > 0).all() and (sig.alpha < 1).all() | |
| def test_chunk_rule_runs_and_reads_chunk_start_weights(device): | |
| cfg = ModelConfig(d_model=32, n_heads=2, chunk=8, rule="chunk") | |
| mem = FastWeightMemory(cfg).to(device) | |
| u = torch.randn(1, 20, 32, device=device) | |
| st = SessionState.zeros(cfg, batch=1, device=device).layers[0] | |
| m, S, M, chunk, sig = mem(u, st.S, st.M, mode="chunk", chunk0=st.chunk) | |
| assert m.shape == (1, 20, 32) and M is not None and M.shape == S.shape | |
| assert chunk is not None and chunk.count == 4 # 20 = 8 + 8 + 4 pending | |
| # first chunk reads from S0 = 0 so its memory output is exactly zero | |
| assert torch.allclose(m[:, :8], torch.zeros_like(m[:, :8])) | |
| nz = (sig.write_norm[0, 0] != 0).nonzero().flatten().tolist() | |
| assert nz == [7, 15] # one write per completed chunk, reported on its last token | |
| def test_chunk_rule_is_stream_equivalent(device, mode): | |
| """Any partition of the token stream into forward calls gives identical results.""" | |
| cfg = ModelConfig(d_model=32, n_heads=2, chunk=8, rule="chunk") | |
| mem = FastWeightMemory(cfg).to(device) | |
| u = torch.randn(1, 37, 32, device=device) | |
| st = SessionState.zeros(cfg, batch=1, device=device).layers[0] | |
| st.S += torch.randn_like(st.S) * 0.1 | |
| m_full, S_full, M_full, c_full, sig_full = mem(u, st.S, st.M, mode=mode, chunk0=st.chunk) | |
| S, M, c = st.S, st.M, st.chunk | |
| outs, writes = [], [] | |
| cuts = [0, 3, 8, 9, 20, 21, 30, 37] | |
| for a, b in zip(cuts[:-1], cuts[1:]): | |
| m, S, M, c, sig = mem(u[:, a:b], S, M, mode=mode, chunk0=c) | |
| outs.append(m) | |
| writes.append(sig.write_norm) | |
| assert torch.allclose(m_full, torch.cat(outs, 1), atol=1e-5, rtol=1e-5) | |
| assert torch.allclose(S_full, S, atol=1e-5) and torch.allclose(M_full, M, atol=1e-5) | |
| assert c.count == c_full.count == 37 % 8 | |
| assert torch.allclose(c.A, c_full.A, atol=1e-5) | |
| assert torch.allclose(sig_full.write_norm, torch.cat(writes, 2), atol=1e-5) | |
| def test_chunk_rule_beta_scale_caps_orthogonalized_write(device): | |
| cfg = ModelConfig(d_model=32, n_heads=2, chunk=8, rule="chunk", chunk_orthogonalize=True) | |
| mem = FastWeightMemory(cfg).to(device) | |
| u = torch.randn(1, 8, 32, device=device) | |
| st = SessionState.zeros(cfg, batch=1, device=device).layers[0] | |
| _, S1, M1, _, _ = mem(u, st.S, st.M, mode="chunk", chunk0=st.chunk, beta_scale=1.0) | |
| _, Sh, Mh, _, _ = mem(u, st.S, st.M, mode="chunk", chunk0=st.chunk, beta_scale=0.5) | |
| _, S0, M0, _, sig0 = mem(u, st.S, st.M, mode="chunk", chunk0=st.chunk, beta_scale=0.0) | |
| assert torch.allclose((Sh - st.S).norm(), 0.5 * (S1 - st.S).norm(), rtol=1e-4) | |
| assert torch.equal(S0, st.S) and torch.equal(M0, st.M) | |
| assert sig0.write_norm.abs().max() == 0 | |
| # warm momentum: scale 0 still writes nothing, momentum only decays, S only decays | |
| _, S2, M2, _, sig2 = mem(u, S1, M1, mode="chunk", chunk0=st.chunk, beta_scale=0.0) | |
| assert torch.allclose(M2, cfg.chunk_momentum * M1) | |
| alpha_chunk = sig2.alpha.mean(dim=-1) | |
| assert torch.allclose(S2, alpha_chunk[..., None, None] * S1, atol=1e-6) | |
| def test_mlp_memory_not_available_yet(): | |
| with pytest.raises(NotImplementedError): | |
| FastWeightMemory(ModelConfig(rule="chunk", memory="mlp")) | |