Model Bridge

SafeLens' model bridge follows the same abstraction pattern that makes TransformerLens scale across many model families: an architecture adapter maps provider-specific module paths onto canonical components such as resid_pre, attn_out, mlp_out, q, k, v, z, and result.

For Transformers-backed models, result is exposed as a TransformerLens-style derived activation. SafeLens captures the input to the attention output projection as z and computes per-head residual-space result vectors with z @ W_O. When patching result, SafeLens calls the user hook on those per-head vectors, sums the per-head residual-space delta, and writes that delta back to the merged attention projection output.

The bridge is independent of TransformerLens. It is used by Transformers-backed wrappers after the model is loaded.

Current architecture adapters:

• llama_like_decoder: Qwen, Qwen2, Qwen3, LLaMA, Mistral, Mixtral, Gemma, OLMo, StableLM, and Yi-style model.layers decoders.
• gpt2_decoder: GPT-2 and DistilGPT2-style transformer.h decoders.
• gpt_neox_decoder: GPT-NeoX and Pythia-style gpt_neox.layers decoders.
• gptj_decoder: GPT-J-style decoder blocks.
• gpt_neo_decoder: GPT-Neo-style decoder blocks.
• joint_qkv_decoder: BLOOM and Falcon-style joint QKV decoders.
• mpt_decoder: MPT-style decoder blocks.
• phi_decoder: Phi-style decoder blocks.
• opt_decoder: OPT-style decoder layers.
• bert_encoder: BERT/RoBERTa-style encoder layers.
• distilbert_encoder: DistilBERT encoder layers.
• audio_encoder: Wav2Vec2/Hubert encoder layers.
• t5_encoder_decoder: initial T5 encoder-stack component mapping.

Attention pattern caching uses returned attention weights when available. Attention pattern patching and raw pre-softmax score hooks use eager softmax instrumentation. If a model runs flash attention or SDPA without a Python torch.softmax call, SafeLens raises a clear error and the model should be run with an eager attention implementation for those hooks.

List the registered architecture adapters:

safelens models list-architectures
safelens models list-architectures --json

Architecture bridge primitives for Transformers-backed model adapters.

The design mirrors the useful part of TransformerLens' model bridge: keep model loading provider-specific, but map each model family into a small canonical component vocabulary that SafeLens hooks and patching code can target.

ArchitectureAdapter

Map one architecture family onto SafeLens' canonical component names.

Source code in src/SafeLens/utils/model_bridge.py

