def add_saving_hooks(
    transformer_graph: TransformerHookGraph,
    activation_location_type_and_pass_types: list[ActivationLocationTypeAndPassType],
    device: torch.device,
    unpadded_prompt_length: int,
    batch_index: int,
) -> dict[ActivationLocationTypeAndPassType, dict[LayerIndex, torch.Tensor]]:
    """This is a helper function that returns a TransformerHooks object containing hooks at naturally existing hook locations
    (e.g. MLP post-activations, rather than autoencoder latent activations) for the given
    scalar derivers, and a dictionary mapping each activation location type with a naturally existing hook to the activations requested
    (before it is filled with activations during forward passes)."""

    requested_activations_by_location_type_and_pass_type = {}

    for activation_location_type_and_pass_type in activation_location_type_and_pass_types:
        (
            hook_fn,
            requested_activations_by_location_type_and_pass_type[
                activation_location_type_and_pass_type
            ],
        ) = create_saving_hook_fn(
            device,
            activation_location_type_and_pass_type=activation_location_type_and_pass_type,
            unpadded_prompt_length=unpadded_prompt_length,
            batch_index=batch_index,
        )

        transformer_graph.append(
            activation_location_type_and_pass_type,
            hook_fn,
        )

    return requested_activations_by_location_type_and_pass_type
This is a helper function that returns a TransformerHooks object containing hooks at naturally existing hook locations
    (e.g. MLP post-activations, rather than autoencoder latent activations) for the given
    scalar derivers, and a dictionary mapping each activation location type with a naturally existing hook to the activations requested
    (before it is filled with activations during forward passes).
neuron_explainer/activation_server/derived_scalar_computation.py