def get_activation_location_type_for_node_type(
    node_type: NodeType, q_k_or_v: ActivationLocationType | None
) -> ActivationLocationType:
    """This returns the activation location associated with a node of a given type, and
    specifying Q, K, or V in the case of attention. This returns an activation location type
    that is sufficient to determine that node's effect on the residual stream (post-softmax
    in the case of Q, K; ATTN_WEIGHTED_SUM_OF_VALUES in the case of V)"""
    match node_type:
        case NodeType.ATTENTION_HEAD:
            assert q_k_or_v is not None
            match q_k_or_v:
                case ActivationLocationType.ATTN_VALUE:
                    return ActivationLocationType.ATTN_WEIGHTED_SUM_OF_VALUES
                case ActivationLocationType.ATTN_QUERY | ActivationLocationType.ATTN_KEY:
                    return ActivationLocationType.ATTN_QK_PROBS
                case _:
                    raise NotImplementedError(q_k_or_v)
        case NodeType.MLP_NEURON:
            assert q_k_or_v is None
            return ActivationLocationType.MLP_POST_ACT
        case NodeType.AUTOENCODER_LATENT:
            assert q_k_or_v is None
            return ActivationLocationType.ONLINE_AUTOENCODER_LATENT
        case NodeType.MLP_AUTOENCODER_LATENT:
            assert q_k_or_v is None
            return ActivationLocationType.ONLINE_MLP_AUTOENCODER_LATENT
        case NodeType.ATTENTION_AUTOENCODER_LATENT:
            assert q_k_or_v is None
            return ActivationLocationType.ONLINE_ATTENTION_AUTOENCODER_LATENT
        case _:
            raise NotImplementedError(node_type)
This returns the activation location associated with a node of a given type, and
    specifying Q, K, or V in the case of attention. This returns an activation location type
    that is sufficient to determine that node's effect on the residual stream (post-softmax
    in the case of Q, K; ATTN_WEIGHTED_SUM_OF_VALUES in the case of V)
neuron_explainer/activations/derived_scalars/edge_attribution.py