def make_final_residual_grad_scalar_source(
    dst_config: DstConfig,
    use_backward_pass: bool,
) -> ScalarSource:
    """
    Many DSTs depend on the residual stream gradient at the last point in the forward pass before the point
    from which the backward pass is run. There are two ways of deriving this residual stream gradient.

    Background on backward passes:
    By default, the backward pass is run starting from some scalar function of the transformer's
    output logits. In this case, the last relevant point in the forward pass is at the very last
    residual stream location in the network (pre- final layer norm).
    A backward pass can also be run from an arbitrary activation in the network. In this case, the
    last relevant point in the forward pass is at the residual stream location immediately preceding
    the layer index of the activation from which the backward pass is run (pre- layer norm for that
    layer).
    The DstConfig object specifies whether the backward pass is the default (trace_config=None)
    or from an activation (trace_config=TraceConfig(node_index=NodeIndex(),...)).
    Note that if all you care about for a particular DST is the gradient at the last point in the forward pass
    (i.e. the first point in the backward pass), then running the full backward pass is actually wasteful.
    If you need to compute gradients with respect to many different activations, it's best just to run the very
    first part of the backward pass if possible. This is what use_backward_pass=False does.

    Two ways of deriving the residual stream gradient:
     - use_backward_pass=True: assume a literal backward pass has been run, outside the DST setup, as specified
        by the DstConfig object. In this case, you can directly use the "raw" residual stream gradient at a location
        inferrable from dst_config.trace_config.
     - use_backward_pass=False: (specific to the case where trace_config is not None)
        do not make assumptions about the literal backward pass that has been run. Take
        the residual stream **activations** (the forward pass) at the location implied by
        dst_config.trace_config. Recompute the activation specified from those residual
        stream activations, and run a small backward pass on the activation, back to those residual stream
        activations.
    """
    if use_backward_pass:
        return make_backward_pass_scalar_source_for_final_residual_grad(dst_config)
    else:
        return make_no_backward_pass_scalar_source_for_final_residual_grad(dst_config)

    Many DSTs depend on the residual stream gradient at the last point in the forward pass before the point
    from which the backward pass is run. There are two ways of deriving this residual stream gradient.

    Background on backward passes:
    By default, the backward pass is run starting from some scalar function of the transformer's
    output logits. In this case, the last relevant point in the forward pass is at the very last
    residual stream location in the network (pre- final layer norm).
    A backward pass can also be run from an arbitrary activation in the network. In this case, the
    last relevant point in the forward pass is at the residual stream location immediately preceding
    the layer index of the activation from which the backward pass is run (pre- layer norm for that
    layer).
    The DstConfig object specifies whether the backward pass is the default (trace_config=None)
    or from an activation (trace_config=TraceConfig(node_index=NodeIndex(),...)).
    Note that if all you care about for a particular DST is the gradient at the last point in the forward pass
    (i.e. the first point in the backward pass), then running the full backward pass is actually wasteful.
    If you need to compute gradients with respect to many different activations, it's best just to run the very
    first part of the backward pass if possible. This is what use_backward_pass=False does.

    Two ways of deriving the residual stream gradient:
     - use_backward_pass=True: assume a literal backward pass has been run, outside the DST setup, as specified
        by the DstConfig object. In this case, you can directly use the "raw" residual stream gradient at a location
        inferrable from dst_config.trace_config.
     - use_backward_pass=False: (specific to the case where trace_config is not None)
        do not make assumptions about the literal backward pass that has been run. Take
        the residual stream **activations** (the forward pass) at the location implied by
        dst_config.trace_config. Recompute the activation specified from those residual
        stream activations, and run a small backward pass on the activation, back to those residual stream
        activations.
    
neuron_explainer/activations/derived_scalars/direct_effects.py