def test_autoencoder_hooks_add_error():
    """Test that the autoencoder hook modifies (or not) the graph correctly"""
    # create transformer and autoencoder
    layer_idx = 4
    n_latents = 10
    model = get_test_model()
    cfg = model.cfg
    autoencoder = Autoencoder(n_latents=n_latents, n_inputs=cfg.d_ff).to(model.device)
    input_tokens = torch.arange(5, device=model.device)[None, :]

    for add_error in [True, False]:
        # create hooks
        hooks = TransformerHooks()
        ae_hooks = AutoencoderHooks(
            encode=unflatten(autoencoder.encode),
            decode=unflatten(autoencoder.decode),
            add_error=add_error,
        )
        hooks.mlp.post_act.append_fwd(AtLayers(layer_idx).append(ae_hooks))

        # make sure to reset autoencoder gradients
        autoencoder.zero_grad()
        assert autoencoder.encoder.weight.grad is None

        # run forward and backward pass, with hooks
        model.zero_grad()
        output, _ = model.forward(input_tokens, hooks=hooks)
        output.sum().backward()
        grads = model.xf_layers[0].mlp.in_layer.weight.grad.clone()

        # run forward and backward pass, without hook
        model.zero_grad()
        output_nohook, _ = model.forward(input_tokens)
        output_nohook.sum().backward()
        grads_nohook = model.xf_layers[0].mlp.in_layer.weight.grad.clone()

        # check forward pass
        # if add_error:
        #     torch.testing.assert_allclose(output, output_nohook, atol=1e-3, rtol=1e-5)
        #     torch.testing.assert_allclose(grads, grads_nohook, atol=1., rtol=1e-1)
        assert add_error == torch.allclose(output, output_nohook, atol=1e-3, rtol=1e-5)
        # check backward pass
        assert add_error == torch.allclose(grads, grads_nohook, atol=1.0, rtol=1e-1)

        # check that we can get gradient on autoencoder latents in all cases
        assert autoencoder.encoder.weight.grad is not None
Test that the autoencoder hook modifies (or not) the graph correctly
neuron_explainer/tests/test_hooks.py