def get_location_within_layer_for_dst(
    dst: DerivedScalarType,
    dst_config: DstConfig,
) -> LocationWithinLayer | None:
    """Determines the location within a layer for DSTs which are not associated with an activation
    location type, or whose location within a layer depends on information in the DstConfig (e.g.
    autoencoder related DSTs). Defining new direct write related DSTs may require additional entries
    here."""
    if dst.location_within_layer is not None:
        # this might be determinable from the DST alone, in which case return it right away
        return dst.location_within_layer
    else:
        match dst.node_type:
            case (
                NodeType.AUTOENCODER_LATENT
                | NodeType.MLP_AUTOENCODER_LATENT
                | NodeType.ATTENTION_AUTOENCODER_LATENT
            ):
                autoencoder_context = dst_config.get_autoencoder_context(dst.node_type)
                if autoencoder_context is not None:
                    return autoencoder_context.dst.location_within_layer
                else:
                    return None
            case NodeType.RESIDUAL_STREAM_CHANNEL:
                match dst:
                    case DerivedScalarType.ATTN_WRITE:
                        return LocationWithinLayer.ATTN
                    case DerivedScalarType.PREVIOUS_LAYER_RESID_POST_MLP:
                        return LocationWithinLayer.END_OF_PREV_LAYER
                    case _:
                        return None
            case _:
                return None
Determines the location within a layer for DSTs which are not associated with an activation
    location type, or whose location within a layer depends on information in the DstConfig (e.g.
    autoencoder related DSTs). Defining new direct write related DSTs may require additional entries
    here.
neuron_explainer/activations/derived_scalars/locations.py