Repository Public Research preview
1import numpy as np2import torch34from wp_ai_0001.alignment.maps import fit_cca, fit_identity, fit_procrustes, fit_random_orthogonal, fit_ridge5from wp_ai_0001.config import OpSeqTaskConfig, ToyModelConfig6from wp_ai_0001.datasets import opseq7from wp_ai_0001.datasets.common import encode_split8from wp_ai_0001.models.decoding import greedy_decode9from wp_ai_0001.models.toy_transformer import ToyTransformer10from wp_ai_0001.transplantation.grafts import ActivationGraft, GraftSpec, NormMatchedRandomGraft, weight_graft111213def test_maps_recover_known_affine_transform():14 rng = np.random.default_rng(0)15 X = rng.standard_normal((2000, 16))16 Q, _ = np.linalg.qr(rng.standard_normal((16, 16)))17 Y = 2.0 * X @ Q + 0.518 for fit in (fit_procrustes, fit_ridge, fit_cca):19 m = fit(X, Y)20 assert m.r2(X, Y) > 0.999, fit.__name__21 assert fit_identity(X, X).r2(X, X) == 1.022 assert fit_random_orthogonal(X, Y, seed=1).r2(X, Y) < 0.5232425def test_maps_handle_width_mismatch():26 rng = np.random.default_rng(1)27 X = rng.standard_normal((500, 12))28 Y = rng.standard_normal((500, 16))29 for fit in (fit_procrustes, fit_ridge, fit_cca, fit_random_orthogonal):30 m = fit(X, Y)31 assert m(X).shape == (500, 16)323334def _models():35 torch.manual_seed(0)36 cfg = OpSeqTaskConfig()37 split = opseq.generate(cfg, 32, seed=21, name="g")38 vocab = opseq.vocab_for(cfg)39 enc = encode_split(split, vocab, 24)40 mc = ToyModelConfig(d_model=32, n_layers=2, n_heads=2, d_ff=64, max_seq_len=24)41 donor = ToyTransformer(mc, len(vocab)).eval()42 torch.manual_seed(1)43 recip = ToyTransformer(mc, len(vocab)).eval()44 return donor, recip, enc454647def test_identity_self_graft_is_noop_and_cross_graft_changes_output():48 donor, recip, enc = _models()49 ids = torch.from_numpy(enc["ids"]); pl = torch.from_numpy(enc["prompt_len"]); na = torch.from_numpy(enc["n_answer"])50 K = enc["answers"].shape[1]51 ref = greedy_decode(recip, ids, pl, na, K)52 X = np.zeros((4, 32)); amap = fit_identity(X, X)53 spec = GraftSpec("activation", "resid_post.0", "resid_post.0", alpha=1.0, targeting="all")54 g = ActivationGraft(spec, recip, amap) # donor == recipient, identity map55 recip.hooks.add("resid_post.0", g)56 out = greedy_decode(recip, ids, pl, na, K)57 recip.hooks.clear()58 assert torch.equal(ref, out) and g.applied > 059 g2 = ActivationGraft(spec, donor, amap)60 recip.hooks.add("resid_post.0", g2)61 out2 = greedy_decode(recip, ids, pl, na, K)62 recip.hooks.clear()63 assert not torch.equal(ref, out2)64 # alpha = 0 is a no-op even with a foreign donor65 g3 = ActivationGraft(GraftSpec("activation", "resid_post.0", "resid_post.0", alpha=0.0), donor, amap)66 recip.hooks.add("resid_post.0", g3)67 out3 = greedy_decode(recip, ids, pl, na, K)68 recip.hooks.clear()69 assert torch.equal(ref, out3)707172def test_norm_matched_random_has_matched_norms():73 donor, recip, enc = _models()74 ids = torch.from_numpy(enc["ids"]); pos = torch.from_numpy(enc["prompt_len"] - 1)75 X = np.zeros((4, 32)); amap = fit_identity(X, X)76 spec = GraftSpec("random", "resid_post.0", "resid_post.0")77 inner = ActivationGraft(spec, donor, amap)78 rnd = NormMatchedRandomGraft(spec, donor, amap, seed=0)79 rows = torch.zeros(ids.shape[0], 32)80 a = inner.replacement(rows, ids, pos)81 b = rnd.replacement(rows, ids, pos)82 assert torch.allclose(a.norm(dim=-1), b.norm(dim=-1), atol=1e-4) and not torch.allclose(a, b)838485def test_weight_graft_copies_and_checks_shapes():86 donor, recip, _ = _models()87 rep = weight_graft(recip, donor, ["attn", "mlp"], 0, 0)88 assert rep["n_params_copied"] > 089 assert torch.equal(recip.blocks[0].attn.qkv.weight, donor.blocks[0].attn.qkv.weight)90 assert not torch.equal(recip.blocks[1].attn.qkv.weight, donor.blocks[1].attn.qkv.weight)91