def make_reconstituted_attribution_fn(
    transformer: Transformer,
    autoencoder_context: AutoencoderContext | None,
    node_type: NodeType,  # the type of the node of interest
    q_k_or_v: (
        ActivationLocationType | None
    ),  # if node_type is ATTENTION_HEAD, this specifies Q, K, or V
    detach_layer_norm_scale: bool,
) -> Callable[[torch.Tensor, torch.Tensor, LayerIndex, PassType], torch.Tensor]:
    """The 'attribution' of a node is taken to be the product of the node's gradient and the node's activation.
    This returns a function to compute the attribution of attention heads (specifically mediated by Q, K, or V), MLP activations,
    or autoencoder activations. The input expected by that function is the residual stream just before the node in question.
    The function returned can be used for further analysis, e.g. computing the gradient of the attribution with respect to the
    input residual stream.

    Note that this can be used to compute the attribution of one OR many downstream subnodes, depending on the tensor_indices_for_grad
    (an empty tensor_indices_for_grad corresponds to the entire layer worth of activations)."""

    assert (q_k_or_v is None) == (
        node_type != NodeType.ATTENTION_HEAD
    )  # for these functions, we require q_k_or_v to be
    # specified if node_type is ATTENTION_HEAD
    match node_type:
        case NodeType.ATTENTION_HEAD:
            match q_k_or_v:
                case ActivationLocationType.ATTN_QUERY | ActivationLocationType.ATTN_KEY:
                    # in all cases but attn value, the attribution fn is the hadamard product of the activation and the gradient
                    # NOTE: q_k_or_v = None covers all non-attention activations
                    match q_k_or_v:
                        case ActivationLocationType.ATTN_QUERY:
                            attention_trace_type = AttentionTraceType.Q
                        case ActivationLocationType.ATTN_KEY:
                            attention_trace_type = AttentionTraceType.K
                        case None:
                            attention_trace_type = AttentionTraceType.QK
                        case _:
                            raise NotImplementedError(q_k_or_v)
                    activation_fn = make_reconstituted_activation_fn(
                        transformer=transformer,
                        autoencoder_context=autoencoder_context,
                        node_type=node_type,
                        pre_or_post_act=PreOrPostAct.POST,
                        detach_layer_norm_scale=detach_layer_norm_scale,
                        attention_trace_type=attention_trace_type,
                    )

                    def attribution_fn(
                        resid: torch.Tensor,
                        grad: torch.Tensor,
                        layer_index: LayerIndex,
                        pass_type: PassType,
                    ) -> torch.Tensor:
                        activation = activation_fn(resid, layer_index, pass_type)
                        assert activation.shape == grad.shape, (
                            activation.shape,
                            grad.shape,
                        )
                        return activation * grad

                case ActivationLocationType.ATTN_VALUE:
                    assert (
                        get_activation_location_type_for_node_type(node_type, q_k_or_v)
                        == ActivationLocationType.ATTN_WEIGHTED_SUM_OF_VALUES
                    )
                    apply_attn_V_act = make_apply_attn_V_act(
                        transformer=transformer,
                        q_k_or_v=q_k_or_v,
                        detach_layer_norm_scale=detach_layer_norm_scale,
                    )

                    def attribution_fn(
                        resid: torch.Tensor,
                        grad: torch.Tensor,
                        layer_index: LayerIndex,
                        pass_type: PassType,
                    ) -> torch.Tensor:
                        attn, V = apply_attn_V_act(resid, layer_index, pass_type)
                        attn_weighted_V = torch.einsum("qkh,khd->qkhd", attn, V)
                        # grad is w/r/t (attn_weighted_V summed over k, or ATTN_WEIGHTED_SUM_OF_VALUES)
                        return torch.einsum("qkhd,qhd->qkh", attn_weighted_V, grad)

                case _:
                    raise NotImplementedError(q_k_or_v)
        case _:
            activation_fn = make_reconstituted_activation_fn(
                transformer=transformer,
                autoencoder_context=autoencoder_context,
                node_type=node_type,
                pre_or_post_act=PreOrPostAct.POST,
                detach_layer_norm_scale=detach_layer_norm_scale,
                attention_trace_type=None,
            )

            def attribution_fn(
                resid: torch.Tensor,
                grad: torch.Tensor,
                layer_index: LayerIndex,
                pass_type: PassType,
            ) -> torch.Tensor:
                activation = activation_fn(resid, layer_index, pass_type)
                assert activation.shape == grad.shape, (
                    activation.shape,
                    grad.shape,
                )
                return activation * grad

    return attribution_fn
The 'attribution' of a node is taken to be the product of the node's gradient and the node's activation.
    This returns a function to compute the attribution of attention heads (specifically mediated by Q, K, or V), MLP activations,
    or autoencoder activations. The input expected by that function is the residual stream just before the node in question.
    The function returned can be used for further analysis, e.g. computing the gradient of the attribution with respect to the
    input residual stream.

    Note that this can be used to compute the attribution of one OR many downstream subnodes, depending on the tensor_indices_for_grad
    (an empty tensor_indices_for_grad corresponds to the entire layer worth of activations).
neuron_explainer/activations/derived_scalars/edge_attribution.py