#!/usr/bin/env python3 """Sprints 3-5 on Tier 0: transplant families, full control battery, sweeps, mechanism tests. Experiment id ``graft_tier0``. Requires the model matrix from scripts/train_models.py. Usage: python scripts/run_transplant.py [--quick] [--device auto] [--seed-base 30000] [--tag main|confirmatory] Every condition is evaluated on the same held-out splits; per-item outputs are saved; CIs are cluster bootstraps over programs; control-adjusted effects are paired on identical items. """ from __future__ import annotations import argparse import gzip import json import sys import time from pathlib import Path import numpy as np import torch ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT / "src")) sys.path.insert(0, str(ROOT / "scripts")) from train_models import MODELS_DIR, ROLES, model_dir # noqa: E402 from wp_ai_0001.alignment.maps import FITTERS, AffineMap, fit_ridge # noqa: E402 from wp_ai_0001.config import OpSeqTaskConfig, canonical_json # noqa: E402 from wp_ai_0001.datasets import opseq # noqa: E402 from wp_ai_0001.datasets.common import Split, assert_disjoint # noqa: E402 from wp_ai_0001.evaluation.metrics import chance_level # noqa: E402 from wp_ai_0001.instrumentation.provenance import Manifest, make_run_id, sha256_text, utc_now # noqa: E402 from wp_ai_0001.models.hooks import Zero # noqa: E402 from wp_ai_0001.models.toy_transformer import resolve_device # noqa: E402 from wp_ai_0001.orchestration import registry # noqa: E402 from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean, cluster_bootstrap_paired_diff, seeds_summary # noqa: E402 from wp_ai_0001.training.train_toy import load as load_model # noqa: E402 from wp_ai_0001.transplantation.experiment import ( # noqa: E402 MAX_SEQ, OMEGA, build_eval_splits, collect_rows, distill, evaluate_condition, fit_adapter_from_donor, omega_step_acc, train_supervised_adapter, ) from wp_ai_0001.transplantation.grafts import ( # noqa: E402 ActivationGraft, AdapterGraft, GraftSpec, NormMatchedRandomGraft, SteeringGraft, SubspaceGraft, copy_embedding_row, weight_graft, ) EVAL = ["target_iid", "control_iid", "comp_heldout"] class Exp: def __init__(self, args): self.args = args self.device = resolve_device(args.device) self.task = ROLES["donor"][0] self.vocab = opseq.vocab_for(self.task) self.n_boot = 300 if args.quick else 2000 self.models: dict[str, object] = {} self.results: dict = {"conditions": {}, "pairs": {}, "notes": {}, "unsupported": {}} self.items: dict = {} # ---------- models ---------- def model(self, key: str): if key not in self.models: m, v = load_model(self._dir(key), self.device) assert v.itos == self.vocab.itos, f"vocab mismatch for {key}" m.eval() self.models[key] = m return self.models[key] @staticmethod def _dir(key: str) -> Path: role, seed = key.rsplit("_s", 1) return model_dir(role, int(seed)) # ---------- evaluation bookkeeping ---------- def record(self, cond: str, pair: str, res: dict, extra: dict | None = None) -> None: rec = {"pair": pair, **(extra or {})} for n, r in res.items(): om = omega_step_acc(self.splits[n], r) valid = ~np.isnan(om) rec[n] = { "accuracy": r.accuracy, "trace_accuracy": r.trace_accuracy, "ci95": cluster_bootstrap_mean(r.correct.astype(float), r.clusters, self.n_boot, seed=0).to_dict(), "omega_step_acc": float(om[valid].mean()) if valid.any() else None, } self.items[f"{cond}|{pair}|{n}"] = {"correct": r.correct, "trace": r.trace_correct, "omega": om, "clusters": r.clusters, "pred": r.predictions} self.results["conditions"].setdefault(cond, {})[pair] = rec def paired(self, cond_a: str, cond_b: str, pair: str, n: str = "target_iid", metric: str = "correct") -> dict | None: ka, kb = f"{cond_a}|{pair}|{n}", f"{cond_b}|{pair}|{n}" if ka not in self.items or kb not in self.items: return None a, b = self.items[ka], self.items[kb] if metric == "omega": va, vb = a["omega"], b["omega"] ok = ~np.isnan(va) & ~np.isnan(vb) cl = [c for c, o in zip(a["clusters"], ok) if o] return cluster_bootstrap_paired_diff(va[ok], vb[ok], cl, self.n_boot, seed=1).to_dict() return cluster_bootstrap_paired_diff(a[metric].astype(float), b[metric].astype(float), a["clusters"], self.n_boot, seed=1).to_dict() # ---------- main ---------- def run(self): a = self.args n_eval, n_align = (300, 600) if a.quick else (1000, 3000) index = json.loads((MODELS_DIR / "index.json").read_text()) train_hashes: set[str] = set() for k, v in index.items(): sp = Split.load(self._dir(k) / "train_split.json") train_hashes |= sp.item_hashes() self.splits = build_eval_splits(self.task, train_hashes, a.seed_base, n_eval, n_align) assert_disjoint(*self.splits.values()) self.results["splits"] = {k: {"n": len(v), "hash": v.hash, "seed": v.seed} for k, v in self.splits.items()} self.results["chance"] = {k: chance_level(self.splits[k]) for k in EVAL} self.results["leakage_audit"] = {"alignment_vs_eval_overlap": 0, "train_vs_eval_overlap": 0, "n_train_hashes": len(train_hashes)} seeds = [0] if a.quick else [0, 1, 2] pairs = [(f"donor_s{i}", f"recipient_s{j}") for i in seeds for j in seeds] self.results["pair_meta"] = {f"{d}->{r}": {"same_init_seed": d.rsplit("_s", 1)[1] == r.rsplit("_s", 1)[1]} for d, r in pairs} matched = [(f"donor_s{i}", f"recipient_s{i}") for i in seeds] hook = "resid_post.0" t0 = time.time() # ---------- baselines ---------- for key in sorted({p for pr in pairs for p in pr} | ({f"unrelated_s{i}" for i in seeds})): self.record("base", key, evaluate_condition(self.model(key), None, None, self.splits, self.vocab, self.device, EVAL + ["len_ood"])) print(f"[{utc_now()}] baselines done ({time.time()-t0:.0f}s)") # ---------- maps on alignment data ---------- self.rows: dict[tuple[str, str], object] = {} def rows(key, h=hook, split="align_target"): k = (key, h, split) if k not in self.rows: self.rows[k] = collect_rows(self.model(key), self.splits[split], self.vocab, h, self.device) return self.rows[k] self.results["maps"] = {} maps: dict[tuple, AffineMap] = {} for d, r in pairs: X, Y = rows(d).X, rows(r).X Xu, Yu = rows(d, split="align_unrelated").X, rows(r, split="align_unrelated").X for fam in ("identity", "procrustes", "ridge", "cca", "random_orth"): m = FITTERS[fam](X, Y) if fam != "random_orth" else FITTERS[fam](X, Y, seed=0) maps[(d, r, fam)] = m self.results["maps"][f"{d}->{r}|{fam}"] = {"r2_fit": m.r2(X, Y), "r2_unrelated_heldout": m.r2(Xu, Yu), "n_params": m.n_params, **{k: v for k, v in m.info.items() if k != "fit_seconds"}} mu = fit_ridge(Xu, Yu, family="unrelated_data") maps[(d, r, "unrelated_data")] = mu mup = FITTERS["procrustes"](Xu, Yu) maps[(d, r, "unrelated_procrustes")] = mup self.results["maps"][f"{d}->{r}|unrelated_procrustes"] = {"r2_fit_unrelated": mup.r2(Xu, Yu), "r2_on_target_rows": mup.r2(X, Y), "n_params": mup.n_params} self.results["maps"][f"{d}->{r}|unrelated_data"] = {"r2_fit_unrelated": mu.r2(Xu, Yu), "r2_on_target_rows": mu.r2(X, Y), "n_params": mu.n_params} print(f"[{utc_now()}] maps fitted ({time.time()-t0:.0f}s)") # ---------- activation grafts: map family × targeting (all pairs) ---------- for d, r in pairs: pair = f"{d}->{r}" R = self.model(r) for fam in ("identity", "procrustes", "ridge", "cca"): for tg in ("all", "omega"): spec = GraftSpec("activation", hook, hook, 1.0, tg) g = ActivationGraft(spec, self.model(d), maps[(d, r, fam)]) self.record(f"act_{fam}_{tg}", pair, evaluate_condition(R, hook, g, self.splits, self.vocab, self.device, EVAL), {"map": fam, "targeting": tg, "applied": g.applied}) # controls spec = GraftSpec("activation", hook, hook, 1.0, "all") self.record("ctl_sham_map", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "random_orth")]), self.splits, self.vocab, self.device, EVAL)) self.record("ctl_unrelated_data_map", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "unrelated_data")]), self.splits, self.vocab, self.device, EVAL)) self.record("act_unrelated_procrustes_all", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "unrelated_procrustes")]), self.splits, self.vocab, self.device, EVAL), {"map": "procrustes fitted on non-Ω data"}) self.record("ctl_norm_matched_random", pair, evaluate_condition(R, hook, NormMatchedRandomGraft(spec, self.model(d), maps[(d, r, "ridge")], seed=0), self.splits, self.vocab, self.device, EVAL)) self.record("ctl_shuffled_token", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(d), maps[(d, r, "ridge")], shuffle_tokens=torch.Generator().manual_seed(0)), self.splits, self.vocab, self.device, EVAL)) # shuffled layer: donor layer 3 -> recipient layer 0 with its own ridge map X3 = rows(d, h="resid_post.3").X m3 = fit_ridge(X3, rows(r).X) self.record("ctl_shuffled_layer", pair, evaluate_condition(R, hook, ActivationGraft(GraftSpec("activation", hook, "resid_post.3", 1.0, "all"), self.model(d), m3), self.splits, self.vocab, self.device, EVAL)) # unrelated-skill donor (competent on vex, not on Ω) u = f"unrelated_s{r.rsplit('_s',1)[1]}" mu = fit_ridge(rows(u).X, rows(r).X) self.record("ctl_unrelated_donor", pair, evaluate_condition(R, hook, ActivationGraft(spec, self.model(u), mu), self.splits, self.vocab, self.device, EVAL)) print(f"[{utc_now()}] activation grafts + controls {pair} ({time.time()-t0:.0f}s)") # ---------- steering, subspace, adapters, distillation, weight grafts (matched seeds) ---------- for d, r in matched: pair = f"{d}->{r}" R, D = self.model(r), self.model(d) rd, rr = rows(d), rows(r) REF = "unrelated_procrustes" # reference map for secondary families (D7): orthogonal, fitted on task-free rows mapped = maps[(d, r, REF)](rd.X) # steering vector: mapped mean(Ω rows) - mean(non-Ω rows) of the donor v = mapped[rd.is_omega].mean(0) - mapped[~rd.is_omega].mean(0) for alpha in ([1.0] if a.quick else [0.5, 1.0, 2.0, 4.0]): self.record(f"steer_a{alpha}", pair, evaluate_condition(R, hook, SteeringGraft(GraftSpec("steering", hook, None, alpha, "all"), v), self.splits, self.vocab, self.device, EVAL), {"alpha": alpha}) # subspace grafts: PCA of (mapped donor - recipient) on Ω rows; norm-matched random subspace of the same k diff = (mapped - rr.X)[rd.is_omega] U, S, Vt = np.linalg.svd(diff - diff.mean(0), full_matrices=False) rng = np.random.default_rng(0) for k in ([4, 32] if a.quick else [1, 2, 4, 8, 16, 32, 64, 128]): basis = Vt[:k].T self.record(f"subspace_k{k}", pair, evaluate_condition(R, hook, SubspaceGraft(GraftSpec("subspace", hook, hook, 1.0, "all"), D, maps[(d, r, REF)], basis), self.splits, self.vocab, self.device, EVAL), {"k": k, "explained_var": float((S[:k] ** 2).sum() / (S ** 2).sum()), "map": REF}) rb = rng.standard_normal((basis.shape[0], k)) self.record(f"subspace_random_k{k}", pair, evaluate_condition(R, hook, SubspaceGraft(GraftSpec("subspace", hook, hook, 1.0, "all"), D, maps[(d, r, REF)], rb), self.splits, self.vocab, self.device, EVAL), {"k": k, "map": REF}) # adapter graft (donor-free at inference): ridge from recipient rows to mapped donor rows ad = fit_adapter_from_donor(rr.X, mapped) self.record("adapter_graft", pair, evaluate_condition(R, hook, AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), ad), self.splits, self.vocab, self.device, EVAL), {"n_params": ad.n_params, "n_fit_items": len(self.splits["align_target"]), "targets": REF}) self.adapters = getattr(self, "adapters", {}) self.adapters[pair] = ad # matched-size supervised adapter (labels, no donor) on the same alignment items steps = 100 if a.quick else 300 sup = train_supervised_adapter(R, hook, self.splits["align_target"], self.vocab, self.device, steps=steps) self.record("baseline_supervised_adapter", pair, evaluate_condition(R, hook, AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), sup), self.splits, self.vocab, self.device, EVAL), {"n_params": sup.n_params, "steps": steps}) # distillation baseline (full recipient fine-tuned on donor distributions, same items) student = distill(R, D, self.splits["align_target"], self.vocab, self.device, steps=steps) self.record("baseline_distillation", pair, evaluate_condition(student, None, None, self.splits, self.vocab, self.device, EVAL), {"steps": steps, "n_params_trained": student.n_params}) del student # weight grafts (same architecture): copies into a fresh copy of the recipient import copy as _copy omega_id = self.vocab.stoi[OMEGA] for name, parts, src_key, src_layer in [ ("weight_L0_attn", ["attn"], d, 0), ("weight_L0_mlp", ["mlp"], d, 0), ("weight_L0_attn_mlp", ["attn", "mlp"], d, 0), ("weight_L0_block_ln", ["attn", "mlp", "ln"], d, 0), ("weight_embed_omega", [], d, 0), ("weight_L0_block_embed", ["attn", "mlp", "ln"], d, 0), ("weight_L012_blocks_embed", ["attn", "mlp", "ln"], d, 0), ("ctl_weight_shuffled_layer", ["attn", "mlp", "ln"], d, 3), ("ctl_weight_unrelated_donor", ["attn", "mlp", "ln"], f"unrelated_s{r.rsplit('_s',1)[1]}", 0), ]: Rc = _copy.deepcopy(R) rep = {"n_params_copied": 0} S = self.model(src_key) if parts: rep = weight_graft(Rc, S, parts, src_layer, 0) if name == "weight_L012_blocks_embed": for l in (1, 2): rep["n_params_copied"] += weight_graft(Rc, S, parts, l, l)["n_params_copied"] if name in ("weight_embed_omega", "weight_L0_block_embed", "weight_L012_blocks_embed"): rep["n_params_copied"] += copy_embedding_row(Rc, S, omega_id) self.record(name, pair, evaluate_condition(Rc, None, None, self.splits, self.vocab, self.device, EVAL), rep) del Rc print(f"[{utc_now()}] steering/subspace/adapter/weights {pair} ({time.time()-t0:.0f}s)") # ---------- sweeps (matched seed 0): layer grid, dose ---------- d, r = matched[0] pair = f"{d}->{r}" R, D = self.model(r), self.model(d) L = R.cfg.n_layers grid = {} for ld in range(L): for lr_ in range(L): hd, hr = f"resid_post.{ld}", f"resid_post.{lr_}" m = FITTERS["procrustes"](rows(d, h=hd, split="align_unrelated").X, rows(r, h=hr, split="align_unrelated").X) res = evaluate_condition(R, hr, ActivationGraft(GraftSpec("activation", hr, hd, 1.0, "all"), D, m), self.splits, self.vocab, self.device, ["target_iid", "control_iid"]) grid[f"{ld}->{lr_}"] = {"target": res["target_iid"].accuracy, "control": res["control_iid"].accuracy, "omega_step": float(np.nanmean(omega_step_acc(self.splits["target_iid"], res["target_iid"])))} self.results["layer_grid"] = {"pair": pair, "map": "unrelated_procrustes", "cells": grid} dose = {} for alpha in ([0.0, 0.5, 1.0] if a.quick else [0.0, 0.1, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 2.0]): res = evaluate_condition(R, hook, ActivationGraft(GraftSpec("activation", hook, hook, alpha, "all"), D, maps[(d, r, "unrelated_procrustes")]), self.splits, self.vocab, self.device, ["target_iid", "control_iid"]) dose[str(alpha)] = {"target": res["target_iid"].accuracy, "control": res["control_iid"].accuracy} self.results["dose"] = {"pair": pair, "map": "unrelated_procrustes", "cells": dose} print(f"[{utc_now()}] sweeps done ({time.time()-t0:.0f}s)") # ---------- mismatched pairs ---------- mis = {} for rk, layers in (("recipient_w96_s0", [(0, 0), (1, 1)]), ("recipient_L6_s0", [(0, 0), (0, 1), (1, 1), (1, 2), (2, 2)])): if not (self._dir(rk) / "model.pt").exists(): self.results["unsupported"][rk] = "model not trained" continue Rm = self.model(rk) base = evaluate_condition(Rm, None, None, self.splits, self.vocab, self.device, EVAL) self.record("base", rk, base) for ld, lr_ in layers: hd, hr = f"resid_post.{ld}", f"resid_post.{lr_}" for fam in ("procrustes", "ridge", "cca", "unrelated_procrustes"): if fam == "unrelated_procrustes": X, Y = rows(d, h=hd, split="align_unrelated").X, rows(rk, h=hr, split="align_unrelated").X m = FITTERS["procrustes"](X, Y) else: X, Y = rows(d, h=hd).X, rows(rk, h=hr).X m = FITTERS[fam](X, Y) self.record(f"mismatch_{fam}_L{ld}to{lr_}", f"{d}->{rk}", evaluate_condition(Rm, hr, ActivationGraft(GraftSpec("activation", hr, hd, 1.0, "all"), D, m), self.splits, self.vocab, self.device, EVAL), {"r2_fit": m.r2(X, Y)}) sham = FITTERS["random_orth"](rows(d, h=hd).X, rows(rk, h=hr).X, seed=0) self.record(f"mismatch_sham_L{ld}to{lr_}", f"{d}->{rk}", evaluate_condition(Rm, hr, ActivationGraft(GraftSpec("activation", hr, hd, 1.0, "all"), D, sham), self.splits, self.vocab, self.device, EVAL)) mis[rk] = "done" self.results["mismatch"] = mis # stride-2 pair (multi-step composition inside one prediction) if (self._dir("donor_s2_s0") / "model.pt").exists() and (self._dir("recipient_s2_s0") / "model.pt").exists(): task2 = ROLES["donor_s2"][0] ex2 = set() for k in ("donor_s2_s0", "recipient_s2_s0"): ex2 |= Split.load(self._dir(k) / "train_split.json").item_hashes() sp2 = build_eval_splits(task2, ex2, a.seed_base + 500, n_eval, n_align) saved = self.splits self.splits = sp2 D2, R2 = self.model("donor_s2_s0"), self.model("recipient_s2_s0") self.rows_s2 = {} self.record("base", "donor_s2_s0", evaluate_condition(D2, None, None, sp2, self.vocab, self.device, EVAL)) self.record("base", "recipient_s2_s0", evaluate_condition(R2, None, None, sp2, self.vocab, self.device, EVAL)) for l in range(L): h = f"resid_post.{l}" X = collect_rows(D2, sp2["align_target"], self.vocab, h, self.device).X Y = collect_rows(R2, sp2["align_target"], self.vocab, h, self.device).X m = fit_ridge(X, Y) self.record(f"stride2_act_ridge_L{l}", "donor_s2_s0->recipient_s2_s0", evaluate_condition(R2, h, ActivationGraft(GraftSpec("activation", h, h, 1.0, "all"), D2, m), sp2, self.vocab, self.device, EVAL), {"r2_fit": m.r2(X, Y)}) Xu = collect_rows(D2, sp2["align_unrelated"], self.vocab, h, self.device).X Yu = collect_rows(R2, sp2["align_unrelated"], self.vocab, h, self.device).X mp = FITTERS["procrustes"](Xu, Yu) self.record(f"stride2_act_unrelated_procrustes_L{l}", "donor_s2_s0->recipient_s2_s0", evaluate_condition(R2, h, ActivationGraft(GraftSpec("activation", h, h, 1.0, "all"), D2, mp), sp2, self.vocab, self.device, EVAL), {"r2_fit": mp.r2(Xu, Yu)}) sham = FITTERS["random_orth"](X, Y, seed=0) self.record(f"stride2_sham_L{l}", "donor_s2_s0->recipient_s2_s0", evaluate_condition(R2, h, ActivationGraft(GraftSpec("activation", h, h, 1.0, "all"), D2, sham), sp2, self.vocab, self.device, EVAL)) self.results["splits_stride2"] = {k: {"n": len(v), "hash": v.hash} for k, v in sp2.items()} self.splits = saved else: self.results["unsupported"]["stride2"] = "models not trained" print(f"[{utc_now()}] mismatched + stride-2 done ({time.time()-t0:.0f}s)") # ---------- mechanism: donor ablation -> loss -> rescue; bidirectional ---------- mech = {} d0, d1, r0 = matched[0][0], "donor_s1", matched[0][1] if not (self._dir(d1) / "model.pt").exists(): d1 = d0 self.results["unsupported"]["rescue"] = "intact donor_s1 missing: rescue source is the ablated model itself (invalid)" Dab = self.model(d0) zero = Zero(positions=None) # ablate attention output of layer 0 at ALL positions -> but only read positions matter # ablation implemented as hook on attn_out.0 restricted to read positions via positions=ctx.positions is not supported by Zero; # use the graft machinery: replace with zeros at read positions through a norm-0 "steering" of -x -> SteeringGraft cannot; so patch: from wp_ai_0001.models.hooks import HookContext class _ZeroAtRead: def __call__(self, x, ctx: HookContext): if ctx.positions is None: return x out = x.clone() out[torch.arange(x.shape[0], device=x.device), ctx.positions] = 0 return out zr = _ZeroAtRead() Dab.hooks.add("attn_out.0", zr) abl = evaluate_condition(Dab, None, None, self.splits, self.vocab, self.device, EVAL) self.record("mech_donor_ablated_attn0", d0, abl) # rescue: graft intact donor_s1's resid_post.0 (ridge-mapped) into the ablated donor_s0 at read positions Xr, Yr = rows(d1).X, rows(d0).X # NOTE: Yr collected on the intact d0 (hook added after) -> re-collect under ablation Yab = collect_rows(Dab, self.splits["align_target"], self.vocab, hook, self.device).X mres = FITTERS["procrustes"](rows(d1, split="align_unrelated").X, rows(d0, split="align_unrelated").X) # task-free Procrustes into the intact donor's coordinates res = evaluate_condition(Dab, hook, ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d1), mres), self.splits, self.vocab, self.device, EVAL) self.record("mech_rescue_from_intact_donor", f"{d1}->{d0}(ablated)", res) shamm = FITTERS["random_orth"](rows(d1, split="align_unrelated").X, rows(d0, split="align_unrelated").X, seed=0) res = evaluate_condition(Dab, hook, ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d1), shamm), self.splits, self.vocab, self.device, EVAL) self.record("mech_rescue_sham", f"{d1}->{d0}(ablated)", res) Dab.hooks.remove("attn_out.0", zr) # bidirectional: recipient -> donor graft (should not harm donor on control, should remove Ω on target if replacing) mrd = fit_ridge(rows(r0).X, rows(d0).X) res = evaluate_condition(self.model(d0), hook, ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(r0), mrd), self.splits, self.vocab, self.device, EVAL) self.record("mech_reverse_recipient_to_donor", f"{r0}->{d0}", res) self.results["mechanism"] = {"ablated_component": "attn_out.0 at read positions (donor_s0)", "rescue_source": d1, "yab_norm_mean": float(np.linalg.norm(Yab, axis=1).mean())} print(f"[{utc_now()}] mechanism done ({time.time()-t0:.0f}s)") # ---------- rule vs cache (H5): maps/adapters fitted on a subset of Ω input values ---------- d, r = matched[0] pair = f"{d}->{r}" rd, rr = rows(d), rows(r) seen_vals = {0, 1, 2, 3, 4, 5} keep = (~rd.is_omega) | np.isin(rd.omega_input, list(seen_vals)) m_seen = fit_ridge(rd.X[keep], rr.X[keep]) mp_seen = FITTERS["procrustes"](rd.X[keep], rr.X[keep]) ad_seen = fit_adapter_from_donor(rr.X[keep], maps[(d, r, "unrelated_procrustes")](rd.X[keep])) # partition target_iid by whether ALL Ω inputs are in the seen set from wp_ai_0001.transplantation.experiment import omega_input_values oin = omega_input_values(self.splits["target_iid"]) seen_mask = np.array([all(v in seen_vals for v in vals) for vals in oin]) unseen_mask = np.array([all(v not in seen_vals for v in vals) for vals in oin]) self.results["rule_vs_cache"] = {"seen_values": sorted(seen_vals), "n_items_seen": int(seen_mask.sum()), "n_items_unseen": int(unseen_mask.sum()), "n_fit_rows": int(keep.sum())} for cname, g in (("h5_act_ridge_seenfit", ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d), m_seen)), ("h5_act_procrustes_seenfit", ActivationGraft(GraftSpec("activation", hook, hook, 1.0, "all"), self.model(d), mp_seen)), ("h5_adapter_seenfit", AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), ad_seen))): res = evaluate_condition(self.model(r), hook, g, self.splits, self.vocab, self.device, ["target_iid"]) self.record(cname, pair, res) om = omega_step_acc(self.splits["target_iid"], res["target_iid"]) cl = res["target_iid"].clusters self.results["rule_vs_cache"][cname] = { "seen": cluster_bootstrap_mean(om[seen_mask], [c for c, k in zip(cl, seen_mask) if k], self.n_boot).to_dict(), "unseen": cluster_bootstrap_mean(om[unseen_mask], [c for c, k in zip(cl, unseen_mask) if k], self.n_boot).to_dict(), } print(f"[{utc_now()}] rule-vs-cache done ({time.time()-t0:.0f}s)") # ---------- persistent graft: save adapter for an independent process ---------- pdir = self.out_dir / "persistent" pdir.mkdir(parents=True, exist_ok=True) ad = self.adapters[f"{matched[0][0]}->{matched[0][1]}"] np.savez(pdir / "adapter_graft.npz", W=ad.W, b=ad.b, hook=np.array(hook), recipient=np.array(matched[0][1])) self.results["persistent"] = {"adapter_file": str((pdir / "adapter_graft.npz").relative_to(ROOT)), "evaluated_in_fresh_process": False} # ---------- derived contrasts ---------- self.results["contrasts"] = {} for d, r in pairs: pair = f"{d}->{r}" c = {} for cond in [k for k in self.results["conditions"] if pair in self.results["conditions"][k] and k != "base"]: gain = self.paired(cond, "base", pair) if f"base|{r}|target_iid" in self.items else None # paired against the recipient baseline: items identical; recipient base stored under key r a_ = self.items.get(f"{cond}|{pair}|target_iid") b_ = self.items.get(f"base|{r}|target_iid") s_ = self.items.get(f"ctl_sham_map|{pair}|target_iid") if a_ is None or b_ is None: continue gain = cluster_bootstrap_paired_diff(a_["correct"].astype(float), b_["correct"].astype(float), a_["clusters"], self.n_boot, seed=2).to_dict() entry = {"transfer_gain": gain} if s_ is not None: adj = a_["correct"].astype(float) - b_["correct"].astype(float) - (s_["correct"].astype(float) - b_["correct"].astype(float)) entry["control_adjusted"] = cluster_bootstrap_mean(adj, a_["clusters"], self.n_boot, seed=3).to_dict() cb = self.items.get(f"base|{r}|control_iid") cc = self.items.get(f"{cond}|{pair}|control_iid") if cb is not None and cc is not None: entry["retention_control"] = cluster_bootstrap_paired_diff(cc["correct"].astype(float), cb["correct"].astype(float), cc["clusters"], self.n_boot, seed=4).to_dict() c[cond] = entry self.results["contrasts"][pair] = c # seed summaries for the main conditions self.results["across_pairs"] = {} for cond, d_ in self.results["conditions"].items(): vals = [v["target_iid"]["accuracy"] for p, v in d_.items() if "target_iid" in v and "->" in p] if len(vals) >= 2: self.results["across_pairs"][cond] = {"target": seeds_summary(vals), "control": seeds_summary([v["control_iid"]["accuracy"] for p, v in d_.items() if "control_iid" in v and "->" in p])} self.results["elapsed_seconds"] = time.time() - t0 def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--quick", action="store_true") ap.add_argument("--device", default="auto") ap.add_argument("--seed-base", type=int, default=30_000) ap.add_argument("--tag", default="main") args = ap.parse_args() exp_id = "graft_tier0" + ("_quick" if args.quick else "") + ("" if args.tag == "main" else f"_{args.tag}") cfg = {"quick": args.quick, "seed_base": args.seed_base, "tag": args.tag, "models": "matrix_v1", "protocol": "v2-unrelated-procrustes-reference"} chash = sha256_text(canonical_json(cfg))[:12] run_id = make_run_id(exp_id, chash) out_dir = ROOT / "experiments" / "summaries" / run_id art_dir = ROOT / "artifacts" / "runs" / run_id out_dir.mkdir(parents=True, exist_ok=True) art_dir.mkdir(parents=True, exist_ok=True) man = Manifest(run_id, exp_id, cfg, ROOT / "experiments" / "manifests") registry.record(run_id, exp_id, "planned", config_hash=chash) exp = Exp(args) exp.out_dir = out_dir man.set_status("running") registry.record(run_id, exp_id, "running") try: exp.run() exp.results["run_id"] = run_id exp.results["device"] = str(exp.device) (out_dir / "results.json").write_text(json.dumps(exp.results, indent=1, default=float)) with gzip.open(art_dir / "items.json.gz", "wt") as f: json.dump({k: {"correct": v["correct"].tolist(), "trace": v["trace"].tolist(), "omega": [None if np.isnan(x) else float(x) for x in v["omega"]], "pred": v["pred"]} for k, v in exp.items.items()}, f) man.data["data_hashes"] = exp.results["splits"] man.data["outputs"] = {"results": str(out_dir.relative_to(ROOT) / "results.json"), "items": str(art_dir.relative_to(ROOT))} man.set_status("completed") registry.record(run_id, exp_id, "completed", results=str(out_dir.relative_to(ROOT) / "results.json")) print(f"[{utc_now()}] completed {run_id} in {exp.results['elapsed_seconds']:.0f}s") except Exception as exc: # noqa: BLE001 man.set_status("failed", repr(exc)) registry.record(run_id, exp_id, "failed", error=repr(exc)) raise if __name__ == "__main__": main()