spb/wpai0001_uqo

1"""Gate C sanity: hooks are a true no-op when disabled/empty, and interventions have effects."""2import torch34from wp_ai_0001.config import OpSeqTaskConfig, ToyModelConfig5from wp_ai_0001.datasets import opseq6from wp_ai_0001.datasets.common import encode_split7from wp_ai_0001.models.hooks import Add, Capture, GaussianNoise, ProjectSubspace, Replace, Zero8from wp_ai_0001.models.toy_transformer import ToyTransformer91011def _setup(seed=0):12    torch.manual_seed(seed)13    cfg = OpSeqTaskConfig()14    split = opseq.generate(cfg, 64, seed=11, name="sanity")15    vocab = opseq.vocab_for(cfg)16    enc = encode_split(split, vocab, 24)17    model = ToyTransformer(ToyModelConfig(d_model=32, n_layers=2, n_heads=2, d_ff=64, max_seq_len=24), len(vocab)).eval()18    return model, torch.from_numpy(enc["ids"]), torch.from_numpy(enc["prompt_len"] - 1), torch.from_numpy(enc["answers"][:, 0])192021def test_no_hooks_is_bitwise_identity():22    model, ids, pos, _ = _setup()23    with torch.no_grad():24        ref = model(ids, pos)25        assert model.hooks.calls == 026        out = model(ids, pos)27    assert torch.equal(ref, out)282930def test_capture_and_disabled_hooks_are_bitwise_identity():31    model, ids, pos, _ = _setup()32    with torch.no_grad():33        ref = model(ids, pos)34        cap = Capture()35        for n in model.hook_names():36            model.hooks.add(n, cap)37        out_cap = model(ids, pos)38        assert set(cap.store) == set(model.hook_names())39        # a modifying hook that is disabled must also be a no-op40        model.hooks.add("resid_post.0", Zero())41        with model.hooks.disabled():42            out_dis = model(ids, pos)43        model.hooks.enabled = False44        out_dis2 = model(ids, pos)45    assert torch.equal(ref, out_cap) and torch.equal(ref, out_dis) and torch.equal(ref, out_dis2)464748def test_identity_object_returned_when_no_hooks():49    model, ids, pos, _ = _setup()50    x = torch.randn(2, 3, 32)51    assert model.hooks("resid_post.0", x) is x52    model.hooks.enabled = False53    model.hooks.add("resid_post.0", Zero())54    assert model.hooks("resid_post.0", x) is x555657def test_interventions_change_outputs_reproducibly():58    model, ids, pos, _ = _setup()59    with torch.no_grad():60        ref = model(ids, pos)61        model.hooks.add("resid_post.0", Zero(positions=pos))62        a = model(ids, pos)63        b = model(ids, pos)64    assert not torch.equal(ref, a) and torch.equal(a, b)65    # unpatched positions are untouched up to the hooked layer: check resid_post.0 itself66    model.hooks.clear()67    cap = Capture()68    model.hooks.add("resid_post.0", cap)69    with torch.no_grad():70        model(ids, pos)71    clean = cap.store["resid_post.0"]72    model.hooks.clear()73    cap2 = Capture()74    model.hooks.add("resid_post.0", Zero(positions=pos))75    model.hooks.add("resid_post.0", cap2)76    with torch.no_grad():77        model(ids, pos)78    z = cap2.store["resid_post.0"]79    B = ids.shape[0]80    assert torch.all(z[torch.arange(B), pos] == 0)81    mask = torch.ones_like(clean, dtype=torch.bool)82    mask[torch.arange(B), pos] = False83    assert torch.equal(z[mask], clean[mask])848586def test_replace_with_own_activation_is_identity():87    model, ids, pos, _ = _setup()88    cap = Capture()89    model.hooks.add("resid_mid.1", cap)90    with torch.no_grad():91        ref = model(ids, pos)92    model.hooks.clear()93    model.hooks.add("resid_mid.1", Replace(cap.store["resid_mid.1"]))94    with torch.no_grad():95        out = model(ids, pos)96    assert torch.allclose(ref, out, atol=0, rtol=0)979899def test_add_noise_project_have_expected_algebra():100    model, ids, pos, _ = _setup()101    x = torch.randn(4, 5, 32)102    v = torch.randn(32)103    y = Add(v, scale=2.0)(x, model.hooks.__class__ and __import__("wp_ai_0001.models.hooks", fromlist=["HookContext"]).HookContext("h"))104    assert torch.allclose(y - x, 2.0 * v.expand_as(x))105    ctx = __import__("wp_ai_0001.models.hooks", fromlist=["HookContext"]).HookContext("h")106    n1, n2 = GaussianNoise(0.5, seed=3)(x, ctx), GaussianNoise(0.5, seed=3)(x, ctx)107    assert torch.equal(n1, n2) and not torch.equal(n1, x)108    basis = torch.randn(32, 4)109    abl = ProjectSubspace(basis, "ablate")(x, ctx)110    keep = ProjectSubspace(basis, "keep")(x, ctx)111    assert torch.allclose(abl + keep, x, atol=1e-5)112    assert torch.allclose(abl @ basis, torch.zeros(4, 5, 4), atol=1e-4)113114115def test_head_patch_touches_only_selected_head():116    model, ids, pos, _ = _setup()117    cap = Capture()118    model.hooks.add("attn_heads.0", cap)119    with torch.no_grad():120        model(ids, pos)121    src = torch.zeros_like(cap.store["attn_heads.0"])122    model.hooks.clear()123    cap2 = Capture()124    model.hooks.add("attn_heads.0", Replace(src, positions=[2], heads=[1]))125    model.hooks.add("attn_heads.0", cap2)126    with torch.no_grad():127        model(ids, pos)128    out = cap2.store["attn_heads.0"]129    ref = cap.store["attn_heads.0"]130    assert torch.all(out[:, 2, 1] == 0)131    assert torch.equal(out[:, 2, 0], ref[:, 2, 0]) and torch.equal(out[:, :2], ref[:, :2]) and torch.equal(out[:, 3:], ref[:, 3:])132