def make_attn_write_sum_heads_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """Returns a ScalarDeriver object for the attention write, between any two tokens,
    summed over all heads."""
    scalar_deriver = make_attn_write_scalar_deriver(dst_config)

    def _sum_heads(attn_write: torch.Tensor) -> torch.Tensor:
        n_tokens_written_to, n_tokens_attended_to, n_heads, d_model = attn_write.shape
        return attn_write.sum(dim=2)

    return scalar_deriver.apply_transform_fn_to_output(
        _sum_heads,
        pass_type_to_transform=PassType.FORWARD,
        output_dst=DerivedScalarType.ATTN_WRITE_SUM_HEADS,
    )
Returns a ScalarDeriver object for the attention write, between any two tokens,
    summed over all heads.
neuron_explainer/activations/derived_scalars/attention.py