def make_unflattened_attn_write_norm_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """Returns a ScalarDeriver object for the attention write norm, between any
    two tokens. The actual attention write vector is the vector sum of many attention
    write vectors. Note that the norm of this vector sum will be less that the sum of the
    attention write norms computed here."""
    model_context = dst_config.get_model_context()
    layer_indices = dst_config.layer_indices or list(range(model_context.n_layers))

    if dst_config.derive_gradients:
        sub_scalar_sources: tuple[ScalarSource, ...] = (
            RawScalarSource(
                activation_location_type=ActivationLocationType.ATTN_QK_PROBS,
                pass_type=PassType.FORWARD,
            ),
            RawScalarSource(
                activation_location_type=ActivationLocationType.ATTN_QK_PROBS,
                pass_type=PassType.BACKWARD,
            ),
            RawScalarSource(
                activation_location_type=ActivationLocationType.ATTN_VALUE,
                pass_type=PassType.FORWARD,
            ),
        )
    else:
        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.UNFLATTENED_ATTN_WRITE_NORM,
        dst_config=dst_config,
        sub_scalar_sources=sub_scalar_sources,
        tensor_calculate_derived_scalar_fn=_make_unflattened_attn_write_norm_tensor_calculate_derived_scalar_fn(
            model_context=model_context,
            layer_indices=layer_indices,
        ),
    )
Returns a ScalarDeriver object for the attention write norm, between any
    two tokens. The actual attention write vector is the vector sum of many attention
    write vectors. Note that the norm of this vector sum will be less that the sum of the
    attention write norms computed here.
neuron_explainer/activations/derived_scalars/attention.py