class ArchitectureAdapter:
    """Map one architecture family onto SafeLens' canonical component names."""

    def __init__(
        self,
        *,
        name: str,
        model_types: Sequence[str],
        component_specs: Sequence[ComponentHookSpec],
        model_name_markers: Sequence[str] = (),
        notes: Sequence[str] = (),
    ) -> None:
        self.name = name
        self.model_types = tuple(model_types)
        self.model_name_markers = tuple(marker.lower() for marker in model_name_markers)
        self.notes = tuple(notes)
        self._specs = {spec.component: spec for spec in component_specs}
        self._aliases: dict[str, str] = {}
        for spec in component_specs:
            for alias in spec.all_names():
                self._aliases[alias] = spec.component

    def supports_model(self, *, model_type: str | None, model_name: str) -> bool:
        lowered_model_type = (model_type or "").lower()
        lowered_model_name = model_name.lower()
        return lowered_model_type in self.model_types or any(
            marker in lowered_model_name for marker in self.model_name_markers
        )

    def supported_components(
        self,
        *,
        include_unsupported: bool = False,
        for_cache: bool | None = None,
    ) -> tuple[str, ...]:
        components: list[str] = []
        for spec in self._specs.values():
            if not include_unsupported and not spec.supported:
                continue
            if for_cache is True and not spec.cacheable:
                continue
            if for_cache is False and not spec.patchable:
                continue
            if for_cache is None and not include_unsupported and not spec.patchable:
                continue
            components.append(spec.component)
        return tuple(components)

    def inspect(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "model_types": list(self.model_types),
            "model_name_markers": list(self.model_name_markers),
            "supported_components": list(self.supported_components()),
            "cacheable_components": list(self.supported_components(for_cache=True)),
            "patchable_components": list(self.supported_components(for_cache=False)),
            "target_components": list(self.supported_components(include_unsupported=True)),
            "notes": list(self.notes),
        }

    def parse_component_ref(self, layer: LayerRef) -> ComponentRef | None:
        if isinstance(layer, int):
            return ComponentRef(layer=layer, component="resid_post", original=layer)
        if isinstance(layer, tuple):
            if not layer:
                return None
            component = str(layer[0])
            layer_index = layer[1] if len(layer) >= 2 else None
            layer_type = str(layer[2]) if len(layer) >= 3 and layer[2] is not None else None
            if not isinstance(layer_index, int):
                return None
            return self._make_ref(
                layer_index,
                _normalize_component(component, layer_type=layer_type),
                layer,
            )
        if not isinstance(layer, str):
            return None

        safe_match = re.fullmatch(r"layer_(\d+)\.([a-zA-Z0-9_]+)", layer)
        if safe_match is not None:
            return self._make_ref(
                int(safe_match.group(1)),
                safe_match.group(2),
                layer,
            )

        block_match = re.fullmatch(
            r"(blocks|encoder|decoder)\.(\d+)\.(?:([a-zA-Z0-9_]+)\.)?hook_([a-zA-Z0-9_]+)",
            layer,
        )
        if block_match is not None:
            stack = block_match.group(1)
            layer_type = block_match.group(3)
            return self._make_ref(
                int(block_match.group(2)),
                _normalize_component(
                    block_match.group(4),
                    layer_type=layer_type,
                    stack=stack if stack in {"encoder", "decoder"} else None,
                ),
                layer,
            )

        return None

    def register_component_hook(
        self,
        model: Any,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        return self.register_component_hook_for_mode(
            model,
            layer,
            hook_fn,
            for_cache=False,
            prepend=prepend,
        )

    def register_component_hook_for_mode(
        self,
        model: Any,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        for_cache: bool,
        prepend: bool = False,
    ) -> Any:
        component_ref = self.parse_component_ref(layer)
        if component_ref is None:
            raise KeyError(f"Layer reference {layer!r} is not a component hook name.")
        spec = self._spec_for_ref(component_ref, for_cache=for_cache)
        module = self.get_component(model, component_ref)
        if spec.value == "attention_scores" or (
            spec.value == "attention_pattern" and not for_cache
        ):
            return _register_attention_softmax_hook(
                module,
                hook_fn,
                component_ref,
                self.name,
                spec,
                prepend=prepend,
            )
        if _is_attention_result_component(spec.component) and not for_cache:
            attention_result_hook = _make_attention_result_output_hook(
                hook_fn,
                component_ref,
                self.name,
                spec,
                model,
            )
            return _ComponentHookHandle(
                _register_module_forward_hook(module, attention_result_hook, prepend=prepend),
                _hook_contexts_for(attention_result_hook),
            )
        if spec.value == "norm_scale" or spec.component.endswith("_normalized"):
            norm_output_hook = _make_component_output_hook(
                hook_fn, component_ref, self.name, spec, model
            )
            return _ComponentHookHandle(
                _register_module_forward_hook(module, norm_output_hook, prepend=prepend),
                _hook_contexts_for(norm_output_hook),
            )
        t5_input_handle = _try_register_t5_attention_input_patch(
            model,
            hook_fn,
            component_ref,
            self.name,
            spec,
            prepend=prepend,
            for_cache=for_cache,
        )
        if t5_input_handle is not None:
            return t5_input_handle
        if spec.mode == "forward_input":
            component_input_hook = _make_component_input_hook(
                hook_fn, component_ref, self.name, spec, model
            )
            return _ComponentHookHandle(
                _register_module_forward_pre_hook(module, component_input_hook, prepend=prepend),
                _hook_contexts_for(component_input_hook),
            )
        component_output_hook = _make_component_output_hook(
            hook_fn, component_ref, self.name, spec, model
        )
        return _ComponentHookHandle(
            _register_module_forward_hook(module, component_output_hook, prepend=prepend),
            _hook_contexts_for(component_output_hook),
        )

    def requires_output_attentions(self, layer: LayerRef) -> bool:
        component_ref = self.parse_component_ref(layer)
        if component_ref is None:
            return False
        spec = self._specs.get(component_ref.component)
        return spec is not None and spec.value in {"attention_pattern", "attention_scores"}

    def get_component(self, model: Any, component_ref: ComponentRef) -> Any:
        spec = self._spec_for_ref(component_ref, for_cache=True)
        attempted_paths: list[str] = []
        for template in spec.module_paths:
            path = template.format(layer=component_ref.layer)
            attempted_paths.append(path)
            try:
                return resolve_module_path(model, path)
            except (AttributeError, IndexError, KeyError, TypeError):
                continue
        attempted = ", ".join(attempted_paths)
        raise KeyError(
            f"Could not resolve component {component_ref.component!r} for architecture "
            f"{self.name!r}. Tried module paths: {attempted}."
        )

    def get_attention_weight(self, model: Any, component: str, layer: int) -> Any:
        """Return a TransformerLens-shaped attention weight tensor for one layer."""
        component_ref = self._make_ref(layer, component, component)
        if component_ref is None:
            raise KeyError(f"Unknown attention component {component!r}.")
        spec = self._spec_for_ref(component_ref, for_cache=True)
        module = self.get_component(model, component_ref)
        weight = getattr(module, "weight", None)
        if weight is None:
            raise KeyError(f"Component {component!r} at layer {layer} has no weight.")
        if spec.activation == "split_qkv_heads":
            base_component = _attention_base_component(component)
            return reshape_joint_qkv_attention_weight(
                weight,
                component=base_component,
                q_heads=head_count_for_component(model, "q"),
                kv_heads=head_count_for_component(model, "k"),
                qkv_layout=spec.qkv_layout,
                packed_axis=preferred_qkv_weight_packed_axis(module, architecture=self.name),
            )
        if spec.activation != "split_heads":
            raise NotImplementedError(
                f"{self.name} cannot expose W_{component.upper()} from "
                f"{spec.activation!r} projections yet."
            )
        base_component = _attention_base_component(component)
        n_heads = head_count_for_component(model, base_component)
        return reshape_attention_weight(
            weight,
            component=base_component,
            n_heads=n_heads,
            packed_axis=preferred_attention_weight_packed_axis(
                module,
                architecture=self.name,
                component=base_component,
            ),
        )

    def get_attention_bias(self, model: Any, component: str, layer: int) -> Any:
        """Return a TransformerLens-shaped attention bias for one layer."""
        component_ref = self._make_ref(layer, component, component)
        if component_ref is None:
            raise KeyError(f"Unknown attention component {component!r}.")
        spec = self._spec_for_ref(component_ref, for_cache=True)
        module = self.get_component(model, component_ref)
        bias = getattr(module, "bias", None)
        if bias is None:
            return zeros_for_attention_bias(model, component)
        if spec.activation == "split_qkv_heads":
            base_component = _attention_base_component(component)
            return reshape_joint_qkv_attention_bias(
                bias,
                component=base_component,
                q_heads=head_count_for_component(model, "q"),
                kv_heads=head_count_for_component(model, "k"),
                qkv_layout=spec.qkv_layout,
            )
        if spec.activation != "split_heads":
            raise NotImplementedError(
                f"{self.name} cannot expose b_{component.upper()} from "
                f"{spec.activation!r} projections yet."
            )
        base_component = _attention_base_component(component)
        return reshape_attention_bias(
            bias,
            component=base_component,
            n_heads=head_count_for_component(model, base_component),
        )

    def get_embedding_weight(self, model: Any, *, positional: bool = False) -> Any:
        """Return token or positional embedding weights for common Transformers layouts."""
        paths = _POSITION_EMBEDDING_MODULE_PATHS if positional else _TOKEN_EMBEDDING_MODULE_PATHS
        kind = "positional embedding" if positional else "token embedding"
        try:
            return _weight_from_first_path(model, paths, kind=kind)
        except KeyError as path_error:
            if positional:
                raise
            embedding_weight = _input_embedding_weight_from_model(model)
            if embedding_weight is not None:
                return embedding_weight
            raise path_error

    def get_mlp_weight(self, model: Any, component: str, layer: int) -> Any:
        """Return a TransformerLens-shaped MLP weight matrix for one layer."""
        if component == "in":
            paths = self._mlp_weight_paths(
                layer,
                canonical_component="pre_linear",
                fallback_component="pre",
            )
        elif component == "gate":
            self._raise_if_mlp_component_unsupported("pre")
            paths = (
                *self._mlp_weight_paths(layer, canonical_component="pre"),
                f"model.language_model.layers.{layer}.mlp.gate",
                f"model.language_model.layers.{layer}.mlp.w1",
                f"model.layers.{layer}.mlp.gate",
                f"model.layers.{layer}.mlp.w1",
            )
        elif component == "out":
            paths = self._mlp_weight_paths(layer, canonical_component="post")
        else:
            raise ValueError(f"Unsupported MLP weight component {component!r}.")
        module, weight = _module_weight_from_first_path(
            model,
            paths,
            kind=f"MLP {component} weight",
        )
        if _is_transformers_conv1d_module(module):
            return weight
        return transpose_2d_weight(weight)

    def get_mlp_bias(self, model: Any, component: str, layer: int) -> Any:
        """Return a TransformerLens-shaped MLP bias vector for one layer."""
        if component == "in":
            paths = self._mlp_weight_paths(
                layer,
                canonical_component="pre_linear",
                fallback_component="pre",
            )
        elif component == "out":
            paths = self._mlp_weight_paths(layer, canonical_component="post")
        else:
            raise ValueError(f"Unsupported MLP bias component {component!r}.")
        zero_axis = 1 if self.name == "gpt2_decoder" else 0
        return _bias_from_first_path(
            model,
            paths,
            kind=f"MLP {component} bias",
            zero_axis=zero_axis,
        )

    def _raise_if_mlp_component_unsupported(self, canonical_component: str) -> None:
        spec = self._specs.get(canonical_component)
        if spec is not None and not spec.supported:
            reason = spec.unsupported_reason or f"component {spec.component!r} is not supported"
            raise NotImplementedError(f"{self.name} does not expose {spec.component!r}: {reason}.")

    def _mlp_weight_paths(
        self,
        layer: int,
        *,
        canonical_component: str,
        fallback_component: str | None = None,
    ) -> tuple[str, ...]:
        spec = self._specs.get(canonical_component)
        if spec is None and fallback_component is not None:
            spec = self._specs.get(fallback_component)
        if spec is None:
            raise KeyError(f"{self.name!r} does not declare MLP component {canonical_component!r}.")
        if not spec.supported:
            reason = spec.unsupported_reason or f"component {spec.component!r} is not supported"
            raise NotImplementedError(f"{self.name} does not expose {spec.component!r}: {reason}.")
        return tuple(template.format(layer=layer) for template in spec.module_paths)

    def _make_ref(self, layer: int, component: str, original: LayerRef) -> ComponentRef | None:
        normalized = component if component in self._aliases else _normalize_component(component)
        if normalized not in self._aliases:
            return None
        canonical_component = self._aliases[normalized]
        spec = self._specs.get(canonical_component)
        transformer_lens_name_override = None
        if spec is not None and spec.transformer_lens_name_template is not None:
            transformer_lens_name_override = spec.transformer_lens_name_template.format(layer=layer)
        return ComponentRef(
            layer=layer,
            component=canonical_component,
            original=original,
            transformer_lens_name_override=transformer_lens_name_override,
        )

    def _spec_for_ref(self, component_ref: ComponentRef, *, for_cache: bool) -> ComponentHookSpec:
        spec = self._specs[component_ref.component]
        if not spec.supported:
            reason = spec.unsupported_reason or f"component {spec.component!r} is not supported"
            raise NotImplementedError(f"{self.name} does not expose {spec.component!r}: {reason}.")
        if for_cache and not spec.cacheable:
            raise NotImplementedError(f"{self.name} cannot cache {spec.component!r}.")
        if not for_cache and not spec.patchable:
            if spec.unsupported_reason:
                raise NotImplementedError(
                    f"{self.name} can cache {spec.component!r}, but cannot patch it: "
                    f"{spec.unsupported_reason}."
                )
            raise NotImplementedError(
                f"{self.name} can cache {spec.component!r}, but cannot patch it "
                "without instrumenting the attention computation before value mixing."
            )
        return spec

get_attention_bias(model, component, layer)

Return a TransformerLens-shaped attention bias for one layer.

Source code in src/SafeLens/utils/model_bridge.py

def get_attention_bias(self, model: Any, component: str, layer: int) -> Any:
    """Return a TransformerLens-shaped attention bias for one layer."""
    component_ref = self._make_ref(layer, component, component)
    if component_ref is None:
        raise KeyError(f"Unknown attention component {component!r}.")
    spec = self._spec_for_ref(component_ref, for_cache=True)
    module = self.get_component(model, component_ref)
    bias = getattr(module, "bias", None)
    if bias is None:
        return zeros_for_attention_bias(model, component)
    if spec.activation == "split_qkv_heads":
        base_component = _attention_base_component(component)
        return reshape_joint_qkv_attention_bias(
            bias,
            component=base_component,
            q_heads=head_count_for_component(model, "q"),
            kv_heads=head_count_for_component(model, "k"),
            qkv_layout=spec.qkv_layout,
        )
    if spec.activation != "split_heads":
        raise NotImplementedError(
            f"{self.name} cannot expose b_{component.upper()} from "
            f"{spec.activation!r} projections yet."
        )
    base_component = _attention_base_component(component)
    return reshape_attention_bias(
        bias,
        component=base_component,
        n_heads=head_count_for_component(model, base_component),
    )

get_attention_weight(model, component, layer)

Return a TransformerLens-shaped attention weight tensor for one layer.

Source code in src/SafeLens/utils/model_bridge.py

def get_attention_weight(self, model: Any, component: str, layer: int) -> Any:
    """Return a TransformerLens-shaped attention weight tensor for one layer."""
    component_ref = self._make_ref(layer, component, component)
    if component_ref is None:
        raise KeyError(f"Unknown attention component {component!r}.")
    spec = self._spec_for_ref(component_ref, for_cache=True)
    module = self.get_component(model, component_ref)
    weight = getattr(module, "weight", None)
    if weight is None:
        raise KeyError(f"Component {component!r} at layer {layer} has no weight.")
    if spec.activation == "split_qkv_heads":
        base_component = _attention_base_component(component)
        return reshape_joint_qkv_attention_weight(
            weight,
            component=base_component,
            q_heads=head_count_for_component(model, "q"),
            kv_heads=head_count_for_component(model, "k"),
            qkv_layout=spec.qkv_layout,
            packed_axis=preferred_qkv_weight_packed_axis(module, architecture=self.name),
        )
    if spec.activation != "split_heads":
        raise NotImplementedError(
            f"{self.name} cannot expose W_{component.upper()} from "
            f"{spec.activation!r} projections yet."
        )
    base_component = _attention_base_component(component)
    n_heads = head_count_for_component(model, base_component)
    return reshape_attention_weight(
        weight,
        component=base_component,
        n_heads=n_heads,
        packed_axis=preferred_attention_weight_packed_axis(
            module,
            architecture=self.name,
            component=base_component,
        ),
    )

get_embedding_weight(model, *, positional=False)

Return token or positional embedding weights for common Transformers layouts.

Source code in src/SafeLens/utils/model_bridge.py

def get_embedding_weight(self, model: Any, *, positional: bool = False) -> Any:
    """Return token or positional embedding weights for common Transformers layouts."""
    paths = _POSITION_EMBEDDING_MODULE_PATHS if positional else _TOKEN_EMBEDDING_MODULE_PATHS
    kind = "positional embedding" if positional else "token embedding"
    try:
        return _weight_from_first_path(model, paths, kind=kind)
    except KeyError as path_error:
        if positional:
            raise
        embedding_weight = _input_embedding_weight_from_model(model)
        if embedding_weight is not None:
            return embedding_weight
        raise path_error

get_mlp_bias(model, component, layer)

Return a TransformerLens-shaped MLP bias vector for one layer.

Source code in src/SafeLens/utils/model_bridge.py

def get_mlp_bias(self, model: Any, component: str, layer: int) -> Any:
    """Return a TransformerLens-shaped MLP bias vector for one layer."""
    if component == "in":
        paths = self._mlp_weight_paths(
            layer,
            canonical_component="pre_linear",
            fallback_component="pre",
        )
    elif component == "out":
        paths = self._mlp_weight_paths(layer, canonical_component="post")
    else:
        raise ValueError(f"Unsupported MLP bias component {component!r}.")
    zero_axis = 1 if self.name == "gpt2_decoder" else 0
    return _bias_from_first_path(
        model,
        paths,
        kind=f"MLP {component} bias",
        zero_axis=zero_axis,
    )

get_mlp_weight(model, component, layer)

Return a TransformerLens-shaped MLP weight matrix for one layer.

Source code in src/SafeLens/utils/model_bridge.py

def get_mlp_weight(self, model: Any, component: str, layer: int) -> Any:
    """Return a TransformerLens-shaped MLP weight matrix for one layer."""
    if component == "in":
        paths = self._mlp_weight_paths(
            layer,
            canonical_component="pre_linear",
            fallback_component="pre",
        )
    elif component == "gate":
        self._raise_if_mlp_component_unsupported("pre")
        paths = (
            *self._mlp_weight_paths(layer, canonical_component="pre"),
            f"model.language_model.layers.{layer}.mlp.gate",
            f"model.language_model.layers.{layer}.mlp.w1",
            f"model.layers.{layer}.mlp.gate",
            f"model.layers.{layer}.mlp.w1",
        )
    elif component == "out":
        paths = self._mlp_weight_paths(layer, canonical_component="post")
    else:
        raise ValueError(f"Unsupported MLP weight component {component!r}.")
    module, weight = _module_weight_from_first_path(
        model,
        paths,
        kind=f"MLP {component} weight",
    )
    if _is_transformers_conv1d_module(module):
        return weight
    return transpose_2d_weight(weight)

ComponentHookContext

Small TransformerLens-style hook object passed to component hooks.

Source code in src/SafeLens/utils/model_bridge.py

class ComponentHookContext:
    """Small TransformerLens-style hook object passed to component hooks."""

    def __init__(self, component_ref: ComponentRef) -> None:
        self.name = component_ref.transformer_lens_name
        self.component = component_ref.component
        self.ctx: dict[str, Any] = {}
        self._layer = component_ref.layer
        self.safelens_name = component_ref.safelens_name

    def layer(self) -> int:
        return self._layer

ComponentHookSpec dataclass

How one canonical component maps to a HuggingFace module path.

Source code in src/SafeLens/utils/model_bridge.py

@dataclass(frozen=True)
class ComponentHookSpec:
    """How one canonical component maps to a HuggingFace module path."""

    component: str
    mode: HookMode
    module_paths: tuple[str, ...]
    value: ComponentValue = "output"
    activation: ComponentActivation = "raw"
    qkv_layout: QKVLayout = "split"
    aliases: tuple[str, ...] = ()
    patchable: bool = True
    cacheable: bool = True
    supported: bool = True
    unsupported_reason: str | None = None
    transformer_lens_name_template: str | None = None

    def all_names(self) -> tuple[str, ...]:
        return (self.component, *self.aliases)

ComponentRef dataclass

One parsed reference to a canonical transformer component.

Source code in src/SafeLens/utils/model_bridge.py

@dataclass(frozen=True)
class ComponentRef:
    """One parsed reference to a canonical transformer component."""

    layer: int
    component: str
    original: LayerRef
    transformer_lens_name_override: str | None = None

    @property
    def safelens_name(self) -> str:
        return f"layer_{self.layer}.{self.component}"

    @property
    def transformer_lens_name(self) -> str:
        if self.transformer_lens_name_override is not None:
            return self.transformer_lens_name_override
        return transformer_lens_component_name(self.component, self.layer)

apply_attention_result_patch(output, original_result, patched_result)

Apply a patched per-head result tensor to the merged attention output.

Source code in src/SafeLens/utils/model_bridge.py

def apply_attention_result_patch(output: Any, original_result: Any, patched_result: Any) -> Any:
    """Apply a patched per-head `result` tensor to the merged attention output."""
    delta = subtract_values(patched_result, original_result)
    merged_delta = sum_attention_heads(delta)
    return add_values(output, merged_delta)

apply_norm_affine(module, normalized, reference_output=None)

Apply PyTorch/HF norm affine weights to a normalized activation.

Source code in src/SafeLens/utils/model_bridge.py

def apply_norm_affine(module: Any, normalized: Any, reference_output: Any | None = None) -> Any:
    """Apply PyTorch/HF norm affine weights to a normalized activation."""
    weight = _first_existing_attr(module, "weight", "w")
    bias = _first_existing_attr(module, "bias", "b")
    try:
        import torch

        if isinstance(normalized, torch.Tensor):
            output = normalized
            if weight is not None:
                output = output * weight.to(dtype=output.dtype, device=output.device)
            if bias is not None:
                output = output + bias.to(dtype=output.dtype, device=output.device)
            if reference_output is not None and isinstance(reference_output, torch.Tensor):
                output = output.to(dtype=reference_output.dtype, device=reference_output.device)
            return output
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(normalized, "shape"):
            numpy_output = np.asarray(normalized)
            if weight is not None:
                numpy_output = numpy_output * np.asarray(weight)
            if bias is not None:
                numpy_output = numpy_output + np.asarray(bias)
            return numpy_output
    except Exception:
        pass
    output = normalized
    if weight is not None:
        output = _multiply_last_dim_nested(output, weight)
    if bias is not None:
        output = _add_last_dim_nested(output, bias)
    return output

architecture_adapter_for_model(model, *, model_name='')

Select a SafeLens architecture adapter for a loaded Transformers model.

Source code in src/SafeLens/utils/model_bridge.py

def architecture_adapter_for_model(model: Any, *, model_name: str = "") -> ArchitectureAdapter:
    """Select a SafeLens architecture adapter for a loaded Transformers model."""
    config = _model_config(model)
    model_type = _config_attr(config, "model_type")
    return architecture_adapter_for_name(model_name=model_name, model_type=model_type)

architecture_adapter_for_name(*, model_name, model_type=None)

Select an architecture adapter from a model name and optional HF model_type.

Source code in src/SafeLens/utils/model_bridge.py

def architecture_adapter_for_name(
    *,
    model_name: str,
    model_type: str | None = None,
) -> ArchitectureAdapter:
    """Select an architecture adapter from a model name and optional HF model_type."""
    resolved_model_name = resolve_transformer_lens_compatible_model_name(model_name)
    lowered_model_type = (model_type or "").lower()
    if lowered_model_type in {"", "qwen2_moe", "qwen3_moe"} and is_qwen_routed_moe_model_name(
        resolved_model_name
    ):
        return ROUTED_MOE_ADAPTER
    for adapter in SUPPORTED_ARCHITECTURE_ADAPTERS:
        if adapter.supports_model(model_type=model_type, model_name=resolved_model_name):
            return adapter
    return GENERIC_DECODER_ADAPTER

attention_head_count(model)

Read the configured query/output attention head count.

Source code in src/SafeLens/utils/model_bridge.py

def attention_head_count(model: Any) -> int | None:
    """Read the configured query/output attention head count."""
    config = _model_config(model)
    for name in ("num_attention_heads", "n_head", "n_heads", "num_heads"):
        value = _config_attr(config, name)
        if value is not None:
            return int(value)
    return None

attention_head_dim(model)

Read the per-query-head dimension used by q/k/v projections.

Source code in src/SafeLens/utils/model_bridge.py

def attention_head_dim(model: Any) -> int | None:
    """Read the per-query-head dimension used by q/k/v projections."""
    config = _model_config(model)
    for name in ("head_dim", "d_head", "d_kv", "kv_channels"):
        value = _config_attr(config, name)
        if value is not None:
            return int(value)
    d_model = _model_hidden_size(model)
    n_heads = attention_head_count(model)
    if d_model is not None and n_heads:
        return d_model // n_heads
    return None

call_component_hook(hook_fn, *, activation, component_ref, architecture, hook_context=None)

Call a user hook with SafeLens component metadata when accepted.

Source code in src/SafeLens/utils/model_bridge.py

def call_component_hook(
    hook_fn: HookFn,
    *,
    activation: Any,
    component_ref: ComponentRef,
    architecture: str,
    hook_context: ComponentHookContext | None = None,
) -> Any:
    """Call a user hook with SafeLens component metadata when accepted."""
    if hook_context is None:
        hook_context = ComponentHookContext(component_ref)
    hook_kwargs = {
        "activation": activation,
        "output": activation,
        "component": component_ref.component,
        "layer": component_ref.layer,
        "hook_name": component_ref.safelens_name,
        "transformer_lens_name": component_ref.transformer_lens_name,
        "architecture": architecture,
        "hook": hook_context,
    }
    return call_user_hook(
        hook_fn,
        hook_kwargs,
        positional_arg_options=((activation, hook_context), (activation,)),
    )

compute_attention_result_activation(activation, model, spec, *, module, component_ref, architecture)

Compute TransformerLens-style per-head result from merged z input.

Source code in src/SafeLens/utils/model_bridge.py

def compute_attention_result_activation(
    activation: Any,
    model: Any,
    spec: ComponentHookSpec,
    *,
    module: Any | None,
    component_ref: ComponentRef | None,
    architecture: str | None,
) -> Any:
    """Compute TransformerLens-style per-head `result` from merged `z` input."""
    if module is None:
        raise ValueError("Computing attention result activations requires the output module.")
    if component_ref is None:
        raise ValueError("Computing attention result activations requires a component ref.")
    weight = getattr(module, "weight", None)
    if weight is None:
        raise KeyError(
            f"Cannot compute {component_ref.safelens_name!r}: output projection has no weight."
        )
    base_component = _attention_base_component(spec.component)
    z_component = "z" if base_component == "result" else base_component
    n_heads = head_count_for_component(model, z_component)
    z_activation = split_heads(activation, n_heads)
    W_O = reshape_attention_weight(
        weight,
        component="z",
        n_heads=n_heads,
        packed_axis=preferred_attention_weight_packed_axis(
            module,
            architecture=architecture or "",
            component="z",
        ),
    )
    from SafeLens.core.analysis import compute_head_results_from_z

    return compute_head_results_from_z(z_activation, W_O)

extract_component_activation(output, spec, model)

Extract the activation value for a component from a module hook output.

Source code in src/SafeLens/utils/model_bridge.py

def extract_component_activation(output: Any, spec: ComponentHookSpec, model: Any) -> Any:
    """Extract the activation value for a component from a module hook output."""
    if spec.value == "output":
        return transform_component_activation(first_output(output), spec, model)
    if spec.value == "attention_pattern":
        pattern = _find_attention_pattern(output)
        if pattern is None:
            raise RuntimeError(
                f"Could not find attention pattern in output for component {spec.component!r}. "
                "Ensure the forward pass was called with output_attentions=True and the "
                "selected Transformers attention implementation returns attention weights."
            )
        return pattern
    if spec.value == "attention_scores":
        scores = _find_attention_scores(output)
        if scores is None:
            raise RuntimeError(
                f"Could not find attention scores in output for component {spec.component!r}. "
                "Use eager attention softmax instrumentation for pre-softmax scores."
            )
        return scores
    if spec.value == "norm_scale":
        scale = norm_scale_from_input(model, inputs=(), output=output)
        if scale is None:
            raise RuntimeError(
                f"Could not compute normalization scale for component {spec.component!r}."
            )
        return scale
    return output

first_attention_head(value)

Select the first head from a TransformerLens input tensor.

Source code in src/SafeLens/utils/model_bridge.py

def first_attention_head(value: Any) -> Any:
    """Select the first head from a TransformerLens input tensor."""
    try:
        import torch

        if hasattr(value, "shape") and isinstance(value, torch.Tensor):
            return value.select(dim=-2, index=0)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(value, "shape"):
            return np.asarray(value).take(0, axis=-2)
    except Exception:
        pass
    shape = getattr(value, "shape", None)
    if shape is not None:
        value = value.tolist()
    return _select_nested_penultimate_dim(value, 0)

first_output(output)

Return the tensor payload from common Transformers module outputs.

Source code in src/SafeLens/utils/model_bridge.py

def first_output(output: Any) -> Any:
    """Return the tensor payload from common Transformers module outputs."""
    if isinstance(output, tuple):
        return output[0]
    if _is_structured_list_output(output):
        return output[0]
    return output

head_count_for_component(model, component)

Read the configured attention head count for one component.

Source code in src/SafeLens/utils/model_bridge.py

def head_count_for_component(model: Any, component: str) -> int:
    """Read the configured attention head count for one component."""
    component = _attention_base_component(component)
    if component in {"k", "v"}:
        n_key_value_heads = key_value_head_count(model)
        if n_key_value_heads is not None:
            return n_key_value_heads
    n_heads = attention_head_count(model)
    if n_heads is not None:
        return n_heads
    config = getattr(model, "config", None)
    raise ValueError(
        f"Could not infer attention head count for component {component!r} "
        f"from {type(config).__name__}."
    )

is_qwen_routed_moe_model_name(model_name)

Return whether a Qwen model name clearly denotes a routed MoE checkpoint.

Source code in src/SafeLens/utils/model_bridge.py

def is_qwen_routed_moe_model_name(model_name: str) -> bool:
    """Return whether a Qwen model name clearly denotes a routed MoE checkpoint."""
    lowered = resolve_transformer_lens_compatible_model_name(model_name).lower()
    if "qwen" not in lowered:
        return False
    return "moe" in lowered or re.search(r"[-_/]a\d+(?:\.\d+)?b(?:[-_/]|$)", lowered) is not None

key_value_head_count(model)

Read the configured key/value head count when it differs from query heads.

Source code in src/SafeLens/utils/model_bridge.py

def key_value_head_count(model: Any) -> int | None:
    """Read the configured key/value head count when it differs from query heads."""
    config = _model_config(model)
    if _is_falcon_multi_query_config(config):
        return 1
    for name in ("num_key_value_heads", "num_kv_heads", "n_head_kv"):
        value = _config_attr(config, name)
        if value is not None:
            return int(value)
    return None

list_architecture_adapters()

Return public metadata for SafeLens' architecture bridge adapters.

Source code in src/SafeLens/utils/model_bridge.py

def list_architecture_adapters() -> list[dict[str, Any]]:
    """Return public metadata for SafeLens' architecture bridge adapters."""
    return [adapter.inspect() for adapter in SUPPORTED_ARCHITECTURE_ADAPTERS]

merge_component_activation(activation, reference, spec, model)

Merge a patched SafeLens component activation back into the raw HF tensor shape.

Source code in src/SafeLens/utils/model_bridge.py

def merge_component_activation(
    activation: Any,
    reference: Any,
    spec: ComponentHookSpec,
    model: Any,
) -> Any:
    """Merge a patched SafeLens component activation back into the raw HF tensor shape."""
    if spec.activation == "split_heads":
        return merge_heads(activation, reference)
    if spec.activation == "split_qkv_heads":
        return merge_qkv_heads(activation, reference, model, spec)
    if spec.activation == "repeat_heads":
        return first_attention_head(activation)
    return activation

merge_heads(activation, reference)

Flatten [batch, pos, head, head_dim] back to the reference final dimension.

Source code in src/SafeLens/utils/model_bridge.py

def merge_heads(activation: Any, reference: Any) -> Any:
    """Flatten `[batch, pos, head, head_dim]` back to the reference final dimension."""
    shape = getattr(activation, "shape", None)
    reshape = getattr(activation, "reshape", None)
    reference_shape = getattr(reference, "shape", None)
    if shape is None or not callable(reshape):
        if not _is_sequence(activation):
            return activation
        nested_shape = _nested_shape(activation)
        if len(nested_shape) < 4:
            return activation
        return _merge_nested_last_two_dims(activation)
    if len(shape) < 4:
        return activation
    hidden_size = int(shape[-2]) * int(shape[-1])
    if reference_shape is not None:
        hidden_size = int(reference_shape[-1])
    return activation.reshape(*shape[:-2], hidden_size)

merge_qkv_heads(activation, reference, model, spec)

Replace one split-head component in a raw joint QKV projection tensor.

Source code in src/SafeLens/utils/model_bridge.py

def merge_qkv_heads(
    activation: Any,
    reference: Any,
    model: Any,
    spec: ComponentHookSpec,
) -> Any:
    """Replace one split-head component in a raw joint QKV projection tensor."""
    component = _attention_base_component(spec.component)
    q_heads = head_count_for_component(model, "q")
    kv_heads = head_count_for_component(model, "k")
    if spec.qkv_layout == "interleaved":
        return merge_interleaved_qkv_heads(
            activation,
            reference,
            component,
            q_heads=q_heads,
            kv_heads=kv_heads,
        )
    q_slice, k_slice, v_slice = split_qkv_slices(reference, q_heads=q_heads, kv_heads=kv_heads)
    component_slice = {"q": q_slice, "k": k_slice, "v": v_slice}[component]
    merged_component = merge_heads(activation, component_slice)
    patched = clone_tensor_like(reference)
    target = {"q": 0, "k": 1, "v": 2}[component]
    start, stop = split_qkv_slice_bounds(reference, q_heads=q_heads, kv_heads=kv_heads)[target]
    if getattr(patched, "shape", None) is None:
        return _replace_nested_last_dim_slice(patched, start, stop, merged_component)
    patched[..., start:stop] = merged_component
    return patched

module_uses_centered_layer_norm(module)

Return whether a norm module mean-centers like LayerNorm.

Source code in src/SafeLens/utils/model_bridge.py

def module_uses_centered_layer_norm(module: Any) -> bool:
    """Return whether a norm module mean-centers like LayerNorm."""
    try:
        import torch

        if isinstance(module, torch.nn.LayerNorm):
            return True
    except Exception:
        pass
    class_name = type(module).__name__.lower()
    if "rms" in class_name:
        return False
    return "layernorm" in class_name or "layer_norm" in class_name or class_name == "layernorm"

norm_module_output_from_scale(module, source, scale, reference_output)

Recompute a norm module output from a patched TL-style scale.

Source code in src/SafeLens/utils/model_bridge.py

def norm_module_output_from_scale(
    module: Any, source: Any, scale: Any, reference_output: Any
) -> Any:
    """Recompute a norm module output from a patched TL-style scale."""
    normalized = normalized_output_from_scale(module, source, scale)
    return apply_norm_affine(module, normalized, reference_output)

norm_scale_from_input(module, inputs=(), output=None)

Return TransformerLens-style norm scale [batch, pos, 1] for a norm module input.

Source code in src/SafeLens/utils/model_bridge.py

def norm_scale_from_input(
    module: Any, inputs: Sequence[Any] = (), output: Any | None = None
) -> Any | None:
    """Return TransformerLens-style norm scale `[batch, pos, 1]` for a norm module input."""
    source = inputs[0] if inputs else output
    if source is None:
        return None
    try:
        import torch

        if isinstance(source, torch.Tensor):
            epsilon = float(getattr(module, "variance_epsilon", getattr(module, "eps", 1e-5)))
            source_float = source.float() if not torch.is_floating_point(source) else source
            if module_uses_centered_layer_norm(module):
                source_float = source_float - source_float.mean(dim=-1, keepdim=True)
            return torch.sqrt(source_float.pow(2).mean(dim=-1, keepdim=True) + epsilon).to(
                dtype=source.dtype,
                device=source.device,
            )
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(source, "shape") and type(source).__module__.split(".")[0] == "numpy":
            array = np.asarray(source)
            epsilon = float(getattr(module, "variance_epsilon", getattr(module, "eps", 1e-5)))
            if module_uses_centered_layer_norm(module):
                array = array - array.mean(axis=-1, keepdims=True)
            return np.sqrt(np.mean(array * array, axis=-1, keepdims=True) + epsilon)
    except Exception:
        pass
    shape = _nested_shape(source)
    if len(shape) < 1:
        return None
    epsilon = float(getattr(module, "variance_epsilon", getattr(module, "eps", 1e-5)))

    def scale_vector(vector: Any) -> list[float]:
        values = [float(item) for item in vector]
        if module_uses_centered_layer_norm(module):
            mean = sum(values) / max(1, len(values))
            values = [value - mean for value in values]
        variance = sum(value * value for value in values) / max(1, len(values))
        return [math.sqrt(variance + epsilon)]

    return _map_nested_vectors(source, scale_vector)

normalized_output_from_scale(module, source, scale=None)

Return normalized norm input before affine weights, matching TL hook_normalized.

Source code in src/SafeLens/utils/model_bridge.py

def normalized_output_from_scale(module: Any, source: Any, scale: Any | None = None) -> Any:
    """Return normalized norm input before affine weights, matching TL hook_normalized."""
    if scale is None:
        scale = norm_scale_from_input(module, inputs=(source,), output=None)
    try:
        import torch

        if isinstance(source, torch.Tensor):
            if not isinstance(scale, torch.Tensor):
                scale = torch.as_tensor(scale, dtype=source.dtype, device=source.device)
            else:
                scale = scale.to(dtype=source.dtype, device=source.device)
            source_float = source.float() if not torch.is_floating_point(source) else source
            if module_uses_centered_layer_norm(module):
                source_float = source_float - source_float.mean(dim=-1, keepdim=True)
            return (source_float / scale).to(dtype=source.dtype, device=source.device)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(source, "shape") or hasattr(scale, "shape"):
            array = np.asarray(source)
            if module_uses_centered_layer_norm(module):
                array = array - array.mean(axis=-1, keepdims=True)
            return array / np.asarray(scale)
    except Exception:
        pass
    if scale is None:
        return source
    return _divide_by_scale_nested(source, scale, centered=module_uses_centered_layer_norm(module))

preferred_attention_weight_packed_axis(module, *, architecture, component)

Return the likely packed axis for architecture-specific attention weights.

Source code in src/SafeLens/utils/model_bridge.py

def preferred_attention_weight_packed_axis(
    module: Any,
    *,
    architecture: str,
    component: str,
) -> int | None:
    """Return the likely packed axis for architecture-specific attention weights."""
    module_type = type(module).__name__.lower()
    if component == "z":
        if architecture == "gpt2_decoder" or module_type == "conv1d":
            return 0
        return 1
    if module_type == "conv1d":
        return 1
    return 0

preferred_qkv_weight_packed_axis(module, *, architecture)

Return the likely packed axis for architecture-specific joint QKV weights.

Source code in src/SafeLens/utils/model_bridge.py

def preferred_qkv_weight_packed_axis(module: Any, *, architecture: str) -> int | None:
    """Return the likely packed axis for architecture-specific joint QKV weights."""
    if architecture == "gpt2_decoder":
        return 1
    module_type = type(module).__name__.lower()
    if module_type == "conv1d":
        return 1
    return None

replace_component_activation(output, patched, spec, model)

Replace the tensor payload while preserving tuple/list module output shape.

Source code in src/SafeLens/utils/model_bridge.py

def replace_component_activation(
    output: Any,
    patched: Any,
    spec: ComponentHookSpec,
    model: Any,
) -> Any:
    """Replace the tensor payload while preserving tuple/list module output shape."""
    if spec.value != "output":
        return output
    raw_output = first_output(output)
    merged = merge_component_activation(patched, raw_output, spec, model)
    if isinstance(output, tuple):
        return (merged, *output[1:])
    if _is_structured_list_output(output):
        return [merged, *output[1:]]
    return merged

reshape_attention_bias(bias, *, component, n_heads)

Convert a packed projection bias to TransformerLens attention bias shape.

Source code in src/SafeLens/utils/model_bridge.py

def reshape_attention_bias(
    bias: Any,
    *,
    component: str,
    n_heads: int,
) -> Any:
    """Convert a packed projection bias to TransformerLens attention bias shape."""
    component = _attention_base_component(component)
    if component == "z":
        return bias
    shape = getattr(bias, "shape", None)
    reshape = getattr(bias, "reshape", None)
    if shape is None or not callable(reshape):
        if not _is_sequence(bias):
            return bias
        if n_heads <= 0 or len(bias) % n_heads != 0:
            raise ValueError(
                f"Cannot split attention bias of length {len(bias)} into {n_heads} heads."
            )
        head_dim = len(bias) // n_heads
        return [
            list(bias[head_index * head_dim : (head_index + 1) * head_dim])
            for head_index in range(n_heads)
        ]
    if len(shape) == 2:
        return bias
    if len(shape) != 1:
        raise ValueError(f"Cannot reshape attention bias for component {component!r}.")
    if n_heads <= 0 or int(shape[0]) % n_heads != 0:
        raise ValueError(
            f"Cannot split attention bias of length {int(shape[0])} into {n_heads} heads."
        )
    return bias.reshape(n_heads, int(shape[0]) // n_heads)

reshape_attention_weight(weight, *, component, n_heads, packed_axis=0)

Convert HF linear weights to TransformerLens attention weight shapes.

Source code in src/SafeLens/utils/model_bridge.py

def reshape_attention_weight(
    weight: Any,
    *,
    component: str,
    n_heads: int,
    packed_axis: int | None = 0,
) -> Any:
    """Convert HF linear weights to TransformerLens attention weight shapes."""
    component = _attention_base_component(component)
    shape = getattr(weight, "shape", None)
    reshape = getattr(weight, "reshape", None)
    if shape is None or not callable(reshape) or len(shape) != 2:
        if not _is_sequence(weight):
            raise ValueError(f"Cannot reshape attention weight for component {component!r}.")
        return _reshape_attention_weight_list(
            weight,
            component=component,
            n_heads=n_heads,
            packed_axis=packed_axis,
        )
    axis = packed_axis
    if axis is None:
        axis = infer_attention_weight_packed_axis(weight, n_heads=n_heads)
    if axis not in {0, 1}:
        raise ValueError(f"packed_axis must be 0, 1, or None, got {packed_axis!r}.")
    packed_dim = int(shape[axis])
    other_dim = int(shape[1 - axis])
    if n_heads <= 0 or packed_dim % n_heads != 0:
        raise ValueError(
            f"Cannot split packed dimension {packed_dim} into {n_heads} heads for {component!r}."
        )
    head_dim = packed_dim // n_heads
    if component in {"q", "k", "v"}:
        if axis == 0:
            return weight.reshape(n_heads, head_dim, other_dim).permute(0, 2, 1)
        return weight.reshape(other_dim, n_heads, head_dim).permute(1, 0, 2)
    if component == "z":
        if axis == 0:
            return weight.reshape(n_heads, head_dim, other_dim)
        return weight.reshape(other_dim, n_heads, head_dim).permute(1, 2, 0)
    raise ValueError(f"Unsupported attention weight component {component!r}.")

reshape_joint_qkv_attention_bias(bias, *, component, q_heads, kv_heads, qkv_layout)

Convert a joint QKV projection bias to TransformerLens attention bias shape.

Source code in src/SafeLens/utils/model_bridge.py

def reshape_joint_qkv_attention_bias(
    bias: Any,
    *,
    component: str,
    q_heads: int,
    kv_heads: int,
    qkv_layout: QKVLayout,
) -> Any:
    """Convert a joint QKV projection bias to TransformerLens attention bias shape."""
    component_bias, n_heads = extract_qkv_bias(
        bias,
        component=component,
        q_heads=q_heads,
        kv_heads=kv_heads,
        qkv_layout=qkv_layout,
    )
    return reshape_attention_bias(component_bias, component=component, n_heads=n_heads)

reshape_joint_qkv_attention_weight(weight, *, component, q_heads, kv_heads, qkv_layout, packed_axis=None)

Convert joint QKV projection weights to TransformerLens attention shapes.

Source code in src/SafeLens/utils/model_bridge.py

def reshape_joint_qkv_attention_weight(
    weight: Any,
    *,
    component: str,
    q_heads: int,
    kv_heads: int,
    qkv_layout: QKVLayout,
    packed_axis: int | None = None,
) -> Any:
    """Convert joint QKV projection weights to TransformerLens attention shapes."""
    shape = getattr(weight, "shape", None)
    reshape = getattr(weight, "reshape", None)
    if shape is None or not callable(reshape) or len(shape) != 2:
        if not _is_sequence(weight):
            raise ValueError(f"Cannot reshape joint QKV weight for component {component!r}.")
        return _reshape_joint_qkv_attention_weight_list(
            weight,
            component=component,
            q_heads=q_heads,
            kv_heads=kv_heads,
            qkv_layout=qkv_layout,
            packed_axis=packed_axis,
        )
    if component not in {"q", "k", "v"}:
        raise ValueError(f"Joint QKV weights only expose q/k/v, got {component!r}.")

    axis = packed_axis
    if axis is None:
        axis = infer_qkv_weight_packed_axis(weight, q_heads=q_heads, kv_heads=kv_heads)
    if axis == 0:
        return reshape_joint_qkv_weight_packed_rows(
            weight,
            component=component,
            q_heads=q_heads,
            kv_heads=kv_heads,
            qkv_layout=qkv_layout,
        )
    if axis == 1:
        return reshape_joint_qkv_weight_packed_columns(
            weight,
            component=component,
            q_heads=q_heads,
            kv_heads=kv_heads,
            qkv_layout=qkv_layout,
        )
    raise ValueError(f"packed_axis must be 0, 1, or None, got {packed_axis!r}.")

reshape_joint_qkv_weight_packed_columns(weight, *, component, q_heads, kv_heads, qkv_layout)

Handle Conv1D-style joint weights shaped [d_model, qkv_out].

Source code in src/SafeLens/utils/model_bridge.py

def reshape_joint_qkv_weight_packed_columns(
    weight: Any,
    *,
    component: str,
    q_heads: int,
    kv_heads: int,
    qkv_layout: QKVLayout,
) -> Any:
    """Handle Conv1D-style joint weights shaped `[d_model, qkv_out]`."""
    packed_dim = int(weight.shape[1])
    d_model = int(weight.shape[0])
    component_weight, n_heads = extract_qkv_weight_columns(
        weight,
        component=component,
        q_heads=q_heads,
        kv_heads=kv_heads,
        qkv_layout=qkv_layout,
    )
    head_dim = packed_dim // (q_heads + 2 * kv_heads)
    return component_weight.reshape(d_model, n_heads, head_dim).permute(1, 0, 2)

reshape_joint_qkv_weight_packed_rows(weight, *, component, q_heads, kv_heads, qkv_layout)

Handle linear-style joint weights shaped [qkv_out, d_model].

Source code in src/SafeLens/utils/model_bridge.py

def reshape_joint_qkv_weight_packed_rows(
    weight: Any,
    *,
    component: str,
    q_heads: int,
    kv_heads: int,
    qkv_layout: QKVLayout,
) -> Any:
    """Handle linear-style joint weights shaped `[qkv_out, d_model]`."""
    packed_dim = int(weight.shape[0])
    d_model = int(weight.shape[1])
    component_weight, n_heads = extract_qkv_weight_rows(
        weight,
        component=component,
        q_heads=q_heads,
        kv_heads=kv_heads,
        qkv_layout=qkv_layout,
    )
    head_dim = packed_dim // (q_heads + 2 * kv_heads)
    return component_weight.reshape(n_heads, head_dim, d_model).permute(0, 2, 1)

resolve_module_path(model, path)

Resolve a dotted module path with integer list indexes.

Source code in src/SafeLens/utils/model_bridge.py

def resolve_module_path(model: Any, path: str) -> Any:
    """Resolve a dotted module path with integer list indexes."""
    target = model
    for part in path.split("."):
        if part.isdigit():
            target = target[int(part)]
        else:
            target = getattr(target, part)
    return target

split_heads(activation, n_heads)

Reshape [batch, pos, hidden] activations into [batch, pos, head, head_dim].

Source code in src/SafeLens/utils/model_bridge.py

def split_heads(activation: Any, n_heads: int) -> Any:
    """Reshape `[batch, pos, hidden]` activations into `[batch, pos, head, head_dim]`."""
    shape = getattr(activation, "shape", None)
    reshape = getattr(activation, "reshape", None)
    if shape is None or not callable(reshape):
        if not _is_sequence(activation):
            return activation
        nested_shape = _nested_shape(activation)
        if len(nested_shape) < 3:
            return activation
        hidden_size = nested_shape[-1]
        if n_heads <= 0 or hidden_size % n_heads != 0:
            return activation
        return _split_nested_last_dim(activation, n_heads)
    if len(shape) < 3:
        return activation
    hidden_size = int(shape[-1])
    if n_heads <= 0 or hidden_size % n_heads != 0:
        raise ValueError(
            f"Cannot split activation with final dimension {hidden_size} into {n_heads} heads."
        )
    return activation.reshape(*shape[:-1], n_heads, hidden_size // n_heads)

split_qkv_heads(activation, model, spec)

Extract one component from a joint QKV projection and split it into heads.

Source code in src/SafeLens/utils/model_bridge.py

def split_qkv_heads(activation: Any, model: Any, spec: ComponentHookSpec) -> Any:
    """Extract one component from a joint QKV projection and split it into heads."""
    component = _attention_base_component(spec.component)
    q_heads = head_count_for_component(model, "q")
    kv_heads = head_count_for_component(model, "k")
    if spec.qkv_layout == "interleaved":
        return split_interleaved_qkv_heads(
            activation,
            component,
            q_heads=q_heads,
            kv_heads=kv_heads,
        )
    q_slice, k_slice, v_slice = split_qkv_slices(activation, q_heads=q_heads, kv_heads=kv_heads)
    component_slice = {"q": q_slice, "k": k_slice, "v": v_slice}[component]
    return split_heads(component_slice, q_heads if component == "q" else kv_heads)

sum_attention_heads(value)

Sum the head axis in a TransformerLens result tensor.

Source code in src/SafeLens/utils/model_bridge.py

def sum_attention_heads(value: Any) -> Any:
    """Sum the head axis in a TransformerLens `result` tensor."""
    try:
        import torch

        if hasattr(value, "shape") and isinstance(value, torch.Tensor):
            return value.sum(dim=-2)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(value, "shape"):
            return np.asarray(value).sum(axis=-2)
    except Exception:
        pass
    shape = getattr(value, "shape", None)
    if shape is not None:
        value = value.tolist()
    return _sum_nested_head_axis(value)

supported_transformer_component_names(*, include_pattern=False, include_attention=False, include_result=True)

Return the canonical component vocabulary exposed by model bridges.

Source code in src/SafeLens/utils/model_bridge.py

def supported_transformer_component_names(
    *,
    include_pattern: bool = False,
    include_attention: bool = False,
    include_result: bool = True,
) -> tuple[str, ...]:
    """Return the canonical component vocabulary exposed by model bridges."""
    supported_components = tuple(
        component
        for component in CANONICAL_TRANSFORMER_COMPONENTS
        if include_result or not _is_attention_result_component(component)
    )
    if include_attention:
        return supported_components
    if include_pattern:
        return tuple(
            component
            for component in supported_components
            if not _is_attention_result_component(component)
            and not _is_attention_scores_component(component)
        )
    return tuple(
        component
        for component in supported_components
        if not _is_attention_result_component(component)
        and not _is_attention_pattern_component(component)
        and not _is_attention_scores_component(component)
    )

transform_component_activation(activation, spec, model, *, module=None, component_ref=None, architecture=None)

Convert raw HF projection tensors into SafeLens component activation shape.

Source code in src/SafeLens/utils/model_bridge.py

def transform_component_activation(
    activation: Any,
    spec: ComponentHookSpec,
    model: Any,
    *,
    module: Any | None = None,
    component_ref: ComponentRef | None = None,
    architecture: str | None = None,
) -> Any:
    """Convert raw HF projection tensors into SafeLens component activation shape."""
    if _is_attention_result_component(spec.component):
        return compute_attention_result_activation(
            activation,
            model,
            spec,
            module=module,
            component_ref=component_ref,
            architecture=architecture,
        )
    if spec.activation == "split_heads":
        return split_heads(
            activation,
            head_count_for_component(model, _attention_base_component(spec.component)),
        )
    if spec.activation == "split_qkv_heads":
        return split_qkv_heads(activation, model, spec)
    if spec.activation == "repeat_heads":
        return repeat_along_head_dimension(
            activation,
            head_count_for_component(model, _attention_base_component(spec.component)),
        )
    if spec.component.endswith("_normalized"):
        return normalized_output_from_scale(module or model, activation)
    return activation

transformer_lens_component_name(component, layer)

Return a TransformerLens-style hook name for a canonical component.

Source code in src/SafeLens/utils/model_bridge.py

def transformer_lens_component_name(component: str, layer: int) -> str:
    """Return a TransformerLens-style hook name for a canonical component."""
    if component in {"q_input", "k_input", "v_input", "attn_in"}:
        return f"blocks.{layer}.hook_{component}"
    if component in {"decoder_q_input", "decoder_k_input", "decoder_v_input", "decoder_attn_in"}:
        return f"decoder.{layer}.hook_{component.removeprefix('decoder_')}"
    prefixed_attention = _prefixed_attention_component(component)
    if prefixed_attention is not None:
        prefix, base_component = prefixed_attention
        layer_type = "attn" if prefix == "decoder" else "cross_attn"
        return f"decoder.{layer}.{layer_type}.hook_{base_component}"
    if component in _ATTENTION_HOOK_COMPONENTS:
        hook_component = "attn_scores" if component == "attn_scores" else component
        return f"blocks.{layer}.attn.hook_{hook_component}"
    if component in {"pre", "pre_linear", "post"}:
        return f"blocks.{layer}.mlp.hook_{component}"
    if component in {"decoder_pre", "decoder_pre_linear", "decoder_post"}:
        return f"decoder.{layer}.mlp.hook_{component.removeprefix('decoder_')}"
    if component in {"ln1_scale", "ln1_normalized"}:
        hook_component = "scale" if component == "ln1_scale" else "normalized"
        return f"blocks.{layer}.ln1.hook_{hook_component}"
    if component in {"ln2_scale", "ln2_normalized"}:
        hook_component = "scale" if component == "ln2_scale" else "normalized"
        return f"blocks.{layer}.ln2.hook_{hook_component}"
    if component in {"decoder_ln1_scale", "decoder_ln1_normalized"}:
        hook_component = "scale" if component == "decoder_ln1_scale" else "normalized"
        return f"decoder.{layer}.ln1.hook_{hook_component}"
    if component in {"decoder_ln2_scale", "decoder_ln2_normalized"}:
        hook_component = "scale" if component == "decoder_ln2_scale" else "normalized"
        return f"decoder.{layer}.ln2.hook_{hook_component}"
    if component in {"decoder_ln3_scale", "decoder_ln3_normalized"}:
        hook_component = "scale" if component == "decoder_ln3_scale" else "normalized"
        return f"decoder.{layer}.ln3.hook_{hook_component}"
    if component.startswith("decoder_"):
        return f"decoder.{layer}.hook_{component.removeprefix('decoder_')}"
    if component in {"cross_attn_in", "cross_attn_out"}:
        return f"decoder.{layer}.hook_{component}"
    if component.startswith("ssm_"):
        return f"blocks.{layer}.ssm.hook_{component.removeprefix('ssm_')}"
    return f"blocks.{layer}.hook_{component}"

transpose_2d_weight(weight)

Return a rank-2 weight transposed without requiring tensor dependencies.

Source code in src/SafeLens/utils/model_bridge.py

def transpose_2d_weight(weight: Any) -> Any:
    """Return a rank-2 weight transposed without requiring tensor dependencies."""
    if hasattr(weight, "T") and getattr(weight, "ndim", 0) == 2:
        return weight.T
    if _is_sequence(weight):
        shape = _nested_shape(weight)
        if len(shape) != 2:
            return weight
        return [list(column) for column in zip(*weight, strict=True)]
    return weight

zeros_for_attention_bias(model, component)

Return a zero attention bias with the TransformerLens shape for one component.

Source code in src/SafeLens/utils/model_bridge.py

def zeros_for_attention_bias(model: Any, component: str) -> Any:
    """Return a zero attention bias with the TransformerLens shape for one component."""
    component = _attention_base_component(component)
    n_heads = head_count_for_component(model, component) if component != "z" else None
    d_model = _model_hidden_size(model)
    if d_model is None:
        raise ValueError(f"Could not infer hidden size for b_{component.upper()}.")
    if component == "z":
        return _zeros_vector(d_model, model)
    d_head = attention_head_dim(model)
    if n_heads is None or n_heads <= 0 or d_head is None or d_head <= 0:
        raise ValueError(
            f"Could not infer head dimension for b_{component.upper()} with "
            f"d_model={d_model}, n_heads={n_heads}, d_head={d_head}."
        )
    return _zeros_matrix(n_heads, d_head, model)