def make_in_edge_attribution_scalar_deriver_factory(
    node_type_for_attribution: NodeType,
    q_k_or_v_for_attribution: ActivationLocationType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    """Returns a function that creates a scalar deriver for the edge attribution from arbitrary node
    to the specified downstream node type / sub node type (MLP, autoencoder, or attention head Q, K, or V).
    """

    sub_node_type_to_output_dst = {
        (NodeType.MLP_NEURON, None): DerivedScalarType.MLP_IN_EDGE_ATTRIBUTION,
        (
            NodeType.AUTOENCODER_LATENT,
            None,
        ): DerivedScalarType.ONLINE_AUTOENCODER_IN_EDGE_ATTRIBUTION,
        (
            NodeType.ATTENTION_HEAD,
            ActivationLocationType.ATTN_QUERY,
        ): DerivedScalarType.ATTN_QUERY_IN_EDGE_ATTRIBUTION,
        (
            NodeType.ATTENTION_HEAD,
            ActivationLocationType.ATTN_KEY,
        ): DerivedScalarType.ATTN_KEY_IN_EDGE_ATTRIBUTION,
        (
            NodeType.ATTENTION_HEAD,
            ActivationLocationType.ATTN_VALUE,
        ): DerivedScalarType.ATTN_VALUE_IN_EDGE_ATTRIBUTION,
    }

    output_dst = sub_node_type_to_output_dst[(node_type_for_attribution, q_k_or_v_for_attribution)]

    def make_in_edge_attribution_scalar_deriver(dst_config: DstConfig) -> ScalarDeriver:
        node_write_scalar_source = make_node_write_scalar_source(dst_config)
        return convert_node_write_scalar_deriver_to_in_edge_attribution(
            node_write_scalar_source=node_write_scalar_source,
            output_dst=output_dst,
            dst_config=dst_config,
            downstream_node_type=node_type_for_attribution,
            downstream_q_k_or_v=q_k_or_v_for_attribution,
        )

    return make_in_edge_attribution_scalar_deriver
Returns a function that creates a scalar deriver for the edge attribution from arbitrary node
    to the specified downstream node type / sub node type (MLP, autoencoder, or attention head Q, K, or V).
    
neuron_explainer/activations/derived_scalars/edge_attribution.py