Buckets:
| """cwm3/probes/dump_latents.py — dump the TEMPORAL trunk latents s_t of the seed-0 base runs for the probe rows (read-only w.r.t. runs). | |
| Loads each BEST checkpoint through models/cwm3.load_base_branch (the run's own cfg + TRAIN-only standardisation statistics), builds the | |
| model inputs with train/batch.py prep_* on a CPU/GPU store (exactly the trainer's code path, need_targets=False) and stores the trunk | |
| output that the emission head reads: | |
| s_V4 f32[n,256] V0_can4s_s0 TCN(X 4 s) output before the mode embedding | |
| s_V8 f32[n,256] V0_can8s_s0 TCN(X8 8 s) output (NaN where x8_ok is False) | |
| s_U f32[n,128] U0_kin_can_noK_s0 GRU(Kx)[-1] | |
| s_E f32[n,128] E0_pooled_s0 LN(readout) of the causal token Transformer; s_Evf f32[n,128] = Linear(LN(vf256)) * vf_ok | |
| F_hat f16[n,5,256] E0_pooled_s0 the F head's predicted future front latent (+1..+5 s), read from [s_E, s_Evf, mode] | |
| chains: pE_onset3 f32[n,6] (E demand chains onset prob within 3 s), pU_onset3 f32[n,4] (hand_return, hand_leave, head_turn, head_return) | |
| -> results/probes/latents_{regime}_{split}.npz. GPU is used when >= 8 GB VRAM is free (coordinator rule), else CPU. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| import common as C # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| from models import heads as hd # noqa: E402 | |
| from models.cwm3 import load_base_branch # noqa: E402 | |
| RUNS = {"V4": "V0_can4s_s0", "V8": "V0_can8s_s0", "U": "U0_kin_can_noK_s0", "E": "E0_pooled_s0"} | |
| LOG = C.OUT / "dump_latents.log" | |
| def log(m): | |
| C.log(m, LOG) | |
| def bank_for_E(cfg): | |
| """the E base's token stream for the store (pooled256 slots for E0_pooled; region tokens otherwise) as a lightweight dict.""" | |
| want = "pooled256" if cfg.get("tokens") == "pooled" else "tok" | |
| z = C.npz_memmap(C.BANKS / "front_tok.npz", want) | |
| t = np.asarray(C.npz_memmap(C.BANKS / "front_tok.npz", "t")) | |
| return {want: np.asarray(z), "t": t} | |
| def trunk_E(model, b): | |
| E = model | |
| tok = b["tok"].float() | |
| B, L, R, _ = tok.shape | |
| x = E.tok_in(tok) | |
| miss = (~b["tok_ok"].bool())[:, :, None, None].float() | |
| x = x * (1 - miss) + E.missing * miss | |
| x = x + E.slot_emb.weight[None, :, None, :] + E.reg_emb.weight[None, None, :, :] | |
| x = x + E.dt_in(b["tok_dt"].float()[:, :, None, None]) | |
| x = x.reshape(B, L * R, -1) | |
| x = torch.cat([x, E.readout[None, None].expand(B, 1, -1)], 1) | |
| h = E.enc(x, mask=E.attn_mask) | |
| z_e = E.norm(h[:, -1]) | |
| z_vf = E.vf(b["vf"].float()) * b["vf_ok"].float()[:, None] | |
| z = torch.cat([z_e, z_vf, E.mode(b["mode"].long())], 1) | |
| return z_e, z_vf, E.head(z).view(-1, hd.H, hd.D_E), E.f_head(z).view(-1, 5, hd.D_F) | |
| def onset3(h_on, z0, tau=15): | |
| """closed-form onset probability within tau steps from the on-hazards: 1 - prod(1 - h_on[:tau]) (inactive anchors).""" | |
| return 1 - torch.prod(1 - h_on[:, :tau], 1) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--sets", default="all") | |
| a = ap.parse_args() | |
| dev = C.torch_device() | |
| log(f"device {dev}") | |
| models = {} | |
| for k, run in RUNS.items(): | |
| m, cfg, st, best = load_base_branch(C.RUNS / run, "cpu") | |
| m = m.to(dev).eval() | |
| models[k] = (m, cfg, {kk: np.asarray(v, np.float32) for kk, v in st.items()}, best) | |
| log(f"loaded {run}/{best} branch={cfg['branch']} variant={cfg['variant']} params={sum(p.numel() for p in m.parameters()):,}") | |
| bankE = bank_for_E(models["E"][1]) | |
| sets = [(r, s) for r in C.REGIMES for s in C.SPLITS] if a.sets == "all" else [tuple(x.split("/")) for x in a.sets.split(",")] | |
| for regime, split in sets: | |
| out = C.OUT / f"latents_{regime}_{split}.npz" | |
| if out.exists(): | |
| log(f"{out.name} exists — skipped") | |
| continue | |
| t0 = time.time() | |
| raw = C.load_raw(regime, split) | |
| n = len(raw["t"]) | |
| res = {"s_V4": np.zeros((n, 256), np.float32), "s_V8": np.full((n, 256), np.nan, np.float32), "s_U": np.zeros((n, 128), np.float32), | |
| "s_E": np.zeros((n, 128), np.float32), "s_Evf": np.zeros((n, 128), np.float32), "F_hat": np.zeros((n, 5, 256), np.float16), | |
| "pE_onset3": np.zeros((n, 6), np.float32), "pU_onset3": np.zeros((n, 4), np.float32), "pV4_onset3": np.zeros((n, 4), np.float32)} | |
| bs = 4096 | |
| # ---- V (4 s and 8 s) | |
| for key in ("V4", "V8"): | |
| m, cfg, st, _ = models[key] | |
| store = bt.GpuStore(raw, bt.KEYS["V"], dev) | |
| valid = raw["x8_ok"].astype(bool) if cfg.get("can_hist") == "8s" else np.ones(n, bool) | |
| for i in range(0, n, bs): | |
| idx = torch.arange(i, min(i + bs, n), device=dev) | |
| b, _ = bt.prep_V(store, idx, st, cfg, need_targets=False) | |
| with torch.no_grad(): | |
| z = m.tcn(b["X"].float()) | |
| if key == "V4": | |
| V = m.head(torch.cat([z, m.mode(b["mode"].long())], 1)).view(-1, hd.H, hd.D_V) | |
| h_on, _ = hd.hazards(hd.split_V(V)["chains"]) | |
| res["pV4_onset3"][i:i + len(idx)] = onset3(h_on.permute(0, 2, 1).reshape(-1, hd.H, 1)[:, :, 0].reshape(len(idx), 4, hd.H).permute(0, 2, 1), None).cpu().numpy() if False else \ | |
| (1 - torch.prod(1 - h_on[:, :15, :], 1)).cpu().numpy() | |
| res["s_" + key][i:i + len(idx)] = z.cpu().numpy() | |
| res["s_" + key][~valid] = np.nan | |
| del store | |
| # ---- U | |
| m, cfg, st, _ = models["U"] | |
| store = bt.GpuStore(raw, bt.KEYS["U"], dev) | |
| for i in range(0, n, bs): | |
| idx = torch.arange(i, min(i + bs, n), device=dev) | |
| b, _ = bt.prep_U(store, idx, st, cfg, need_targets=False) | |
| with torch.no_grad(): | |
| _, h = m.gru(b["Kx"].float()) | |
| hu = h[-1] | |
| U = m.head(torch.cat([hu, m.mode(b["mode"].long())], 1)).view(-1, hd.H, hd.D_U) | |
| h_on, h_off = hd.hazards(hd.split_U(U)["chains"]) | |
| p = torch.stack([1 - torch.prod(1 - h_on[:, :15, 0], 1), 1 - torch.prod(1 - h_off[:, :15, 0], 1), | |
| 1 - torch.prod(1 - h_on[:, :15, 1], 1), 1 - torch.prod(1 - h_off[:, :15, 1], 1)], 1) | |
| res["s_U"][i:i + len(idx)] = hu.cpu().numpy() | |
| res["pU_onset3"][i:i + len(idx)] = p.cpu().numpy() | |
| del store | |
| # ---- E | |
| m, cfg, st, _ = models["E"] | |
| store = bt.GpuStore(raw, bt.KEYS["E"], dev, bankE) | |
| for i in range(0, n, bs): | |
| idx = torch.arange(i, min(i + bs, n), device=dev) | |
| b, _ = bt.prep_E(store, idx, st, cfg, need_targets=False) | |
| z_e, z_vf, E, Fh = trunk_E(m, b) | |
| h_on, _ = hd.hazards(hd.split_E(E)["chains"]) | |
| res["s_E"][i:i + len(idx)] = z_e.cpu().numpy() | |
| res["s_Evf"][i:i + len(idx)] = z_vf.cpu().numpy() | |
| res["F_hat"][i:i + len(idx)] = Fh.cpu().numpy().astype(np.float16) | |
| res["pE_onset3"][i:i + len(idx)] = (1 - torch.prod(1 - h_on[:, :15, :], 1)).cpu().numpy() | |
| del store | |
| # sanity: F cosine against the stored future on Fok rows should match the run's metrics_quick f_cos means (same rows subset) | |
| F = raw["F"].astype(np.float32) | |
| Fok = raw["Fok"].astype(bool) | |
| cos = (res["F_hat"].astype(np.float32) * F).sum(-1) / (np.linalg.norm(res["F_hat"].astype(np.float32), axis=-1) * np.linalg.norm(F, axis=-1) + 1e-6) | |
| fcos = [float(np.nanmean(np.where(Fok[:, k], cos[:, k], np.nan))) for k in range(5)] | |
| np.savez(out, **res, f_cos_check=np.array(fcos), runs=np.array([f"{RUNS[k]}:{models[k][3]}" for k in RUNS]), device=np.array(str(dev))) | |
| log(f"{out.name}: n={n:,} f_cos(+1..+5)={np.round(fcos, 3).tolist()} {time.time() - t0:.0f}s") | |
| log("DONE") | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.05 kB
- Xet hash:
- 72e747981d458d3e39f2d61977df18bc8df4dc506c8c2f125fe62863f7bb2dea
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.