def get_activation_indices_and_reconstitute_direct_effect_fn(
    model_context: ModelContext,
    multi_autoencoder_context: MultiAutoencoderContext | None,
    trace_config: TraceConfig,
    loss_fn_for_backward_pass: Callable[[torch.Tensor], torch.Tensor] | None,
) -> tuple[
    ActivationIndex,
    ActivationIndex,
    Callable[[torch.Tensor, torch.Tensor | None], torch.Tensor],
]:
    """
    For use with AttentionTraceType.V. This function returns the ActivationIndex corresponding to the
    residual stream before the upstream (attention V) node as well as the ActivationIndex corresponding to the
    residual stream before the downstream node, or before the loss. It also returns a function taking two arguments,
    both residual stream activations, which computes the direct effect of the upstream activation on the downstream
    activation.
    If trace_config.attention_trace_config.downstream_trace_config is a normal TraceConfig, then it specifies a downstream
    node's activation, which will be used to compute a gradient. If it is None, it is assumed that the output logits (in
    their entirety) are being reconstructed instead, and a loss function computed, to compute the gradient.
    """
    assert trace_config.attention_trace_type == AttentionTraceType.V
    if trace_config.downstream_trace_config is None:
        assert loss_fn_for_backward_pass is not None
        assert isinstance(model_context, StandardModelContext)  # for typechecking
        logit_reconstituter = LogitReconstituter(
            model_context=model_context,
            detach_layer_norm_scale=DETACH_LAYER_NORM_SCALE,
        )
        downstream_activation_index_for_reconstituter = (
            logit_reconstituter.get_residual_activation_index()
        )

        def reconstitute_gradient_fn(downstream_resid: torch.Tensor) -> torch.Tensor:
            reconstitute_gradient = logit_reconstituter.make_reconstitute_gradient_of_loss_fn(
                loss_fn=loss_fn_for_backward_pass
            )
            return reconstitute_gradient(downstream_resid)

    else:
        if trace_config.node_type.is_autoencoder_latent:
            assert multi_autoencoder_context is not None
            autoencoder_context = multi_autoencoder_context.get_autoencoder_context(
                trace_config.node_type
            )
            assert autoencoder_context is not None
        else:
            autoencoder_context = None
        act_reconstituter = ActivationReconstituter.from_trace_config(
            transformer=model_context.get_or_create_model(),
            autoencoder_context=autoencoder_context,
            trace_config=trace_config.downstream_trace_config,
        )
        downstream_trace_config = trace_config.downstream_trace_config
        downstream_activation_index_for_reconstituter = act_reconstituter.get_residual_activation_index_for_node_index(
            # convert trace_config to node index
            downstream_trace_config.node_index
        )

        def reconstitute_gradient_fn(
            downstream_resid: torch.Tensor,
        ) -> torch.Tensor:
            reconstitute_gradient_with_args = (
                act_reconstituter.make_reconstitute_gradient_fn_for_trace_config(
                    trace_config=downstream_trace_config
                )
            )
            return reconstitute_gradient_with_args(
                downstream_resid,
                downstream_trace_config.layer_index,
                downstream_trace_config.pass_type,
            )

    assert trace_config.layer_index is not None
    direct_effect_reconstituter = AttentionDirectEffectReconstituter(
        model_context=model_context,
        layer_indices=[trace_config.layer_index],
        detach_layer_norm_scale=DETACH_LAYER_NORM_SCALE,
    )
    upstream_activation_index_for_reconstituter = direct_effect_reconstituter.get_residual_activation_index_for_node_index(
        # convert trace_config to node index
        trace_config.node_index
    )

    upstream_scalar_hook = direct_effect_reconstituter.make_scalar_hook_for_node_index(
        trace_config.node_index
    )

    def reconstitute_direct_effect_fn(
        upstream_resid: torch.Tensor,
        downstream_resid: torch.Tensor | None,
    ) -> torch.Tensor:
        assert downstream_resid is not None
        gradient = reconstitute_gradient_fn(downstream_resid).detach()
        activations = direct_effect_reconstituter.reconstitute_activations(
            resid=upstream_resid,
            grad=gradient,
            layer_index=trace_config.layer_index,
            pass_type=trace_config.pass_type,
        )
        return upstream_scalar_hook(activations)

    return (
        upstream_activation_index_for_reconstituter,
        downstream_activation_index_for_reconstituter,
        reconstitute_direct_effect_fn,
    )

    For use with AttentionTraceType.V. This function returns the ActivationIndex corresponding to the
    residual stream before the upstream (attention V) node as well as the ActivationIndex corresponding to the
    residual stream before the downstream node, or before the loss. It also returns a function taking two arguments,
    both residual stream activations, which computes the direct effect of the upstream activation on the downstream
    activation.
    If trace_config.attention_trace_config.downstream_trace_config is a normal TraceConfig, then it specifies a downstream
    node's activation, which will be used to compute a gradient. If it is None, it is assumed that the output logits (in
    their entirety) are being reconstructed instead, and a loss function computed, to compute the gradient.
    
neuron_explainer/activation_server/derived_scalar_computation.py