def test_autoencoder_hooks():
    """Test autoencoder hook with a simple cache saving function"""
    # create transformer and autoencoder
    n_latents = 10
    model = get_test_model()
    cfg = model.cfg
    autoencoder = Autoencoder(n_latents=n_latents, n_inputs=cfg.d_ff).to(model.device)

    # create a save function to hook to the autoencoder
    cache = {}

    def store_latents(latents, layer, **kwargs):
        cache[layer] = latents.detach().clone()
        return latents

    # create hooks
    ae_hooks = AutoencoderHooks(
        encode=unflatten(autoencoder.encode), decode=unflatten(autoencoder.decode)
    )
    ae_hooks.latents.append_fwd(store_latents)
    hooks = TransformerHooks()
    hooks.mlp.post_act.append_fwd(ae_hooks)

    # run the model
    input_tokens = torch.arange(5, device=model.device)[None, :]
    model.forward(input_tokens, hooks=hooks)

    # check that the latents were stored for all layers
    assert list(cache.keys()) == list(range(cfg.n_layers))
    assert cache[0].shape == (*input_tokens.shape, n_latents)

    # do the same, but only on a single layer
    layer_idx = 4
    cache = {}
    ae_hooks = AutoencoderHooks(
        encode=unflatten(autoencoder.encode), decode=unflatten(autoencoder.decode)
    )
    ae_hooks.latents.append_fwd(store_latents)
    hooks = TransformerHooks()
    hooks.mlp.post_act.append_fwd(AtLayers(layer_idx).append(ae_hooks))
    model.forward(input_tokens, hooks=hooks)
    assert list(cache.keys()) == [layer_idx]
    assert cache[layer_idx].shape == (*input_tokens.shape, n_latents)

    # do the same, but on a backward pass
    ae_hooks = AutoencoderHooks(
        encode=unflatten(autoencoder.encode), decode=unflatten(autoencoder.decode)
    )
    ae_hooks.latents.append_fwd(store_latents)
    hooks = TransformerHooks()
    hooks.mlp.post_act.append_bwd(AtLayers(layer_idx).append(ae_hooks))
    model.zero_grad()
    X, _ = model.forward(input_tokens, hooks=hooks)
    X.sum().backward()
    assert list(cache.keys()) == [layer_idx]
    assert cache[layer_idx].shape == (*input_tokens.shape, n_latents)
Test autoencoder hook with a simple cache saving function
neuron_explainer/tests/test_hooks.py