def make_attn_write_to_latent_summed_over_heads_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """Returns the attention value weighted by the post-softmax attention between each pair of tokens,
     summed over the heads, and projected to the gradient of autoencoder latent wrt input.

    Output has shape (n_tokens, n_attended_to_tokens).
    """

    attn_write_to_latent_scalar_deriver = make_attn_write_to_latent_scalar_deriver(
        dst_config
    )  # derive scalar of shape (n_tokens, n_attended_to_tokens, n_heads)

    return attn_write_to_latent_scalar_deriver.apply_transform_fn_to_output(
        lambda activations: activations.sum(dim=-1)[..., None],
        pass_type_to_transform=PassType.FORWARD,
        output_dst=DerivedScalarType.ATTN_WRITE_TO_LATENT_SUMMED_OVER_HEADS,
    )  # derive scalar of shape (n_tokens, n_attended_to_tokens, 1)
Returns the attention value weighted by the post-softmax attention between each pair of tokens,
     summed over the heads, and projected to the gradient of autoencoder latent wrt input.

    Output has shape (n_tokens, n_attended_to_tokens).
    
neuron_explainer/activations/derived_scalars/attention.py