def test_autoencoder_backward_hooks(add_error):
    """Test autoencoder hook, running a backward pass *from* a latent"""
    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, :]

    cache = {}

    def store_latents_no_detach(latents, layer, **kwargs):
        cache[layer] = latents
        return latents

    ae_hooks = AutoencoderHooks(
        encode=unflatten(autoencoder.encode),
        decode=unflatten(autoencoder.decode),
        add_error=add_error,
    )
    ae_hooks.latents.append_fwd(store_latents_no_detach)
    hooks = TransformerHooks()
    hooks.mlp.post_act.append_fwd(AtLayers(layer_idx).append(ae_hooks))
    model.zero_grad()
    X, _ = model.forward(input_tokens, hooks=hooks)

    assert model.xf_layers[0].mlp.in_layer.weight.grad is None
    latents = cache[layer_idx]
    latents[:, 0].sum().backward()  # backward from the first latent
    assert model.xf_layers[0].mlp.in_layer.weight.grad is not None
Test autoencoder hook, running a backward pass *from* a latent
neuron_explainer/tests/test_hooks.py