def replace_activation_index_using_reconstituter(
    model_context: ModelContext,
    multi_autoencoder_context: MultiAutoencoderContext | None,
    batched_ds_computation_params: list[DerivedScalarComputationParams],
) -> tuple[
    list[ActivationIndex | None],
    list[ActivationIndex | None],
    list[Callable[[torch.Tensor, torch.Tensor | None], torch.Tensor]],
]:
    """
    Where trace_config occurs in batched_ds_computation_params, convert it to the
    upstream_activation_index (and optionally also downstream_activation_index)
    corresponding to the preceding residual stream required by a Reconstituter. Also
    return a function, generated from the Reconstituter, to obtain the activation
    corresponding to the original trace_config from the residual stream.
    """
    batched_reconstitute_activation_fn: list[
        Callable[[torch.Tensor, torch.Tensor | None], torch.Tensor]
    ] = []
    batched_upstream_activation_index_to_grab: list[ActivationIndex | None] = []
    batched_downstream_activation_index_to_grab: list[ActivationIndex | None] = []
    for ds_computation_params_index in range(len(batched_ds_computation_params)):
        ds_computation_params = batched_ds_computation_params[ds_computation_params_index]
        trace_config = ds_computation_params.trace_config
        if trace_config is not None:
            if trace_config.attention_trace_type == AttentionTraceType.V:
                (
                    upstream_activation_index_for_reconstituter,
                    downstream_activation_index_for_reconstituter,
                    reconstitute_activation_fn,
                ) = get_activation_indices_and_reconstitute_direct_effect_fn(
                    model_context=model_context,
                    multi_autoencoder_context=multi_autoencoder_context,
                    trace_config=trace_config,
                    loss_fn_for_backward_pass=ds_computation_params.loss_fn_for_backward_pass,
                )
            else:
                (
                    upstream_activation_index_for_reconstituter,
                    reconstitute_activation_fn,
                ) = get_activation_index_and_reconstitute_activation_fn(
                    transformer=model_context.get_or_create_model(),
                    multi_autoencoder_context=multi_autoencoder_context,
                    trace_config=trace_config,
                )
                downstream_activation_index_for_reconstituter = None
        else:
            upstream_activation_index_for_reconstituter = None
            downstream_activation_index_for_reconstituter = None

            def dummy_reconstitute_activation_fn(
                resid: torch.Tensor, grad: torch.Tensor | None
            ) -> torch.Tensor:
                raise NotImplementedError("This function should not be called")

            reconstitute_activation_fn = dummy_reconstitute_activation_fn
        if upstream_activation_index_for_reconstituter is not None:
            assert (
                upstream_activation_index_for_reconstituter.activation_location_type.node_type
                == NodeType.RESIDUAL_STREAM_CHANNEL
            )
        if downstream_activation_index_for_reconstituter is not None:
            assert (
                downstream_activation_index_for_reconstituter.activation_location_type.node_type
                == NodeType.RESIDUAL_STREAM_CHANNEL
            )
        batched_upstream_activation_index_to_grab.append(
            upstream_activation_index_for_reconstituter
        )
        batched_downstream_activation_index_to_grab.append(
            downstream_activation_index_for_reconstituter
        )
        batched_reconstitute_activation_fn.append(reconstitute_activation_fn)
    return (
        batched_upstream_activation_index_to_grab,
        batched_downstream_activation_index_to_grab,
        batched_reconstitute_activation_fn,
    )

    Where trace_config occurs in batched_ds_computation_params, convert it to the
    upstream_activation_index (and optionally also downstream_activation_index)
    corresponding to the preceding residual stream required by a Reconstituter. Also
    return a function, generated from the Reconstituter, to obtain the activation
    corresponding to the original trace_config from the residual stream.
    
neuron_explainer/activation_server/derived_scalar_computation.py