def convert_node_write_scalar_deriver_to_in_edge_activation(
    node_write_scalar_source: ScalarSource,
    output_dst: DerivedScalarType,
    dst_config: DstConfig,
    downstream_activation_location_type: ActivationLocationType,
    downstream_q_or_k: ActivationLocationType | None,
) -> ScalarDeriver:
    """Converts a scalar deriver for a write vector from some upstream node type to a scalar deriver for
    in edge activation for downstream nodes of some type (MLP, autoencoder, or attention head). In the
    case of attention heads, this is split up by subnode (Q or K)."""

    model_context = dst_config.get_model_context()
    autoencoder_context = dst_config.get_autoencoder_context()
    assert isinstance(model_context, StandardModelContext)
    transformer = model_context.get_or_create_model()
    reconstituter = ActivationReconstituter.from_activation_location_type(
        transformer=transformer,
        autoencoder_context=autoencoder_context,
        activation_location_type=downstream_activation_location_type,
        q_or_k=downstream_q_or_k,
    )
    return reconstituter.make_jvp_scalar_deriver(
        write_scalar_source=node_write_scalar_source,
        dst_config=dst_config,
        output_dst=output_dst,
    )
Converts a scalar deriver for a write vector from some upstream node type to a scalar deriver for
    in edge activation for downstream nodes of some type (MLP, autoencoder, or attention head). In the
    case of attention heads, this is split up by subnode (Q or K).
neuron_explainer/activations/derived_scalars/edge_activation.py