Buckets:
| """cwm3/tests/test_v2_followups.py — contracts of the 2026-09-11 follow-ups: kin_version switch (v2 causal packs), the regionpca256 road stream | |
| routing (pack v2 only) and the EnvBranch side-vector width (d_side 512 in place of vf256), variant-level recipe merge (select_LF false).""" | |
| from pathlib import Path | |
| import sys | |
| import numpy as np | |
| import pytest | |
| import torch | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train"): | |
| sys.path.insert(0, str(p)) | |
| import cwm3_dataset as ds # noqa: E402 | |
| from models.branches import EnvBranch # noqa: E402 | |
| def test_kin_version_switch_resolves_v2_dir(): | |
| try: | |
| ds.set_kin_source("causal_routes", "v2") | |
| f = ds.files("manual", "val") | |
| assert f["kinhist"].parent.name == "targets_kin_causal_routes_v2" and f["kin"].parent.name == "targets_kin_causal_routes_v2" | |
| ds.set_kin_source("oldin_ctgt", "v2") | |
| f = ds.files("manual", "val") | |
| assert f["kin"].parent.name == "targets_kin_causal_routes_v2" and f["kinhist"].parent.name == "targets_kin" # old inputs, v2 causal targets | |
| ds.set_kin_source("causal_routes", "v1") | |
| assert ds.files("manual", "val")["kin"].parent.name == "targets_kin_causal_routes" | |
| finally: | |
| ds.set_kin_source("cwm2", "v1") | |
| def test_regionpca256_is_pack_v2_only_and_row_aligned(): | |
| p = ds.PACK_V2 / "regionpca256_manual_val.npz" | |
| if not p.exists(): | |
| pytest.skip("regionpca256 stream not built") | |
| with pytest.raises(AssertionError): | |
| ds.load_split("manual", "val", need={"regionpca256"}) # pack v1 -> refused | |
| try: | |
| ds.set_pack_version("v2") | |
| raw = ds.load_split("manual", "val", need={"regionpca256", "vf256"}, max_anchors=256, seed=1, verbose=lambda m: None) | |
| assert raw["regionpca256"].shape == (256, 512) and raw["regionpca256"].dtype == np.float16 and raw["regionpca256_ok"].shape == (256,) | |
| assert raw["vf256"].shape == (256, 256) and raw["meta"]["pack_version"] == "v2" | |
| finally: | |
| ds.set_pack_version("v1") | |
| def test_envbranch_side_512_and_default_identity(): | |
| torch.manual_seed(0) | |
| e256 = EnvBranch(tokens="pooled", n_slots=16) | |
| e512 = EnvBranch(tokens="pooled", n_slots=16, d_side=512) | |
| n = lambda m: sum(p.numel() for p in m.parameters()) | |
| assert n(e256) == 880_604 # the E0 pooled parameter count (train.log of every pooled run) | |
| assert n(e512) - n(e256) == 256 * 128 + 2 * 256 # Linear(512 vs 256 -> 128) + LayerNorm(512 vs 256) | |
| B = 8 | |
| b = {"tok": torch.randn(B, 16, 1, 256), "tok_ok": torch.ones(B, 16, dtype=torch.bool), "tok_dt": -torch.arange(16).flip(0).float()[None].expand(B, -1) * 0.5, | |
| "vf": torch.randn(B, 512), "vf_ok": torch.ones(B, dtype=torch.bool), "mode": torch.zeros(B, dtype=torch.long)} | |
| out = e512(b) | |
| assert out["E"].shape == (B, 25, 12) and out["F"].shape == (B, 5, 256) | |
| with pytest.raises(RuntimeError): | |
| e256(b) # a 512-d side vector cannot be fed to the vf256 branch by accident | |
| def test_variant_recipe_merge_and_synthetic_stream(): | |
| sys.argv = ["x"] | |
| import train_cwm3 as t | |
| c = t.stage_cfg("E0", "regionpca256") | |
| assert c["recipe"]["select_LF"] is False and c["recipe"]["batch"] == 1024 and c["side"] == "regionpca256" and c["tokens"] == "pooled" | |
| assert t.model_kwargs(c)["d_side"] == 512 and "regionpca256" in t.need_streams(c) and "vf256" not in t.need_streams(c) | |
| c0 = t.stage_cfg("E0", "pooled") | |
| assert "select_LF" not in c0["recipe"] and t.model_kwargs(c0)["d_side"] == 256 | |
| raw = ds.synthetic_raw(8, need=t.need_streams(c)) | |
| assert raw["regionpca256"].shape == (8, 512) and raw["regionpca256_ok"].shape == (8,) | |
| def test_kclip_and_input_ln_contracts(): | |
| """coordinator 2026-09-11: new runs clamp at 8 (recorded), old cfgs (no kclip) keep 5; the input LayerNorm is opt-in and parameter-free when off.""" | |
| import batch as bt | |
| from models.branches import DriverBranch | |
| assert bt.K_CLIP == 8.0 and bt.K_CLIP_LEGACY == 5.0 and bt.k_clip({}) == 5.0 and bt.k_clip({"kclip": 8}) == 8.0 | |
| sys.argv = ["x"] | |
| import train_cwm3 as t | |
| for v, clip, ln in (("kin_can_noK", 8.0, False), ("kin_can_noK_clip8", 8.0, False), ("kin_can_noK_clip8_ln", 8.0, True)): | |
| c = t.stage_cfg("U0", v) | |
| assert float(c["recipe"].get("kclip", bt.K_CLIP)) == clip and t.model_kwargs(c)["in_ln"] is ln and c["lambda"] == {"K": 0.0} and c["recipe"]["batch"] == 4096 | |
| torch.manual_seed(0) | |
| u0 = DriverBranch(cabin="none"); u1 = DriverBranch(cabin="none", in_ln=True) | |
| n = lambda m: sum(p.numel() for p in m.parameters()) | |
| assert n(u0) == 159_290 and n(u1) - n(u0) == 2 * 12 # the U0 noK parameter count (train.log of every noK run); LN(12) affine | |
| b = {"Kx": torch.randn(8, 40, 18), "mode": torch.zeros(8, dtype=torch.long)} | |
| assert u1(b)["U"].shape == (8, 25, 26) | |
| # scale invariance of the LN path: scaling the 12 kinematics channels by 3 leaves the LN output unchanged (up to eps), the CAN channels untouched | |
| kx2 = b["Kx"].clone(); kx2[..., :12] *= 3.0 | |
| z1 = torch.cat([u1.in_ln(b["Kx"][..., :12]), b["Kx"][..., 12:]], -1); z2 = torch.cat([u1.in_ln(kx2[..., :12]), kx2[..., 12:]], -1) | |
| assert torch.allclose(z1, z2, atol=1e-3) | |
| def test_p_head_contracts(): | |
| """coordinator 2026-09-11 (plan §3 road-pressure output): E0 `*_P` variants emit the six P chains beside the D chains; loss parts chain_P_* enter the | |
| selection objective; the pressure stream carries P/Pvalid/P0/P0valid; a variant without p_head is bitwise the old module.""" | |
| import batch as bt | |
| from models import heads as hd | |
| from models.branches import EnvBranch | |
| sys.argv = ["x"] | |
| import train_cwm3 as t | |
| c = t.stage_cfg("E0", "pooled_P") | |
| assert c["p_head"] is True and c["lambda"] == {"F": 0.1, "P": 1.0} and c["recipe"]["select_LF"] is False and "pressure" in t.need_streams(c) | |
| c2 = t.stage_cfg("E0", "regionpca256_P") | |
| assert t.model_kwargs(c2) == {**t.model_kwargs(c), "d_side": 512} and c2["side"] == "regionpca256" | |
| torch.manual_seed(0) | |
| e0 = EnvBranch(tokens="pooled", n_slots=16); e1 = EnvBranch(tokens="pooled", n_slots=16, p_head=True) | |
| n = lambda m: sum(p.numel() for p in m.parameters()) | |
| assert n(e0) == 880_604 and n(e1) - n(e0) == (128 + 128 + 16) * 256 + 256 + 256 * 25 * 12 + 25 * 12 # flat_head(d_z=272 -> 256 -> 300) | |
| B = 8 | |
| b = {"tok": torch.randn(B, 16, 1, 256), "tok_ok": torch.ones(B, 16, dtype=torch.bool), "tok_dt": -torch.arange(16).flip(0).float()[None].expand(B, -1) * 0.5, | |
| "vf": torch.randn(B, 256), "vf_ok": torch.ones(B, dtype=torch.bool), "mode": torch.zeros(B, dtype=torch.long)} | |
| out = e1(b) | |
| assert out["P"].shape == (B, 25, 12) and "P" not in e0(b) | |
| raw = ds.synthetic_raw(B, need=t.need_streams(c)) | |
| assert raw["P"].shape == (B, 25, 6) and raw["P0valid"].shape == (B, 6) | |
| st = bt.compute_stats(raw, {"mu": np.zeros(3, np.float32), "sd": np.ones(3, np.float32)}) | |
| assert "p0_base" in st and st["p0_base"].shape == (6,) | |
| store = bt.GpuStore(raw, bt.KEYS["E"], "cpu", ds.synthetic_bank()) | |
| bb, tg = bt.prep_E(store, torch.arange(B), st, c) | |
| assert tg["zP"].shape == (B, 25, 6) and tg["zPm"].dtype == torch.bool and tg["zP_prev"].shape == (B, 25, 6) | |
| loss, parts = bt.LOSS["E"](e1(bb), tg, c["lambda"]) | |
| assert all(f"chain_P_{k}" in parts for k in hd.P_CHAINS) and torch.isfinite(loss) | |
| sel_d_only = sum(float(v) for k, v in parts.items() if k.startswith("chain_") and not k.startswith("chain_P_")) | |
| assert hd.selection_objective(parts, False) > sel_d_only # P chains are part of the selection objective | |
Xet Storage Details
- Size:
- 7.84 kB
- Xet hash:
- 57c6b785ff3b816d1fbee4d8fc4ca23b6bcedc1d5246df76687854a826d2afe3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.