def get_transformer_graph_hooks_and_activation_caches(
    multi_autoencoder_context: MultiAutoencoderContext | None,
    batched_ds_computation_params: list[DerivedScalarComputationParams],
    batched_activation_index_for_backward_pass_by_name: list[dict[str, ActivationIndex | None]],
) -> tuple[
    TransformerHookGraph,
    list[dict[ActivationLocationTypeAndPassType, dict[LayerIndex, torch.Tensor]]],
    list[dict[str, torch.Tensor | None]],
]:
    """This is a helper function that returns:
    1. a TransformerHookGraph object containing hooks for the given
    scalar derivers (this can be passed to Transformer using the as_transformer_hooks method to add
    hooks to the transformer forward and backward pass)
    2. dictionaries mapping each activation location type to the activations requested
    (before it is filled with activations during forward passes), one dictionary per batch element
    3. dictionaries each containing just one value, the scalar tensor on which the backward pass can be run (
    unlike the activations in 2, this tensor is still attached to the pytorch model), one dictionary per batch element
    """

    batched_activation_location_type_and_pass_types = [
        params.activation_location_type_and_pass_types for params in batched_ds_computation_params
    ]
    batched_device = [params.device_for_raw_activations for params in batched_ds_computation_params]
    batched_ablation_specs = [params.ablation_specs for params in batched_ds_computation_params]
    batched_prompt_lengths = [params.prompt_length for params in batched_ds_computation_params]

    # This is a callable that can be used similarly to a Hooks object.
    transformer_graph = TransformerHookGraph()

    """This step constructs the activation location types needed (for the case where they don't already exist). Any hooks to be added
    to that location type can then be appended to transformer_graph in the normal way.
    """

    # steps:
    # 1. add autoencoder graph to transformer_graph (injected autoencoders specified in multi_autoencoder_context are
    #    hooked in a second set of forward hooks, called after ablating and saving hooks, and before activation grabbing hooks)
    #
    # for each batched element:
    #     2. add ablating hooks
    #     3. add activation grabbing hooks (storing without detaching, for backward pass). These come last, after ablating and saving hooks.
    #     4. add saving hooks (storing with detaching)
    #
    # Because the autoencoder is hooked in a second set of hooks, followed by the grabbing hooks, they are always called last:
    # - fwd hooks: ablate activations, save activations, ...
    # - bwd hooks: ablate gradients, save gradients, ...
    # - fwd2 hooks: autoencoder (convert to latent, ablate latents, grab latents, save latents, convert back to activations), grab activations, ...

    # add autoencoder graph
    if multi_autoencoder_context is not None:
        has_multiple_autoencoders = (
            len(multi_autoencoder_context.autoencoder_context_by_node_type) > 1
        )
        for (
            node_type,
            autoencoder_context,
        ) in multi_autoencoder_context.autoencoder_context_by_node_type.items():
            subgraph = AutoencoderHookGraph(
                autoencoder_context, is_one_of_multiple_autoencoders=has_multiple_autoencoders
            )
            subgraph_name = f"{node_type.value}"
            transformer_graph.inject_subgraph(subgraph, subgraph_name)

    (
        batched_requested_activations_by_location_type_and_pass_type,
        batched_requested_attached_activation_for_backward_pass_by_name,
    ) = ([], [])
    assert (
        len(batched_activation_location_type_and_pass_types)
        == len(batched_ablation_specs)
        == len(batched_activation_index_for_backward_pass_by_name)
        == len(batched_prompt_lengths)
    )
    for i in range(len(batched_activation_location_type_and_pass_types)):
        # add ablating hooks
        ablation_spec = batched_ablation_specs[i]
        if ablation_spec is not None:
            add_ablating_hooks(transformer_graph, ablation_spec, batch_index=i)

        # add activation grabbing hooks; the grabbed activations are stored in dicts, keyed by the
        # string name assigned to them
        requested_attached_activation_for_backward_pass_by_name: dict[str, torch.Tensor | None] = {}
        for (
            name,
            activation_index_for_backward_pass,
        ) in batched_activation_index_for_backward_pass_by_name[i].items():
            if activation_index_for_backward_pass is not None:
                requested_attached_activation_for_backward_pass_by_name = add_grabbing_hook_for_backward_pass(
                    requested_attached_activation_for_backward_pass_by_name,
                    name,
                    transformer_graph=transformer_graph,
                    activation_index_for_backward_pass=activation_index_for_backward_pass,
                    batch_index=i,
                    append_to_fwd2=True,  # append to fwd2 when using Reconstituter to obtain gradients, so that
                    # the preceding residual stream backward pass hooks can be called after running .backward()
                    # from the grabbed activation
                )
            else:
                requested_attached_activation_for_backward_pass_by_name[name] = None

        # add saving hooks
        requested_activations_by_location_type_and_pass_type = add_saving_hooks(
            transformer_graph=transformer_graph,
            activation_location_type_and_pass_types=batched_activation_location_type_and_pass_types[
                i
            ],
            device=batched_device[i],
            unpadded_prompt_length=batched_prompt_lengths[i],
            batch_index=i,
        )

        batched_requested_activations_by_location_type_and_pass_type.append(
            requested_activations_by_location_type_and_pass_type
        )
        batched_requested_attached_activation_for_backward_pass_by_name.append(
            requested_attached_activation_for_backward_pass_by_name
        )

    return (
        transformer_graph,
        batched_requested_activations_by_location_type_and_pass_type,
        batched_requested_attached_activation_for_backward_pass_by_name,
    )
This is a helper function that returns:
    1. a TransformerHookGraph object containing hooks for the given
    scalar derivers (this can be passed to Transformer using the as_transformer_hooks method to add
    hooks to the transformer forward and backward pass)
    2. dictionaries mapping each activation location type to the activations requested
    (before it is filled with activations during forward passes), one dictionary per batch element
    3. dictionaries each containing just one value, the scalar tensor on which the backward pass can be run (
    unlike the activations in 2, this tensor is still attached to the pytorch model), one dictionary per batch element
    
neuron_explainer/activation_server/derived_scalar_computation.py