def test_autoencoder_hooks_ablation(add_error):
    """To test ablation, create an autoencoder with large weights at one latent, and ablate it."""
    # create transformer and autoencoder
    layer_idx = 4
    latent_idx = 1
    n_latents = 10
    model = get_test_model()
    cfg = model.cfg
    autoencoder = Autoencoder(n_latents=n_latents, n_inputs=cfg.d_ff).to(model.device)
    autoencoder.latent_bias.data[latent_idx] = +1000  # make sure the latent activates strongly
    input_tokens = torch.arange(5, device=model.device)[None, :]

    # create an ablation hook
    def ablate_latent(latents, layer, **kwargs):
        latents[:, :, latent_idx] = 0
        return latents

    # create a cache hook, to check downstream of the autoencoder hook
    def store_mlp_post_act(x, layer, **kwargs):
        cache[layer] = x.detach().clone().max().item()
        return x

    for ablate in [False, True]:
        cache = {}
        ae_hooks = AutoencoderHooks(
            encode=unflatten(autoencoder.encode),
            decode=unflatten(autoencoder.decode),
            add_error=add_error,
        )
        if ablate:
            ae_hooks.latents.append_fwd(ablate_latent)
        hooks = TransformerHooks()
        hooks.mlp.post_act.append_all(
            AtLayers(layer_idx).append(ae_hooks).append(store_mlp_post_act)
        )

        output, _ = model.forward(input_tokens, hooks=hooks)

        # if not add_error, the MLP post_act is large with no ablation
        # (because the latent activates strongly), and small with ablation
        # if add_error, the MLP post_act is small with no ablation
        # (because the autoencoder returns the identity), and large with ablation
        if ablate != add_error:
            assert cache[layer_idx] < 100
        else:
            assert cache[layer_idx] > 100
To test ablation, create an autoencoder with large weights at one latent, and ablate it.
neuron_explainer/tests/test_hooks.py