Buckets:
| """cwm3/tests/test_trunk_head.py — Experiment B helpers (train/train_trunk_warning_head.py; pre-registration docs/decisions.md 2026-09-18T03:34:28Z). | |
| (a) the pre-registered head = train_warning_head.make_model('mlp') on 641 inputs has EXACTLY 641*64 + 64 + 64 + 1 = 41,153 parameters (<= 50,000) after a fit; | |
| (b) assemble(): f16 trunk block [N, 640] + regime bit -> f32 [N, 641] with the bit last and the trunk order preserved; wrong widths are rejected; | |
| (c) posmap_rows(): label-row positions map back to the targets-file rows; a missing row raises; | |
| (d) write_layout(): the W-layout contract (meta json, uid, t, w_names, w f16, w_mask all True) round-trips through np.load; | |
| (e) the extraction feature order equals CWM3Joint's trunk keys (z_v, h_u, z_e, z_vf) and their widths sum to 640 on a synthetic joint model; | |
| (f) no forbidden openpilot DMS identifier in the new files. | |
| Run: cd cwm3 && python -m pytest tests/test_trunk_head.py -q | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import re | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import pytest # noqa: E402 | |
| import torch # noqa: E402 | |
| torch.set_num_threads(4) | |
| CWM3 = Path(__file__).resolve().parents[1] | |
| for p in (CWM3, CWM3 / "data", CWM3 / "train", CWM3 / "tests"): | |
| sys.path.insert(0, str(p)) | |
| from models.cwm3 import CWM3Joint # noqa: E402 (before train_warning_head: its cwm2 sys.path entries shadow the `models` package) | |
| from test_no_dms_inputs import FORBIDDEN # noqa: E402 | |
| from test_zero_init_identity import _batch, _fresh_branches # noqa: E402 | |
| import train_trunk_warning_head as TT # noqa: E402 | |
| import train_warning_head as TW # noqa: E402 | |
| def test_mlp_head_parameter_count(): | |
| rng = np.random.default_rng(0) | |
| X = rng.normal(size=(200, TT.D_FEAT + 1)).astype(np.float32) | |
| y = (rng.random(200) < 0.3).astype(int) | |
| m = TW.make_model("mlp", {"alpha": 1e-2}, 0).fit(X, y) | |
| n = TT.mlp_param_count(m[-1]) | |
| assert n == 641 * 64 + 64 + 64 + 1 == 41153 and n <= TT.MAX_PARAMS | |
| assert m[-1].hidden_layer_sizes == (64,) | |
| def test_assemble_order_and_dtype(): | |
| feat = (np.arange(3 * TT.D_FEAT, dtype=np.float32).reshape(3, TT.D_FEAT) / 1000).astype(np.float16) | |
| X = TT.assemble(feat, np.array([0, 1, 1], np.int8)) | |
| assert X.shape == (3, TT.D_FEAT + 1) and X.dtype == np.float32 | |
| assert np.array_equal(X[:, -1], [0, 1, 1]) and np.allclose(X[:, :TT.D_FEAT], feat.astype(np.float32)) | |
| with pytest.raises(AssertionError): | |
| TT.assemble(np.zeros((3, 10), np.float16), np.zeros(3)) | |
| def test_posmap_rows(): | |
| cache_rows = np.array([15, 63, 101, 148, 865]) | |
| assert TT.posmap_rows(cache_rows, [101, 15]).tolist() == [2, 0] | |
| with pytest.raises(KeyError): | |
| TT.posmap_rows(cache_rows, [7]) | |
| def test_write_layout_contract(tmp_path): | |
| uid = np.array(["r1", "r1", "r2"], dtype=object) | |
| t = np.array([0.2, 0.4, 1.0]) | |
| p = TT.write_layout(tmp_path / "pred_layout_manual_test.npz", uid, t, {TT.COLS["mlp"]: np.array([0.1, 0.5, 0.9]), TT.COLS["lr"]: np.array([0.2, 0.2, 0.2])}, {"arm": "abl_W_trk"}) | |
| with np.load(p, allow_pickle=True) as z: | |
| assert list(map(str, z["w_names"])) == [TT.COLS["mlp"], TT.COLS["lr"]] | |
| assert z["w"].dtype == np.float16 and z["w"].shape == (3, 2) and z["w_mask"].all() and len(z["uid"]) == 3 | |
| assert json.loads(str(z["meta"]))["arm"] == "abl_W_trk" | |
| assert abs(float(z["w"][1, 0]) - 0.5) < 1e-3 | |
| def test_trunk_keys_match_model_widths(): | |
| V, U, E = _fresh_branches() | |
| b, _ = _batch() | |
| m = CWM3Joint(V, U, E, interactions=("env2veh", "drv2veh"), decoder="direct", k_head=True, seed=0).eval() | |
| with torch.no_grad(): | |
| tr = m(b, return_trunks=True)["trunks"] | |
| assert all(k in tr for k in TT.TRUNK_KEYS) | |
| assert sum(tr[k].shape[1] for k in TT.TRUNK_KEYS) == TT.D_FEAT == 256 + 128 + 128 + 128 | |
| cat = torch.cat([tr[k].float() for k in TT.TRUNK_KEYS], 1) | |
| assert cat.shape == (b["V"]["X"].shape[0], TT.D_FEAT) | |
| def test_no_forbidden_tokens(): | |
| pat = re.compile("|".join(re.escape(t) for t in FORBIDDEN)) | |
| for f in ("train/train_trunk_warning_head.py", "tests/test_trunk_head.py"): | |
| assert not pat.search((CWM3 / f).read_text()), f | |
Xet Storage Details
- Size:
- 4.53 kB
- Xet hash:
- 0ed33cb14333845f9ae0f5f8ab1d1501104fd547549c3203ee8d481678feb691
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.