def make_in_edge_activation_scalar_deriver_factory(
    activation_location_type: ActivationLocationType,
    q_or_k: 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 activation location type / sub activation location type (MLP post act,
    autoencoder latent, attention head Q or K).
    """

    sub_node_type_to_output_dst = {
        (ActivationLocationType.MLP_POST_ACT, None): DerivedScalarType.MLP_IN_EDGE_ACTIVATION,
        (
            ActivationLocationType.ONLINE_AUTOENCODER_LATENT,
            None,
        ): DerivedScalarType.ONLINE_AUTOENCODER_IN_EDGE_ACTIVATION,
        (
            ActivationLocationType.ATTN_QK_PROBS,
            ActivationLocationType.ATTN_QUERY,
        ): DerivedScalarType.ATTN_QUERY_IN_EDGE_ACTIVATION,
        (
            ActivationLocationType.ATTN_QK_PROBS,
            ActivationLocationType.ATTN_KEY,
        ): DerivedScalarType.ATTN_KEY_IN_EDGE_ACTIVATION,
    }

    output_dst = sub_node_type_to_output_dst[(activation_location_type, q_or_k)]

    def make_in_edge_activation_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_activation(
            node_write_scalar_source=node_write_scalar_source,
            output_dst=output_dst,
            dst_config=dst_config,
            downstream_activation_location_type=activation_location_type,
            downstream_q_or_k=q_or_k,
        )

    return make_in_edge_activation_scalar_deriver
Returns a function that creates a scalar deriver for the edge attribution from arbitrary node
    to the specified downstream activation location type / sub activation location type (MLP post act,
    autoencoder latent, attention head Q or K).
    
neuron_explainer/activations/derived_scalars/edge_activation.py