"""Helpers for residual-stream sites that are not always direct modules."""

from __future__ import annotations

from torch import nn

from ..models.adapters import ModelAdapter


def resolve_post_attention_anchor(adapter: ModelAdapter, layer: nn.Module) -> nn.Module:
    """
    Resolve the module whose input corresponds to the residual stream after attention.

    For architectures with a dedicated pre-MLP norm, the input to that norm is the
    cleanest place to observe the post-attention residual stream.
    """

    anchor = adapter.get_pre_mlp_norm(layer)
    if anchor is None:
        raise ValueError(
            "This model adapter does not expose a clean post-attention residual site."
        )
    return anchor

