def make_attn_write_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """Returns a ScalarDeriver object for the attention write, between any two tokens,
    and for each head."""
    model_context = dst_config.get_model_context()
    layer_indices = dst_config.layer_indices or list(range(model_context.n_layers))

    assert not dst_config.derive_gradients
    sub_scalar_sources = (
        RawScalarSource(
            activation_location_type=ActivationLocationType.ATTN_QK_PROBS,
            pass_type=PassType.FORWARD,
        ),
        RawScalarSource(
            activation_location_type=ActivationLocationType.ATTN_VALUE, pass_type=PassType.FORWARD
        ),
    )

    return ScalarDeriver(
        dst=DerivedScalarType.ATTN_WRITE,
        dst_config=dst_config,
        sub_scalar_sources=sub_scalar_sources,
        tensor_calculate_derived_scalar_fn=make_attn_write_tensor_calculate_derived_scalar_fn(
            model_context=model_context,
            layer_indices=layer_indices,
        ),
    )
Returns a ScalarDeriver object for the attention write, between any two tokens,
    and for each head.
neuron_explainer/activations/derived_scalars/attention.py