def test_forward_backward_hooks():
    """Test hooks with a simple cache saving function"""

    # create transformer
    model = get_test_model()
    cfg = model.cfg

    # create a save function to hook to the autoencoder
    forward_cache, backward_cache = {}, {}

    def store_forward(xx, layer, **kwargs):
        forward_cache[layer] = xx.detach().clone()
        return xx

    def store_backward(xx, layer, **kwargs):
        backward_cache[layer] = xx.detach().clone()
        return xx

    # create hooks
    hooks = TransformerHooks()
    hooks.mlp.post_act.append_fwd(store_forward)
    hooks.mlp.post_act.append_bwd(store_backward)

    # run the model
    input_tokens = torch.arange(5, device=model.device)[None, :]
    model.zero_grad()
    X, _ = model.forward(input_tokens, hooks=hooks)
    X.sum().backward()
    n_neurons = model.xf_layers[0].mlp.out_layer.weight.shape[1]

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

    # check that the gradient were stored for all layers
    assert list(backward_cache.keys()) == list(range(cfg.n_layers))[::-1]
    assert backward_cache[0].shape == (*input_tokens.shape, n_neurons)
Test hooks with a simple cache saving function
neuron_explainer/tests/test_hooks.py