Spaces:
Sleeping
Sleeping
| """The coordinate candidate under the contract: fresh state per measurement, freeze as the | |
| no-adaptation control, revert exact, and both modes producing a serializable report.""" | |
| from __future__ import annotations | |
| import json | |
| import pytest | |
| import torch | |
| from plastic.data.mechanisms import mechanism_batch, split_combinations | |
| from plastic.eval.contract import ContractSpec, run_contract | |
| from plastic.eval.coordinate_learner import CoordinateLearner | |
| from plastic.model.coordinate import CoordinateConfig, CoordinateDynamics | |
| def _tiny(**kw) -> CoordinateDynamics: | |
| torch.manual_seed(0) | |
| chunk = kw.pop("chunk", 16) | |
| cfg = CoordinateConfig(d_model=32, n_heads=2, n_layers=2, chunk=chunk, **kw) | |
| return CoordinateDynamics(cfg) | |
| def _spec() -> ContractSpec: | |
| # one 32-step episode per scored row; chunk 16 puts one update boundary inside it | |
| return ContractSpec(seq_len=32, eval_batch=2, stream_episodes=4, probe_steps=20) | |
| def test_contract_refuses_a_chunk_as_long_as_the_episode(): | |
| learner = CoordinateLearner(_tiny(chunk=32), mode="frozen") | |
| with pytest.raises(ValueError, match="no fast-update boundary"): | |
| run_contract(learner, _spec(), seed=1) | |
| def test_rows_are_measured_independently(): | |
| learner = CoordinateLearner(_tiny(), mode="frozen") | |
| b = _batch(4) | |
| both = learner.step_mse(b, adapt=True) | |
| from dataclasses import replace | |
| alone = [ | |
| learner.step_mse(replace(b, inputs=b.inputs[i : i + 1], target_delta=b.target_delta[i : i + 1], worlds=b.worlds[i : i + 1], reset_flag=b.reset_flag[i : i + 1]), adapt=True) | |
| for i in range(2) | |
| ] | |
| assert torch.allclose(both[0], alone[0][0], atol=1e-6) and torch.allclose(both[1], alone[1][0], atol=1e-6) | |
| def _batch(seed: int = 0): | |
| train, heldout = split_combinations(k=2, n_heldout=5, seed=0) | |
| g = torch.Generator().manual_seed(seed) | |
| return mechanism_batch(2, seq_len=32, episodes_per_seq=1, combos=heldout, policy="gaussian", rng=g) | |
| def test_step_mse_starts_fresh_every_call_and_freeze_differs_from_adapt(): | |
| learner = CoordinateLearner(_tiny(), mode="frozen") | |
| b = _batch() | |
| a1 = learner.step_mse(b, adapt=True) | |
| a2 = learner.step_mse(b, adapt=True) | |
| f1 = learner.step_mse(b, adapt=False) | |
| assert a1.shape == (2, 32) and torch.equal(a1, a2), "a measurement must not depend on earlier calls" | |
| assert not torch.equal(a1, f1), "fast updates must change predictions after the first chunk" | |
| # within the first chunk nothing has been proposed yet, so adapt and freeze agree there | |
| assert torch.allclose(a1[:, :16], f1[:, :16]) | |
| def test_fast_signals_summary_reports_proposals_only_when_stepped(): | |
| learner = CoordinateLearner(_tiny(), mode="frozen") | |
| b = _batch(3) | |
| assert learner.fast_signals_summary() is None | |
| learner.step_mse(b, adapt=True) | |
| s = learner.fast_signals_summary() | |
| # a 32-step episode with chunk 16 crosses two boundaries, at steps 16 and 32 | |
| assert s["boundaries"] == 2 and s["stepped_fraction"] == 1.0 | |
| assert s["inner_loss_before"] is not None and s["inner_loss_after"] is not None | |
| assert len(s["dW_norm_by_layer"]) == 2 == len(s["dtheta_norm_by_layer"]) == len(s["eta_by_layer"]) | |
| assert s["r_peak_inf"] <= s["bound_peak"] + 1e-6 | |
| learner.step_mse(b, adapt=False) | |
| f = learner.fast_signals_summary() | |
| assert f["stepped_fraction"] == 0.0 and f["inner_loss_before"] is None and f["dW_norm_by_layer"] is None | |
| def test_fast_updates_off_makes_adapt_equal_freeze(): | |
| learner = CoordinateLearner(_tiny(fast_updates=False), mode="frozen") | |
| b = _batch(1) | |
| assert torch.allclose(learner.step_mse(b, adapt=True), learner.step_mse(b, adapt=False)) | |
| def test_contract_runs_and_reverts_exactly(mode): | |
| model = _tiny() | |
| learner = CoordinateLearner(model, mode=mode, lr=1e-3, steps=2) | |
| report = run_contract(learner, _spec(), seed=3) | |
| json.dumps(report) | |
| assert report["revert"]["ok"] is True, report["revert"] | |
| assert report["compute"]["parameters"] == model.num_params() | |
| if mode == "frozen": | |
| assert report["acceptance"]["n_good"] == 0 | |
| assert all(row["delta_mse"] == 0.0 for row in report["transfer"].values()) | |
| else: | |
| assert report["acceptance"] == {"accepted_good": 1.0, "refused_bad": 0.0, "n_good": 2, "n_bad": 1} | |
| assert report["decisions"][0]["accepted"] is True | |
| def _moved_keys(model: CoordinateDynamics, batch) -> set[str]: | |
| """Which fast-parameter keys changed after one adapting pass from the initial state.""" | |
| from plastic.model.coordinate import THETA_KEY, W_KEYS | |
| x, y = batch.inputs, batch.target_delta | |
| state0 = model.init_state(int(x.shape[0])) | |
| with torch.no_grad(): | |
| _, state1, _ = model(x, state0, mode="chunk", target_delta=y, commit=True) | |
| moved: set[str] = set() | |
| for layer0, layer1 in zip(state0.omega, state1.omega): | |
| for k in tuple(W_KEYS) + (THETA_KEY,): | |
| if not torch.equal(layer0[k], layer1[k]): | |
| moved.add("theta" if k == THETA_KEY else "W") | |
| return moved | |
| def test_decay_only_and_coordinates_only_switches_move_only_their_parameters(): | |
| # at initialization the fast step barely changes predictions (about 1e-5), so distinctness | |
| # is checked on which fast parameters move, not on the outputs | |
| b = _batch(2) | |
| assert _moved_keys(_tiny(), b) == {"W", "theta"} | |
| assert _moved_keys(_tiny(adapt_W=False), b) == {"theta"} | |
| assert _moved_keys(_tiny(adapt_theta=False), b) == {"W"} | |
| assert _moved_keys(_tiny(fast_updates=False), b) == set() | |