def make_apply_attn_V_act(
    transformer: Transformer,
    q_k_or_v: ActivationLocationType | None,
    detach_layer_norm_scale: bool,
) -> Callable[[torch.Tensor, LayerIndex, PassType], tuple[torch.Tensor, torch.Tensor]]:
    """Used in functions that require reconstituting some or all of the attention head
    operation. Supports specifying a stop grad through all but one of Q, K, and V; or
    if q_k_or_v is None, then all of Q, K, and V are backprop'ed through."""

    def apply_attn_V_act(
        resid: torch.Tensor,
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        assert pass_type == PassType.FORWARD
        transformer_layer = transformer.xf_layers[layer_index]

        attn_input = apply_layer_norm(
            resid.unsqueeze(0),
            transformer_layer.ln_1,
            detach_layer_norm_scale=detach_layer_norm_scale,
        )  # add batch dimension

        hooks = AttentionHooks()
        add_q_k_or_v_detach_hook(hooks, q_k_or_v)
        get_hook, get_v = make_hook_getter()
        hooks.v.append_fwd(get_hook)
        get_hook, get_attn = make_hook_getter()
        hooks.qk_probs.append_fwd(get_hook)
        # avoid v_out expense
        hooks.v.append_fwd(zero_batch_dim_hook)
        hooks.qk_probs.append_fwd(zero_batch_dim_hook)
        transformer_layer.attn.forward(attn_input, hooks=hooks)
        # remove batch dimensions
        return get_attn()[0], get_v()[0]

    return apply_attn_V_act
Used in functions that require reconstituting some or all of the attention head
    operation. Supports specifying a stop grad through all but one of Q, K, and V; or
    if q_k_or_v is None, then all of Q, K, and V are backprop'ed through.
neuron_explainer/activations/derived_scalars/reconstituted.py