How to modify activations

import torch
from contextlib import contextmanager

@contextmanager
def intervene(module, fn):
    """Attach an editing hook for the duration of the block, then remove it.
    A context manager because a forgotten hook silently corrupts every
    subsequent forward pass."""
    def hook(mod, inputs, output):
        if isinstance(output, tuple):
            return (fn(output[0]),) + output[1:]     # attention returns a tuple
        return fn(output)
    handle = module.register_forward_hook(hook)
    try:
        yield
    finally:
        handle.remove()

# --- the four replacements, as editing functions -----------------------------

def zero(position=None):
    def fn(x):
        x = x.clone()                    # never edit in place
        if position is None:
            return torch.zeros_like(x)
        x[:, position] = 0
        return x
    return fn

def mean_ablate(mean_vector, position=None):
    def fn(x):
        x = x.clone()
        if position is None:
            x[:] = mean_vector
        else:
            x[:, position] = mean_vector
        return x
    return fn

def patch_from(donor, position=None):
    """donor: activation captured from a different input, same shape."""
    def fn(x):
        x = x.clone()
        if position is None:
            return donor.clone()
        x[:, position] = donor[:, position]
        return x
    return fn

def scale(factor, position=None):
    def fn(x):
        x = x.clone()
        if position is None:
            return x * factor
        x[:, position] *= factor
        return x
    return fn

# --- usage ------------------------------------------------------------------

prompt = tok("The Eiffel Tower is in", return_tensors="pt")

with torch.no_grad():
    baseline = model(**prompt).logits

with intervene(model.transformer.h[8].mlp, zero()):
    with torch.no_grad():
        ablated = model(**prompt).logits

# Assert the intervention actually fired. Silently-inert hooks — wrong module
# path, edit not returned, no clone — are the most common failure here.
assert not torch.allclose(baseline, ablated), "intervention had no effect"

import torch, torch.nn.functional as F

@torch.no_grad()
def effect(model, prompt, target_id, module, editor, position=-1):
    """Change in the target token's logit — the standard effect size."""
    base = model(**prompt).logits[0, position]
    with intervene(module, editor):
        edited = model(**prompt).logits[0, position]

    return {
        # Logit difference is preferred over probability difference: it is not
        # squashed near 0 or 1, so effects stay comparable across prompts.
        "logit_delta": float(edited[target_id] - base[target_id]),
        "prob_before": float(F.softmax(base, -1)[target_id]),
        "prob_after": float(F.softmax(edited, -1)[target_id]),
        # How much did EVERYTHING change? A large KL with a small target delta
        # means you damaged the model generally rather than the behaviour.
        "kl": float(F.kl_div(F.log_softmax(edited, -1),
                             F.log_softmax(base, -1),
                             log_target=True, reduction="sum")),
    }

# Sweep every layer to localise, and always run a control prompt where the
# component should NOT matter. An effect present in both is not specific.
for L in range(model.config.n_layer):
    e = effect(model, prompt, paris_id, model.transformer.h[L].mlp, zero())
    c = effect(model, control, control_id, model.transformer.h[L].mlp, zero())
    print(f"layer {L:2}  target {e['logit_delta']:+6.2f}  "
          f"control {c['logit_delta']:+6.2f}  kl {e['kl']:.3f}")