Buckets:
| """cwm3/tests/test_identity_invariance.py — models/identity.py contracts: gradient reversal, paired init of the opt-in heads (baseline parameters | |
| unchanged), CAUSALITY of the per-device normalisation (statistics of a row never depend on later rows of its route, nor on other routes), | |
| prep_U_devnorm changes only the 11 kinematics input channels, and the state-supervision targets / loss shapes.""" | |
| 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)) | |
| from models import identity as idn # noqa: E402 | |
| from models import build_branch # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| import cwm3_dataset as ds # noqa: E402 | |
| def test_grad_reverse_flips_gradient(): | |
| x = torch.randn(4, 3, requires_grad=True) | |
| y = idn.grad_reverse(x, 0.3) | |
| assert torch.equal(y, x) | |
| y.sum().backward() | |
| assert torch.allclose(x.grad, torch.full_like(x, -0.3)) | |
| def test_adv_head_keeps_paired_init_and_emits_logits(): | |
| base = build_branch("U", 0, hid=128, hidden=128, dropout=0.0, cabin="none", in_ln=True) | |
| arm = build_branch("U", 0, hid=128, hidden=128, dropout=0.0, cabin="none", in_ln=True, adv_devices=96) | |
| sb, sa = base.state_dict(), arm.state_dict() | |
| assert set(sa) - set(sb) == {"adv.weight", "adv.bias"} | |
| assert all(torch.equal(sb[k], sa[k]) for k in sb) | |
| b = {"Kx": torch.randn(5, 40, 18), "mode": torch.zeros(5, dtype=torch.long)} | |
| out = arm(b) | |
| assert out["ADV"].shape == (5, 96) and torch.equal(out["U"], base(b)["U"]) | |
| l = idn.adv_loss(out["ADV"], torch.tensor([0, 5, -1, -1, 95])) | |
| assert torch.isfinite(l) and float(idn.adv_loss(out["ADV"], torch.full((5,), -1))) == 0.0 | |
| def test_f_state_head_keeps_paired_init(): | |
| kw = dict(tokens="pooled", n_slots=16, d=128, n_layers=2, n_heads=4, d_ff=256, hidden=256, dropout=0.1, p_head=True) | |
| base, arm = build_branch("E", 1, **kw), build_branch("E", 1, **kw, f_state=len(idn.FS_NAMES)) | |
| sb, sa = base.state_dict(), arm.state_dict() | |
| assert set(sa) - set(sb) == {"f_state.weight", "f_state.bias"} and all(torch.equal(sb[k], sa[k]) for k in sb) | |
| def _toy_raw(seed=0): | |
| rng = np.random.default_rng(seed) | |
| N = 300 | |
| uid = np.array([f"r{i % 3}" for i in range(N)]) | |
| mode = np.where(np.arange(N) % 2 == 0, 0, 1).astype(np.int8) | |
| t = np.zeros(N) | |
| for g in set(zip(uid, mode)): | |
| m = (uid == g[0]) & (mode == g[1]) | |
| t[m] = np.arange(m.sum()) * 0.2 + rng.uniform(0, 0.05) | |
| Kh = rng.normal(0, 1, (N, 40, 11)).astype(np.float32) + rng.normal(0, 3, (N, 1, 11)).astype(np.float32) * 0 # iid rows | |
| Kh[rng.random((N, 40)) < 0.05] = np.nan | |
| return {"uid": uid, "mode": mode, "t": t, "Kh": Kh.astype(np.float16), "uid_code": np.unique(uid, return_inverse=True)[1]} | |
| def test_devnorm_is_causal_and_route_local(): | |
| raw = _toy_raw() | |
| st = {"kmu": np.zeros(11, np.float32), "ksd": np.ones(11, np.float32)} | |
| idn.add_devnorm(raw, st, n0=5.0, sd_floor=0.2) | |
| mu0, sd0 = raw["Kdn_mu"].copy(), raw["Kdn_sd"].copy() | |
| assert mu0.shape == sd0.shape == (300, 11) and np.isfinite(mu0).all() and (sd0 > 0).all() | |
| # perturb the LATEST 30 % of rows of route r1 / mode 0: earlier rows of that route and every other route must be unchanged | |
| m = (raw["uid"] == "r1") & (raw["mode"] == 0) | |
| tt = raw["t"][m] | |
| late = m & (raw["t"] > np.quantile(tt, 0.7)) | |
| raw2 = dict(raw) | |
| Kh2 = raw["Kh"].astype(np.float32).copy() | |
| Kh2[late] += 50.0 | |
| raw2["Kh"] = Kh2.astype(np.float16) | |
| idn.add_devnorm(raw2, st, n0=5.0, sd_floor=0.2) | |
| early_or_other = ~late | |
| assert np.allclose(raw2["Kdn_mu"][early_or_other], mu0[early_or_other]) and np.allclose(raw2["Kdn_sd"][early_or_other], sd0[early_or_other]) | |
| assert not np.allclose(raw2["Kdn_mu"][late], mu0[late]) | |
| # the first row of a route equals the prior when its own anchor kinematics are missing; a huge pseudo-count returns the pool statistics | |
| idn.add_devnorm(raw, st, n0=1e9, sd_floor=0.0) | |
| assert np.allclose(raw["Kdn_mu"], 0.0, atol=1e-6) and np.allclose(raw["Kdn_sd"], 1.0, atol=1e-6) | |
| def test_devnorm_expanding_mean_matches_bruteforce(): | |
| raw = _toy_raw(3) | |
| kmu, ksd = np.full(11, 0.5, np.float32), np.full(11, 2.0, np.float32) | |
| mu, sd = idn.causal_devnorm(raw["Kh"][:, -1, :].astype(np.float32), raw["uid_code"], raw["mode"].astype(np.int64), raw["t"], kmu, ksd, n0=3.0, sd_floor=0.0) | |
| K0 = raw["Kh"][:, -1, :].astype(np.float64) | |
| for i in np.random.default_rng(0).choice(len(K0), 20, replace=False): | |
| same = (raw["uid_code"] == raw["uid_code"][i]) & (raw["mode"] == raw["mode"][i]) & (raw["t"] <= raw["t"][i]) | |
| k = K0[same] | |
| fin = np.isfinite(k) | |
| n = fin.sum(0) | |
| S = np.where(fin, k, 0).sum(0) | |
| S2 = np.where(fin, k ** 2, 0).sum(0) | |
| m_ref = (3.0 * 0.5 + S) / (3.0 + n) | |
| v_ref = (3.0 * (4.0 + 0.25) + S2) / (3.0 + n) - m_ref ** 2 | |
| assert np.allclose(mu[i], m_ref, atol=1e-4) and np.allclose(sd[i], np.sqrt(np.maximum(v_ref, 0)) + 1e-6, atol=1e-4) | |
| def test_prep_U_devnorm_changes_only_kinematics_channels(): | |
| raw = ds.synthetic_raw(256, seed=0, mode=0, need={"X", "Kh", "kin"}) | |
| raw["uid_code"] = np.unique(raw["uid"], return_inverse=True)[1] | |
| st = bt.compute_stats(raw, {"mu": np.zeros(3, np.float32), "sd": np.ones(3, np.float32)}) | |
| idn.add_devnorm(raw, st, n0=50, sd_floor=0.2) | |
| cfg = {"branch": "U", "recipe": {"kclip": 8, "u_devnorm": True}, "kclip": 8.0, "cabin": "none"} | |
| store = bt.GpuStore(raw, bt.KEYS["U"] + idn.store_keys(cfg), "cpu") | |
| idx = torch.arange(64) | |
| b0, _ = bt.prep_U(store, idx, st, cfg, need_targets=False) | |
| b1, _ = idn.prep_for(cfg)(store, idx, st, cfg, need_targets=False) | |
| assert b1["Kx"].shape == (64, 40, 18) | |
| assert torch.equal(b0["Kx"][..., 11:], b1["Kx"][..., 11:]) # valid bit, CAN, Δt untouched | |
| assert not torch.allclose(b0["Kx"][..., :11], b1["Kx"][..., :11]) | |
| assert b1["Kx"][..., :11].abs().max() <= 8.0 and torch.isfinite(b1["Kx"]).all() | |
| assert idn.prep_for({"branch": "U", "recipe": {}}) is bt.prep_U and idn.store_keys({"recipe": {}}) == [] | |
| def test_fs_targets_and_loss_shapes(): | |
| B = 16 | |
| store = type("S", (), {})() | |
| store.t = {"M": torch.ones(B, 25, dtype=torch.bool), "Pvalid": torch.rand(B, 25, 6) > 0.2, "R": torch.randint(0, 3, (B, 25)), "Rvalid": torch.ones(B, 25, dtype=torch.bool), | |
| "Y": torch.rand(B, 25, 5)} | |
| tg = {"zP": (torch.rand(B, 25, 6) > 0.9).float()} | |
| y, m = idn.fs_targets(store, torch.arange(B), tg) | |
| assert y.shape == (B, 5, 10) and m.shape == (B, 5, 10) and m.dtype == torch.bool | |
| assert torch.equal(y[..., 6], (store.t["R"][:, idn.FS_STEPS] >= 1).float()) and torch.equal(y[..., 8], (store.t["Y"][:, idn.FS_STEPS, 2] > 0.5).float()) | |
| l = idn.fs_loss(torch.zeros(B, 5, 10), y, m) | |
| assert torch.isclose(l, torch.tensor(np.log(2.0), dtype=l.dtype), atol=1e-5) | |
Xet Storage Details
- Size:
- 7.08 kB
- Xet hash:
- d29d4b38584edc07305f366e1b4f8a51a3c77005e93f2d4cf94c56e70a30809d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.