def make_token_attribution_scalar_deriver(dst_config: DstConfig) -> ScalarDeriver:
    """This computes an attribution value for each token in the sequence, the inner product of act and grad for the embedding.
    This corresponds to the estimated importance of this token to the final prediction."""
    activation_location_type = ActivationLocationType.RESID_POST_EMBEDDING
    dst = DerivedScalarType.TOKEN_ATTRIBUTION

    act_times_grad_scalar_deriver = make_scalar_deriver_factory_for_act_times_grad(
        activation_location_type=activation_location_type,
        dst=dst,
    )(dst_config)

    def sum_over_last_dim(tensor: torch.Tensor) -> torch.Tensor:
        return tensor.sum(dim=-1)[:, None]  # singleton final dim

    return act_times_grad_scalar_deriver.apply_transform_fn_to_output(
        sum_over_last_dim, pass_type_to_transform=PassType.FORWARD, output_dst=dst
    )
This computes an attribution value for each token in the sequence, the inner product of act and grad for the embedding.
    This corresponds to the estimated importance of this token to the final prediction.
neuron_explainer/activations/derived_scalars/residual.py