Model Wrapper

Model wrappers isolate the rest of the pipeline from model loading details.

Available wrappers:

• DummyModelWrapper: test and CI wrapper with no external dependencies.
• HuggingFaceModelWrapper: loads models directly with Transformers.
• LocalModelWrapper: loads a local Transformers-compatible model directory.
• Qwen3DenseModelWrapper: adapts Qwen3 dense models up to 35B with SafeLens component hooks for residual streams, MLP output, attention output, and q/k/v/z head vectors.
• TransformerLensCompatibleModelWrapper: independent Transformers-based adapter for model families mirrored from the TransformerLens official support table. It uses SafeLens architecture adapters for component hooks and does not import or require TransformerLens.
• ModelScopeModelWrapper: downloads a ModelScope snapshot, then loads it with Transformers.

Select a wrapper through model.source in the YAML config:

model:
  source: modelscope
  name: Qwen/Qwen2.5-0.5B-Instruct

Adapters also declare static capabilities through ModelAdapterRegistry.

safelens models list-supported --json

Static inspection does not load model weights:

safelens inspect-model --model Qwen/Qwen3-8B --json
safelens inspect-model --model gpt2 --json
safelens models list-transformerlens --json

Minimal Python usage:

from SafeLens.core.base import ModelLoadConfig
from SafeLens.utils import build_model_wrapper

model = build_model_wrapper(ModelLoadConfig(source="dummy", name="dummy"))
model.load_model()
output, cache = model.run_with_cache({"text": "hello"}, layers=[0])

Transformers-backed wrappers additionally accept TransformerLens-style cache options: names_filter expands over supported hook names, pos_slice slices the position axis before storage, detach/clone/device prepare cached activations for memory-friendly analysis, remove_batch_dim removes a singleton batch axis from cached activations, and return_cache_object=True returns an ActivationCache. When called with token ids or text directly, run_with_cache(tokens) follows TransformerLens and caches all architecture-bridge component hooks by default, returning an ActivationCache. Pre-softmax attn_scores are intentionally not enabled by the default cache because many Transformers models use SDPA/flash paths rather than a Python torch.softmax; request them explicitly with names_filter when an eager attention path is available. Mapping inputs keep the older SafeLens default of no implicit cache hooks; pass cache_all=True to opt into the same component cache. TransformerLens-style attention result names are supported when the architecture exposes z and W_O; patching them writes the summed per-head result delta back to the merged attention output.

output, cache = model.run_with_cache(
    {"text": "hello"},
    names_filter=lambda name: name.endswith("hook_resid_post"),
    pos_slice=-1,
    device="cpu",
    return_cache_object=True,
    remove_batch_dim=True,
)

tokens = model.to_tokens("hello")
logits, cache = model.run_with_cache(tokens)
resid_pre = cache["resid_pre", 0]
logits, cache = model.run_with_cache(tokens, names_filter="blocks.0.hook_resid_post")
loss, cache = model.run_with_cache(tokens, "loss")
raw_output, cache = model.run_with_cache(tokens, return_type="model_output")
raw_output, full_cache = model.run_with_cache({"input_ids": tokens}, cache_all=True)

Prefer layers=... for SafeLens layer selection; the extra positional arguments on run_with_cache() and run_with_hooks() follow TransformerLens forward ordering, beginning with return_type and loss_per_token.

The same wrappers can install persistent cache hooks, matching TransformerLens' cache_all() and cache_some() workflows. These hooks are permanent by default, so a plain reset_hooks() keeps them active while clearing temporary hooks; use reset_hooks(including_permanent=True) or remove_hooks() to stop caching.

cache = model.cache_some(lambda name: name.endswith("hook_resid_post"))
model(tokens)
resid_post = cache["resid_post", 0]
model.reset_hooks(including_permanent=True)

Transformers-backed wrappers support a TransformerLens-style temporary hook entrypoint. run_with_hooks() and direct calls such as model(tokens) do not install default cache hooks; use run_with_cache() when activations should be stored. Persistent hooks registered with add_hook() can be cleared with either reset_hooks() or remove_hooks(). Like TransformerLens, add_perma_hook() registers hooks that survive a default reset_hooks() call; pass including_permanent=True or call remove_hooks() to clear them. The hooks() context manager installs temporary hooks around arbitrary wrapper calls.

logits = model.run_with_hooks(
    tokens,
    fwd_hooks=[("blocks.0.hook_resid_post", my_hook)],
)
raw_output = model.run_with_hooks(tokens, fwd_hooks=[], return_type="model_output")
logits = model(tokens)
model.add_perma_hook("blocks.0.hook_resid_post", debug_hook)
model.reset_hooks()
model.reset_hooks(including_permanent=True)

Transformers-backed wrappers also expose common TransformerLens-style text and attribution helpers when a tokenizer or output embedding is available:

tokens = model.to_tokens("hello", prepend_bos=False)
batch_tokens = model.to_tokens(["short", "longer"])
text = model.to_string(tokens)
texts = model.to_string([[1, 2], [3, 4]])
pieces = model.to_str_tokens("hello")
token_id = model.to_single_token(" hello")
token_text = model.to_single_str_token(token_id)
position = model.get_token_position(" hello", "well hello there")
directions = model.tokens_to_residual_directions(tokens)
single_direction = model.tokens_to_residual_directions(" hello")

The same wrappers expose a lightweight TransformerLens-style cfg view for common analysis loops:

cfg = model.cfg
for layer in range(cfg.n_layers):
    for head in range(cfg.n_heads):
        ...

For supported Transformers layouts, the wrapper exposes TransformerLens-shaped weight matrices. Split and joint-QKV attention projections are normalized into the same shapes, including GPT-2-style Conv1D packed columns and GPT-NeoX/Pythia-style packed rows:

W_E = model.W_E      # [vocab, d_model]
W_U = model.W_U      # [d_model, vocab]
W_pos = model.W_pos  # [pos, d_model], when the architecture has learned positions
W_Q = model.W_Q  # [layer, head, d_model, d_head]
W_K = model.W_K
W_V = model.W_V
W_O = model.W_O  # [layer, head, d_head, d_model]
W_in = model.W_in    # [layer, d_model, d_mlp]
W_out = model.W_out  # [layer, d_mlp, d_model]

Model wrapper implementations and hook helpers.

DummyModelWrapper

Bases: ModelWrapper

Small in-memory model wrapper used by tests and architecture demos.

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

class DummyModelWrapper(ModelWrapper):
    """Small in-memory model wrapper used by tests and architecture demos."""

    def __init__(self, name: str = "dummy") -> None:
        self.name = name
        self.loaded = False
        self._hooks: list[tuple[LayerRef, HookFn]] = []

    def load_model(self) -> DummyModelWrapper:
        self.loaded = True
        return self

    def add_hook(
        self,
        layer: LayerRef,
        hook_fn: HookFn | None = None,
        *,
        hook: HookFn | None = None,
        dir: str = "fwd",
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
    ) -> _RemovableHandle:
        _ = dir, is_permanent, level, prepend
        resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
        item = (layer, resolved_hook)
        self._hooks.append(item)
        return _RemovableHandle(lambda: self._remove_hook(item))

    def run_with_cache(
        self,
        batch: Batch,
        layers: Sequence[LayerRef] | None = None,
        *,
        names_filter: NamesFilter = None,
        return_cache_object: bool = False,
        remove_batch_dim: bool = False,
        **_kwargs: Any,
    ) -> tuple[dict[str, Any], dict[str, Any] | ActivationCache]:
        if not self.loaded:
            self.load_model()

        candidate_layers = list(layers or [layer for layer, _ in self._hooks])
        selected_layers = _filter_hook_names(
            [activation_name_for_layer(layer) for layer in candidate_layers],
            names_filter,
        )
        cache = {name: {"batch": dict(batch)} for name in selected_layers}
        model_output = {
            "text": batch.get("text") or batch.get("prompt") or "",
            "risk_score": float(batch.get("risk_score", 0.0)),
        }

        for layer, hook_fn in list(self._hooks):
            name = activation_name_for_layer(layer)
            if name in selected_layers:
                activation = cache.get(name, {"batch": dict(batch)})
                patched = _call_dummy_hook(
                    hook_fn,
                    layer=layer,
                    batch=batch,
                    cache=cache,
                    activation=activation,
                )
                if patched is not None:
                    cache[name] = patched

        return model_output, _format_cache_result(
            cache,
            model=self,
            return_cache_object=return_cache_object,
            remove_batch_dim=remove_batch_dim,
        )

    def generate(self, prompt: str | Sequence[str], **generation_kwargs: Any) -> str | list[str]:
        _ = generation_kwargs
        return f"{prompt} [dummy generation]"

    def generate_stream(
        self,
        prompt: str | Sequence[str],
        *,
        max_new_tokens: int = 10,
        max_tokens_per_yield: int = 25,
        **generation_kwargs: Any,
    ) -> Iterable[str | list[str]]:
        generated = self.generate(prompt, max_new_tokens=max_new_tokens, **generation_kwargs)
        if isinstance(generated, list):
            yield generated
        else:
            for start in range(0, len(generated), max_tokens_per_yield):
                yield generated[start : start + max_tokens_per_yield]

    def remove_hooks(self) -> None:
        self._hooks.clear()

    def _remove_hook(self, item: tuple[LayerRef, HookFn]) -> None:
        if item in self._hooks:
            self._hooks.remove(item)

HuggingFaceModelWrapper

Bases: ModelWrapper

Transformers-based model wrapper with forward hook support.

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

class HuggingFaceModelWrapper(ModelWrapper):
    """Transformers-based model wrapper with forward hook support."""

    def __init__(
        self,
        name: str,
        dtype: str = "float32",
        device: str | None = None,
        revision: str | None = None,
        cache_dir: str | None = None,
        trust_remote_code: bool = False,
        load_kwargs: dict[str, Any] | None = None,
        tokenizer_kwargs: dict[str, Any] | None = None,
        pretrained_path: str | None = None,
        pretrained_path_is_local: bool = False,
    ) -> None:
        self.name = name
        self.dtype = dtype
        self.device = device
        self.revision = revision
        self.cache_dir = cache_dir
        self.trust_remote_code = trust_remote_code
        self.load_kwargs = load_kwargs or {}
        self.tokenizer_kwargs = tokenizer_kwargs or {}
        self.pretrained_path = pretrained_path
        self._pretrained_path_is_local = bool(pretrained_path_is_local)
        self.model: Any = None
        self.tokenizer: Any = None
        self._tokenizer_load_error: Exception | None = None
        self._hooks: list[Any] = []
        self.is_caching = False
        self._temporary_backward_hook_depth = 0
        self._attention_hook_count = 0
        self._run_requires_output_attentions = False
        self.default_prepend_bos: bool | None = None
        self._transformer_lens_runtime_flags = dict(_TRANSFORMER_LENS_RUNTIME_FLAG_DEFAULTS)

    @property
    def cfg(self) -> TransformerLensConfigView:
        """Return a TransformerLens-style normalized config view."""
        return _make_transformer_lens_config_view(
            self._require_model(),
            model_name=self.name,
            device=self.device,
            dtype=self.dtype,
            tokenizer=self.tokenizer,
            runtime_flags=self._transformer_lens_runtime_flags,
        )

    def set_use_attn_result(self, use_attn_result: bool) -> None:
        """Mirror TransformerLens' runtime switch for per-head attention results."""
        self._set_transformer_lens_runtime_flag("use_attn_result", use_attn_result)

    def set_use_split_qkv_input(self, use_split_qkv_input: bool) -> None:
        """Mirror TransformerLens' runtime switch for split Q/K/V input hooks."""
        self._set_transformer_lens_runtime_flag("use_split_qkv_input", use_split_qkv_input)

    def set_use_hook_mlp_in(self, use_hook_mlp_in: bool) -> None:
        """Mirror TransformerLens' runtime switch for MLP input hooks."""
        if bool(use_hook_mlp_in) and self.cfg.attn_only:
            raise AssertionError("Cannot use hook_mlp_in on attention-only models.")
        self._set_transformer_lens_runtime_flag("use_hook_mlp_in", use_hook_mlp_in)

    def set_use_attn_in(self, use_attn_in: bool) -> None:
        """Mirror TransformerLens' runtime switch for attention input hooks."""
        if bool(use_attn_in) and self._has_grouped_query_attention():
            raise AssertionError(
                "Cannot use attn_in hooks when key/value heads are grouped; "
                "SafeLens exposes the config switch but does not synthesize missing HF hooks."
            )
        self._set_transformer_lens_runtime_flag("use_attn_in", use_attn_in)

    def set_ungroup_grouped_query_attention(self, ungroup_grouped_query_attention: bool) -> None:
        """Mirror TransformerLens' grouped-query attention compatibility switch."""
        self._set_transformer_lens_runtime_flag(
            "ungroup_grouped_query_attention",
            ungroup_grouped_query_attention,
        )

    def _set_transformer_lens_runtime_flag(self, name: str, value: bool) -> None:
        if name not in _TRANSFORMER_LENS_RUNTIME_FLAG_DEFAULTS:
            raise KeyError(f"Unknown TransformerLens runtime flag {name!r}.")
        self._transformer_lens_runtime_flags[name] = bool(value)

    def _has_grouped_query_attention(self) -> bool:
        cfg = self.cfg
        return (
            cfg.n_heads is not None
            and cfg.n_key_value_heads is not None
            and cfg.n_key_value_heads < cfg.n_heads
        )

    def _uses_decoder_text_input_semantics(self) -> bool:
        return transformer_lens_model_kind(self.name) == "decoder"

    def _transformer_lens_model_kind(
        self,
        *,
        pretrained_path: str | None = None,
        probe_pretrained_path: bool = False,
    ) -> str:
        _ = probe_pretrained_path
        model = getattr(self, "model", None)
        config = _core_model_config(_config_attr(model, "config"))
        model_type = _config_attr(config, "model_type")
        if isinstance(model_type, str) and model_type:
            return transformer_lens_model_kind(model_type)
        candidate = pretrained_path or self.pretrained_path or self.name
        return transformer_lens_model_kind(candidate)

    def check_hooks_to_add(
        self,
        hook_point: Any,
        hook_point_name: str,
        hook: Any,
        dir: str = "fwd",
        is_permanent: bool = False,
        prepend: bool = False,
    ) -> None:
        """Validate TransformerLens runtime-hook flags before adding hooks."""
        _ = hook_point, hook, dir, is_permanent, prepend
        if hook_point_name.endswith("attn.hook_result") or hook_point_name.endswith("hook_result"):
            assert (
                self.cfg.use_attn_result
            ), f"Cannot add hook {hook_point_name} if use_attn_result_hook is False"
        if hook_point_name.endswith(("hook_q_input", "hook_k_input", "hook_v_input")):
            assert (
                self.cfg.use_split_qkv_input
            ), f"Cannot add hook {hook_point_name} if use_split_qkv_input is False"
        if hook_point_name.endswith("mlp_in"):
            assert (
                self.cfg.use_hook_mlp_in
            ), f"Cannot add hook {hook_point_name} if use_hook_mlp_in is False"
        if hook_point_name.endswith("attn_in"):
            assert (
                self.cfg.use_attn_in
            ), f"Cannot add hook {hook_point_name} if use_attn_in is False"

    def get_pos_offset(self, past_kv_cache: Any, batch_size: int) -> int:
        """Return the positional offset implied by a TransformerLens KV cache."""
        if past_kv_cache is None:
            return 0
        cached_batch_size = _past_kv_cache_batch_size(past_kv_cache)
        if cached_batch_size is not None:
            assert int(cached_batch_size) == int(batch_size)
        return _past_kv_cache_length(past_kv_cache)

    def get_residual(
        self,
        embed: Any,
        pos_offset: int,
        prepend_bos: bool | None = None,
        attention_mask: Any | None = None,
        tokens: Any | None = None,
        return_shortformer_pos_embed: bool = True,
        device: Any = None,
    ) -> Any:
        """Convert token embeddings into the first residual stream."""
        model = self._require_model()
        resolved_device = self.device if device is None else device
        if tokens is None:
            tokens = _ones_token_batch_like_embedding(embed, device=resolved_device)
        else:
            tokens = _coerce_token_model_input(
                _ensure_token_batch_dim(tokens), device=resolved_device
            )

        position_type = _infer_positional_embedding_type(model)
        if position_type == "standard":
            pos_module = _positional_embeddings_module(model)
            if pos_module is None:
                residual = embed
                shortformer_pos_embed = None
            else:
                pos_embed = _call_position_embedding_module(
                    pos_module,
                    tokens,
                    pos_offset=pos_offset,
                    attention_mask=attention_mask,
                    device=resolved_device,
                )
                residual = _add_tensor_like_values(embed, pos_embed)
                shortformer_pos_embed = None
        elif position_type == "shortformer":
            pos_module = _positional_embeddings_module(model)
            if pos_module is None:
                raise NotImplementedError(
                    "SafeLens cannot compute shortformer positional embeddings without "
                    "a resolvable positional embedding module."
                )
            shortformer_pos_embed = _call_position_embedding_module(
                pos_module,
                tokens,
                pos_offset=pos_offset,
                attention_mask=attention_mask,
                device=resolved_device,
            )
            residual = embed
        elif position_type in {"rotary", "alibi", None}:
            residual = embed
            shortformer_pos_embed = None
        else:
            raise ValueError(f"Invalid positional_embedding_type {position_type!r}.")

        if return_shortformer_pos_embed:
            return residual, shortformer_pos_embed
        return residual

    def input_to_embed(
        self,
        input: Any,
        prepend_bos: bool | None = None,
        padding_side: str | None = None,
        truncate: bool | None = None,
        attention_mask: Any | None = None,
        past_kv_cache: Any | None = None,
    ) -> tuple[Any, Any, Any | None, Any | None]:
        """Convert token/text input to TL-style `(residual, tokens, pos_embed, mask)`."""
        model = self._require_model()
        if isinstance(input, Mapping):
            if prepend_bos is None and "prepend_bos" in input:
                prepend_bos = input["prepend_bos"]
            if padding_side is None and "padding_side" in input:
                padding_side = input["padding_side"]
            if truncate is None and "truncate" in input:
                truncate = input["truncate"]
            if attention_mask is None and "attention_mask" in input:
                attention_mask = input["attention_mask"]
            if "input_ids" in input:
                input = input["input_ids"]
            elif "tokens" in input:
                input = input["tokens"]
            elif "token_ids" in input:
                input = input["token_ids"]
            else:
                text = _text_or_prompt_value(input)
                if text is not None:
                    input = text
        if isinstance(input, str) or _is_text_batch(input):
            resolved_truncate = True if truncate is None else bool(truncate)
            text_input = cast(str | Sequence[str], input)
            tokens = self.to_tokens(
                text_input,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
                truncate=resolved_truncate,
            )
        else:
            tokens = input
        tokens = _coerce_token_model_input(_ensure_token_batch_dim(tokens), device=self.device)

        if attention_mask is None:
            effective_padding_side = str(
                padding_side or getattr(self.tokenizer, "padding_side", "right")
            )
            needs_mask = effective_padding_side == "left" or past_kv_cache is not None
            if needs_mask and self.tokenizer is not None:
                resolved_prepend_bos = _resolve_default_prepend_bos(self, prepend_bos)
                attention_mask = _attention_mask_from_tokens(
                    tokens,
                    _tokenizer_effective_pad_token_id(self.tokenizer),
                    prepend_bos=resolved_prepend_bos,
                    padding_side=effective_padding_side,
                    bos_token_id=getattr(self.tokenizer, "bos_token_id", None),
                )
        if attention_mask is not None:
            attention_mask = _coerce_token_model_input(
                _ensure_token_batch_dim(attention_mask),
                device=self.device,
            )
            if _shape_of_token_ids(attention_mask) != _shape_of_token_ids(tokens):
                raise AssertionError(
                    f"Attention mask shape {_shape_of_token_ids(attention_mask)!r} "
                    f"does not match tokens shape {_shape_of_token_ids(tokens)!r}."
                )
            append_attention_mask = getattr(past_kv_cache, "append_attention_mask", None)
            if callable(append_attention_mask):
                attention_mask = append_attention_mask(attention_mask)
            elif past_kv_cache is not None:
                attention_mask = _extend_attention_mask_for_past_cache(
                    attention_mask, past_kv_cache
                )

        pos_offset = self.get_pos_offset(past_kv_cache, _token_batch_size(tokens))
        embed_module = _input_embeddings_module(model)
        if embed_module is None:
            raise RuntimeError("Could not resolve an input embedding module.")
        embed = embed_module(tokens)
        residual, shortformer_pos_embed = self.get_residual(
            embed,
            pos_offset,
            prepend_bos=prepend_bos,
            attention_mask=attention_mask,
            tokens=tokens,
            return_shortformer_pos_embed=True,
            device=self.device,
        )
        return residual, tokens, shortformer_pos_embed, attention_mask

    def set_tokenizer(self, tokenizer: Any, default_padding_side: str | None = None) -> None:
        """Set the tokenizer used by TransformerLens-style token helpers."""
        if default_padding_side not in {"right", "left", None}:
            raise AssertionError(
                f"padding_side must be 'right', 'left' or None, got {default_padding_side!r}"
            )
        self.tokenizer = tokenizer
        if default_padding_side is not None:
            _set_attr_if_possible(tokenizer, "padding_side", default_padding_side)
        elif getattr(tokenizer, "padding_side", None) is None:
            _set_attr_if_possible(tokenizer, "padding_side", "right")

        eos_token = getattr(tokenizer, "eos_token", None)
        if eos_token is None:
            eos_token = "<|endoftext|>"
            _set_attr_if_possible(tokenizer, "eos_token", eos_token)
        if getattr(tokenizer, "pad_token", None) is None:
            _set_attr_if_possible(tokenizer, "pad_token", eos_token)
        if getattr(tokenizer, "bos_token", None) is None:
            _set_attr_if_possible(tokenizer, "bos_token", eos_token)
        eos_token_id = getattr(tokenizer, "eos_token_id", None)
        if getattr(tokenizer, "pad_token_id", None) is None and eos_token_id is not None:
            _set_attr_if_possible(tokenizer, "pad_token_id", eos_token_id)
        if getattr(tokenizer, "bos_token_id", None) is None and eos_token_id is not None:
            _set_attr_if_possible(tokenizer, "bos_token_id", eos_token_id)
        self.tokenizer_prepends_bos = _tokenizer_prepends_bos(tokenizer)

    def to(self, device_or_dtype: Any, print_details: bool = True) -> HuggingFaceModelWrapper:
        """Move the underlying model or change dtype, returning ``self`` like PyTorch/TL."""
        _ = print_details
        target = _coerce_to_target(device_or_dtype)
        to_fn = getattr(self.model, "to", None)
        if callable(to_fn):
            to_fn(target)
        _update_wrapper_device_dtype(self, device_or_dtype, target)
        return self

    def cuda(self, device: int | Any | None = None) -> HuggingFaceModelWrapper:
        if isinstance(device, int):
            return self.to(f"cuda:{device}")
        if device is None:
            return self.to("cuda")
        return self.to(device)

    def cpu(self) -> HuggingFaceModelWrapper:
        return self.to("cpu")

    def mps(self) -> HuggingFaceModelWrapper:
        return self.to("mps")

    def move_model_modules_to_device(self) -> HuggingFaceModelWrapper:
        """Compatibility no-op unless a wrapper device has been set."""
        if self.device is not None:
            return self.to(self.device)
        return self

    def init_weights(self) -> None:
        """Unsupported HookedTransformer weight-initialization helper."""
        raise NotImplementedError(
            "SafeLens wraps pretrained Transformers modules and does not initialize "
            "HookedTransformer-format weights."
        )

    def load_and_process_state_dict(self, *args: Any, **kwargs: Any) -> None:
        """Unsupported HookedTransformer weight-processing helper."""
        _ = args, kwargs
        raise NotImplementedError(
            "SafeLens' dependency-free wrapper does not load or process "
            "HookedTransformer-format state dictionaries."
        )

    def fill_missing_keys(self, *args: Any, **kwargs: Any) -> None:
        """Unsupported HookedTransformer state-dict helper."""
        _ = args, kwargs
        raise NotImplementedError(
            "SafeLens' dependency-free wrapper does not mutate HookedTransformer-format "
            "state dictionaries."
        )

    def fold_layer_norm(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
        """Fold LayerNorm/RMSNorm affine parameters into reader weights in place."""
        _ = args
        fold_biases_override = kwargs.pop("fold_biases", None)
        center_weights_override = kwargs.pop("center_weights", None)
        if kwargs:
            unexpected = ", ".join(sorted(kwargs))
            raise TypeError(f"Unexpected fold_layer_norm keyword argument(s): {unexpected}.")
        model = self._require_model()
        normalization_type = self.cfg.normalization_type or "LN"
        is_rms_norm = normalization_type in {"RMS", "RMSPre"}
        fold_biases = (
            (not is_rms_norm) if fold_biases_override is None else bool(fold_biases_override)
        )
        center_weights = (
            (not is_rms_norm) if center_weights_override is None else bool(center_weights_override)
        )
        _fold_layer_norm_weights_in_model(
            model,
            model_name=self.name,
            fold_biases=fold_biases,
            center_weights=center_weights,
            rmsnorm_uses_offset=self.cfg.rmsnorm_uses_offset,
        )
        return self

    def center_writing_weights(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
        """Center weights and biases that write directly to the residual stream."""
        _ = args, kwargs
        model = self._require_model()
        if _is_olmo2_post_norm_model(model):
            return self
        config = getattr(model, "config", None)
        _center_module_weight(_input_embeddings_module(model), axis=-1)
        try:
            positional_embeddings = _positional_embeddings_module(model)
        except KeyError:
            positional_embeddings = None
        if positional_embeddings is not None:
            _center_module_weight(positional_embeddings, axis=-1)

        native_w_o = getattr(model, "W_O", None)
        if native_w_o is not None:
            model.W_O = _center_residual_stream_weight(native_w_o, d_model=self.cfg.d_model)
        native_b_o = getattr(model, "b_O", None)
        if native_b_o is not None:
            model.b_O = _center_residual_stream_bias(native_b_o)
        native_w_out = getattr(model, "W_out", None)
        if native_w_out is not None:
            model.W_out = _center_residual_stream_weight(native_w_out, d_model=self.cfg.d_model)
        native_b_out = getattr(model, "b_out", None)
        if native_b_out is not None:
            model.b_out = _center_residual_stream_bias(native_b_out)

        adapter = architecture_adapter_for_model(model, model_name=self.name)
        for layer in range(_infer_model_layers(model)):
            try:
                attention_ref = adapter.parse_component_ref(
                    transformer_lens_component_name("z", layer)
                )
                if attention_ref is not None:
                    attention_module = adapter.get_component(model, attention_ref)
                    _center_attention_output_module(
                        attention_module,
                        architecture=adapter.name,
                    )
            except (KeyError, NotImplementedError, ValueError):
                pass
            if _config_attr(config, "attn_only", False):
                continue
            try:
                mlp_ref = adapter.parse_component_ref(
                    transformer_lens_component_name("post", layer)
                )
                if mlp_ref is not None:
                    mlp_module = adapter.get_component(model, mlp_ref)
                    _center_mlp_output_module(mlp_module, d_model=self.cfg.d_model)
            except (KeyError, NotImplementedError, ValueError):
                pass
        return self

    def center_unembed(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
        """Center unembedding directions, matching TransformerLens weight processing."""
        _ = args, kwargs
        model = self._require_model()
        native_weight = getattr(model, "W_U", None)
        if native_weight is not None:
            centered = _center_unembed_weight(native_weight, weight_layout="d_model_vocab")
            try:
                model.W_U = centered
            except Exception as exc:
                raise RuntimeError("Could not update native TransformerLens W_U.") from exc
            native_bias = getattr(model, "b_U", None)
            if native_bias is not None:
                try:
                    model.b_U = _center_bias_like(native_bias)
                except Exception as exc:
                    raise RuntimeError("Could not update native TransformerLens b_U.") from exc
            return self

        embeddings = self._output_embeddings()
        weight = getattr(embeddings, "weight", None)
        if weight is None:
            raise RuntimeError("Could not find unembedding weights to center.")
        centered = _center_unembed_weight(weight, weight_layout="vocab_d_model")
        updated = False
        try:
            embeddings.weight = centered
            updated = True
        except Exception as exc:
            data = getattr(weight, "data", None)
            if data is None:
                raise RuntimeError("Could not update output embedding weight.") from exc
            data.copy_(centered)
            updated = True
        if updated and hasattr(model, "_weight"):
            try:
                model._weight = centered
            except Exception:
                pass
        bias = getattr(embeddings, "bias", None)
        if bias is not None:
            centered_bias = _center_bias_like(bias)
            _set_module_bias(embeddings, centered_bias)
            if hasattr(model, "_bias"):
                try:
                    model._bias = centered_bias
                except Exception:
                    pass
        return self

    def fold_value_biases(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
        """Fold attention value biases into output biases and clear value biases."""
        _ = args, kwargs
        model = self._require_model()
        native_b_v = getattr(model, "b_V", None)
        native_w_o = getattr(model, "W_O", None)
        if native_b_v is not None and native_w_o is not None:
            native_b_o = getattr(model, "b_O", None)
            folded_bias, zero_b_v = _fold_value_bias_tensor_like(
                native_b_v,
                native_w_o,
                native_b_o,
                target_heads=self.cfg.n_heads,
            )
            model.b_O = folded_bias
            model.b_V = zero_b_v

        adapter = architecture_adapter_for_model(model, model_name=self.name)
        for layer in range(_infer_model_layers(model)):
            try:
                value_ref = adapter.parse_component_ref(transformer_lens_component_name("v", layer))
                output_ref = adapter.parse_component_ref(
                    transformer_lens_component_name("z", layer)
                )
                if value_ref is None or output_ref is None:
                    continue
                value_spec = adapter._spec_for_ref(value_ref, for_cache=True)
                value_module = adapter.get_component(model, value_ref)
                output_module = adapter.get_component(model, output_ref)
                _fold_value_bias_modules(
                    model,
                    value_module=value_module,
                    output_module=output_module,
                    value_spec=value_spec,
                    architecture=adapter.name,
                )
            except (KeyError, NotImplementedError, ValueError):
                pass
        return self

    def refactor_factored_attn_matrices(
        self,
        *args: Any,
        **kwargs: Any,
    ) -> HuggingFaceModelWrapper:
        """Refactor native attention QK/OV matrices into SVD-based factorizations."""
        _ = args, kwargs
        model = self._require_model()
        if _uses_rotary_embeddings(model):
            raise AssertionError(
                "You can't refactor the QK circuit when using rotary embeddings "
                "(the QK matrix depends on query/key position)."
            )
        refactored = False
        native_w_q = getattr(model, "W_Q", None)
        native_w_k = getattr(model, "W_K", None)
        if native_w_q is not None and native_w_k is not None:
            native_b_q = getattr(model, "b_Q", None)
            native_b_k = getattr(model, "b_K", None)
            refactored_w_q, refactored_b_q, refactored_w_k, refactored_b_k = _refactor_qk_matrices(
                native_w_q,
                native_b_q,
                native_w_k,
                native_b_k,
            )
            model.W_Q = refactored_w_q
            model.W_K = refactored_w_k
            if native_b_q is not None:
                model.b_Q = refactored_b_q
            if native_b_k is not None:
                model.b_K = refactored_b_k
            refactored = True
        native_w_v = getattr(model, "W_V", None)
        native_w_o = getattr(model, "W_O", None)
        if native_w_v is not None and native_w_o is not None:
            native_b_v = getattr(model, "b_V", None)
            if native_b_v is not None:
                native_b_o = getattr(model, "b_O", None)
                folded_bias, zero_b_v = _fold_value_bias_tensor_like(
                    native_b_v,
                    native_w_o,
                    native_b_o,
                    target_heads=self.cfg.n_heads,
                )
                model.b_O = folded_bias
                model.b_V = zero_b_v
            refactored_w_v, refactored_w_o = _refactor_ov_matrices(native_w_v, native_w_o)
            model.W_V = refactored_w_v
            model.W_O = refactored_w_o
            refactored = True
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        for layer in range(_infer_model_layers(model)):
            if _refactor_split_attention_layer(model, adapter=adapter, layer=layer):
                refactored = True
            elif _refactor_joint_qkv_attention_layer(model, adapter=adapter, layer=layer):
                refactored = True
        if not refactored:
            raise NotImplementedError(
                "refactor_factored_attn_matrices currently requires native TL-shaped "
                "W_Q/W_K or W_V/W_O attributes, HF split q/k/v/o projection modules, "
                "or joint QKV projection modules with matching query/key/value head counts. "
                "Grouped-query refactor writeback for this experimental TL pass is not "
                "implemented yet."
            )
        return self

    def process_weights_(
        self,
        fold_ln: bool = True,
        center_writing_weights: bool = True,
        center_unembed: bool = True,
        fold_value_biases: bool = True,
        refactor_factored_attn_matrices: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> HuggingFaceModelWrapper:
        """Run supported in-place TransformerLens-style weight processing passes."""
        _ = args
        if kwargs:
            unexpected = ", ".join(sorted(kwargs))
            raise TypeError(f"Unexpected process_weights_ keyword argument(s): {unexpected}.")
        model = self._require_model()
        has_shortformer_positional_embeddings = (
            _infer_positional_embedding_type(model) == "shortformer"
        )
        has_olmo2_post_norm = _is_olmo2_post_norm_model(model)
        center_writing_weights_by_default = _can_center_writing_weights_by_default(model)
        should_center_writing_weights = (
            center_writing_weights
            and not has_shortformer_positional_embeddings
            and not has_olmo2_post_norm
            and center_writing_weights_by_default
        )
        if fold_ln and not has_shortformer_positional_embeddings and not has_olmo2_post_norm:
            self.fold_layer_norm()
        if should_center_writing_weights:
            self.center_writing_weights()
        if (
            center_unembed
            and not has_shortformer_positional_embeddings
            and not _has_output_logits_soft_cap(model)
            and self._transformer_lens_model_kind() not in {"encoder", "audio_encoder"}
        ):
            self.center_unembed()
        if fold_value_biases:
            self.fold_value_biases()
            if should_center_writing_weights:
                self._center_attention_output_biases_after_value_folding()
        if refactor_factored_attn_matrices:
            self.refactor_factored_attn_matrices()
        return self

    def _center_attention_output_biases_after_value_folding(self) -> None:
        model = self._require_model()
        native_b_o = getattr(model, "b_O", None)
        if native_b_o is not None:
            model.b_O = _center_residual_stream_bias(native_b_o)

        adapter = architecture_adapter_for_model(model, model_name=self.name)
        for layer in range(_infer_model_layers(model)):
            try:
                attention_ref = adapter.parse_component_ref(
                    transformer_lens_component_name("z", layer)
                )
                if attention_ref is None:
                    continue
                attention_module = adapter.get_component(model, attention_ref)
                _center_module_bias(attention_module)
            except (KeyError, NotImplementedError, ValueError):
                continue

    def load_sample_training_dataset(self, *args: Any, **kwargs: Any) -> None:
        """Store an empty sample dataset placeholder for TL notebook compatibility."""
        _ = args, kwargs
        self.dataset: list[Any] = []

    def sample_datapoint(self, *args: Any, **kwargs: Any) -> Any:
        """Return one datapoint from a previously loaded sample dataset."""
        _ = args, kwargs
        dataset = getattr(self, "dataset", None)
        if not dataset:
            raise ValueError("No sample training dataset is loaded.")
        try:
            import random

            return random.choice(dataset)
        except Exception:
            return dataset[0]

    def parameters(self, *args: Any, **kwargs: Any) -> Any:
        """Proxy parameter iteration to the wrapped model when available."""
        model = self._require_model()
        parameters = getattr(model, "parameters", None)
        if callable(parameters):
            return parameters(*args, **kwargs)
        return iter(())

    def named_parameters(self, *args: Any, **kwargs: Any) -> Any:
        """Proxy named parameter iteration to the wrapped model when available."""
        model = self._require_model()
        named_parameters = getattr(model, "named_parameters", None)
        if callable(named_parameters):
            return named_parameters(*args, **kwargs)
        return iter(())

    @property
    def n_params_total(self) -> int:
        """Return the wrapped model's total parameter count."""
        model = self._require_model()
        num_parameters = getattr(model, "num_parameters", None)
        if callable(num_parameters):
            try:
                return int(cast(Any, num_parameters()))
            except TypeError:
                pass
        return sum(_parameter_numel(parameter) for parameter in self.parameters())

    @property
    def W_U(self) -> Any:
        """Return an unembedding matrix shaped `[d_model, vocab]` when available."""
        weight = self._output_embedding_weight()
        return transpose_2d_weight(weight)

    @property
    def b_U(self) -> Any:
        """Return unembedding bias shaped `[vocab]` when available, else zeros."""
        embeddings = self._output_embeddings()
        bias = getattr(embeddings, "bias", None)
        if bias is not None:
            return bias
        model = self._require_model()
        bias = getattr(model, "b_U", None)
        if bias is not None:
            return bias
        return zeros_like_last_dim(self._output_embedding_weight(), axis=0)

    @property
    def W_E(self) -> Any:
        """Return token embedding weights shaped `[vocab, d_model]`."""
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        return adapter.get_embedding_weight(model)

    @property
    def W_pos(self) -> Any:
        """Return positional embedding weights when available."""
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        return adapter.get_embedding_weight(model, positional=True)

    @property
    def W_E_pos(self) -> Any:
        """Return concatenated token and positional embeddings."""
        return _concat_first_dim(self.W_E, self.W_pos)

    @property
    def W_Q(self) -> Any:
        """Return query weights shaped `[layer, head, d_model, d_head]`."""
        return self._stack_attention_weights("q")

    @property
    def W_K(self) -> Any:
        """Return key weights shaped `[layer, head, d_model, d_head]`."""
        return self._stack_attention_weights("k")

    @property
    def W_V(self) -> Any:
        """Return value weights shaped `[layer, head, d_model, d_head]`."""
        return self._stack_attention_weights("v")

    @property
    def W_O(self) -> Any:
        """Return output weights shaped `[layer, head, d_head, d_model]`."""
        return self._stack_attention_weights("z")

    @property
    def b_Q(self) -> Any:
        """Return query biases shaped `[layer, head, d_head]`."""
        return self._stack_attention_biases("q")

    @property
    def b_K(self) -> Any:
        """Return key biases shaped `[layer, head, d_head]`."""
        return self._stack_attention_biases("k")

    @property
    def b_V(self) -> Any:
        """Return value biases shaped `[layer, head, d_head]`."""
        return self._stack_attention_biases("v")

    @property
    def b_O(self) -> Any:
        """Return attention output biases shaped `[layer, d_model]`."""
        return self._stack_attention_biases("z")

    @property
    def QK(self) -> FactoredMatrix:
        """Return the TransformerLens-style QK circuit as a factored matrix."""
        w_q = self.W_Q
        w_k = _repeat_key_value_heads_to_query_heads(
            self.W_K,
            target_heads=_stacked_attention_head_count(w_q),
            tensor_name="W_K",
        )
        return FactoredMatrix(w_q, transpose(w_k))

    @property
    def OV(self) -> FactoredMatrix:
        """Return the TransformerLens-style OV circuit as a factored matrix."""
        w_o = self.W_O
        w_v = _repeat_key_value_heads_to_query_heads(
            self.W_V,
            target_heads=_stacked_attention_head_count(w_o),
            tensor_name="W_V",
        )
        return FactoredMatrix(w_v, w_o)

    @property
    def W_U_U(self) -> Any:
        """Return the left singular vectors of the unembedding matrix."""
        return _svd_component(self.W_U, "U")

    @property
    def W_U_S(self) -> Any:
        """Return the singular values of the unembedding matrix."""
        return _svd_component(self.W_U, "S")

    @property
    def W_U_V(self) -> Any:
        """Return the right singular vectors of the unembedding matrix."""
        return _svd_component(self.W_U, "V")

    @property
    def W_in(self) -> Any:
        """Return MLP input weights shaped `[layer, d_model, d_mlp]`."""
        return self._stack_mlp_weights("in")

    @property
    def W_gate(self) -> Any:
        """Return gated-MLP gate weights shaped `[layer, d_model, d_mlp]`."""
        return self._stack_mlp_weights("gate")

    @property
    def W_out(self) -> Any:
        """Return MLP output weights shaped `[layer, d_mlp, d_model]`."""
        return self._stack_mlp_weights("out")

    @property
    def b_in(self) -> Any:
        """Return MLP input biases shaped `[layer, d_mlp]`."""
        return self._stack_mlp_biases("in")

    @property
    def b_out(self) -> Any:
        """Return MLP output biases shaped `[layer, d_model]`."""
        return self._stack_mlp_biases("out")

    def tl_parameters(self) -> dict[str, Any]:
        """Return a TransformerLens-style parameter mapping for analysis helpers."""
        parameters: dict[str, Any] = {}
        try:
            parameters["embed.W_E"] = self.W_E
        except KeyError:
            pass
        try:
            parameters["unembed.W_U"] = self.W_U
            parameters["unembed.b_U"] = self.b_U
        except RuntimeError:
            pass
        try:
            parameters["pos_embed.W_pos"] = self.W_pos
        except KeyError:
            pass

        cfg = self.cfg
        n_layers = int(cfg.n_layers or 0)
        stacked: dict[str, Any] = {}
        parameter_getters: tuple[tuple[str, Callable[[], Any]], ...] = (
            ("blocks.{layer}.attn.W_Q", lambda: self.W_Q),
            ("blocks.{layer}.attn.W_K", lambda: self.W_K),
            ("blocks.{layer}.attn.W_V", lambda: self.W_V),
            ("blocks.{layer}.attn.W_O", lambda: self.W_O),
            ("blocks.{layer}.attn.b_Q", lambda: self.b_Q),
            ("blocks.{layer}.attn.b_K", lambda: self.b_K),
            ("blocks.{layer}.attn.b_V", lambda: self.b_V),
            ("blocks.{layer}.attn.b_O", lambda: self.b_O),
        )
        for template, getter in parameter_getters:
            try:
                stacked[template] = getter()
            except (KeyError, RuntimeError, ValueError):
                pass
        try:
            stacked.update(
                {
                    "blocks.{layer}.mlp.W_in": self.W_in,
                    "blocks.{layer}.mlp.W_out": self.W_out,
                    "blocks.{layer}.mlp.b_in": self.b_in,
                    "blocks.{layer}.mlp.b_out": self.b_out,
                }
            )
        except (KeyError, RuntimeError, ValueError):
            pass
        try:
            stacked["blocks.{layer}.mlp.W_gate"] = self.W_gate
        except (KeyError, RuntimeError, ValueError):
            pass
        for template, value in stacked.items():
            for layer_index in range(min(n_layers, _stack_first_dim(value))):
                parameters[template.format(layer=layer_index)] = value[layer_index]
        try:
            parameters.update(self._layer_norm_parameters())
        except (KeyError, RuntimeError, ValueError, NotImplementedError):
            pass
        return parameters

    def _layer_norm_parameters(self) -> dict[str, Any]:
        """Return TransformerLens-style per-layer norm affine weights when present."""
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        n_layers = int(self.cfg.n_layers or 0)
        norm_parameters: dict[str, Any] = {}
        for component, norm_name in (("ln1_scale", "ln1"), ("ln2_scale", "ln2")):
            for layer_index in range(n_layers):
                try:
                    norm_module = _norm_module_for_layer(
                        model,
                        adapter=adapter,
                        layer=layer_index,
                        component=component,
                    )
                except (KeyError, NotImplementedError, ValueError, AttributeError, IndexError):
                    continue
                weight = _norm_affine_weight(
                    norm_module,
                    rmsnorm_uses_offset=self.cfg.rmsnorm_uses_offset,
                )
                if weight is not None:
                    norm_parameters[f"blocks.{layer_index}.{norm_name}.w"] = weight
        return norm_parameters

    def accumulated_bias(
        self,
        layer: int,
        mlp_input: bool = False,
        include_mlp_biases: bool = True,
    ) -> Any:
        """Return accumulated attention/MLP output biases before a layer."""
        b_o = self.b_O
        n_layers = _stack_first_dim(b_o)
        if layer < 0 or layer > n_layers:
            raise ValueError(f"layer must be between 0 and {n_layers}, got {layer}.")

        accumulated = zeros_like_last_dim(b_o, axis=-1)
        b_out = self.b_out if include_mlp_biases else None
        for layer_index in range(layer):
            accumulated = _add_tensor_like_values(accumulated, b_o[layer_index])
            if b_out is not None:
                accumulated = _add_tensor_like_values(accumulated, b_out[layer_index])
        if mlp_input:
            assert layer < n_layers, "Cannot include attn_bias from beyond the final layer"
            accumulated = _add_tensor_like_values(accumulated, b_o[layer])
        return accumulated

    def all_composition_scores(self, mode: str) -> Any:
        """Return all TransformerLens-style head composition scores."""
        left = self.OV
        if mode == "Q":
            right = self.QK
        elif mode == "K":
            right = self.QK.T
        elif mode == "V":
            right = self.OV
        else:
            raise ValueError(f"mode must be one of ['Q', 'K', 'V'] not {mode}")

        scores = composition_scores(left, right, broadcast_dims=True)
        return _mask_composition_scores_to_future_layers(scores)

    def all_head_labels(self) -> list[str]:
        """Return TransformerLens-style labels for all attention heads."""
        weight_shape = shape_of(self.W_Q)
        if len(weight_shape) < 2:
            raise ValueError(f"Could not infer layer/head counts from W_Q shape {weight_shape}.")
        n_layers, n_heads = int(weight_shape[0]), int(weight_shape[1])
        return [f"L{layer}H{head}" for layer in range(n_layers) for head in range(n_heads)]

    def load_model(self) -> Any:
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except ImportError as exc:
            raise ImportError(
                "HuggingFaceModelWrapper requires model dependencies. "
                "Install them with `pip install -e '.[models]'`."
            ) from exc

        dtype_map = {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
            "auto": "auto",
        }
        torch_dtype = dtype_map.get(self.dtype, self.dtype)
        pretrained_path = self._resolve_pretrained_path()
        pretrained_kwargs = self._pretrained_kwargs()

        self.tokenizer = self._load_text_tokenizer(
            AutoTokenizer,
            pretrained_path,
            pretrained_kwargs,
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            pretrained_path,
            dtype=torch_dtype,
            trust_remote_code=self.trust_remote_code,
            **pretrained_kwargs,
            **self.load_kwargs,
        )
        if self.device is not None:
            self.model.to(self.device)
        self.model.eval()
        return self.model

    @classmethod
    def from_pretrained(cls, model_name: str, **kwargs: Any) -> HuggingFaceModelWrapper:
        """Build and load a wrapper from a pretrained Transformers model id/path."""
        wrapper = cls(
            name=model_name,
            dtype=str(kwargs.pop("dtype", "float32")),
            device=kwargs.pop("device", None),
            revision=kwargs.pop("revision", None),
            cache_dir=kwargs.pop("cache_dir", None),
            trust_remote_code=bool(kwargs.pop("trust_remote_code", False)),
            load_kwargs=dict(kwargs.pop("load_kwargs", {})),
            tokenizer_kwargs=dict(kwargs.pop("tokenizer_kwargs", {})),
            pretrained_path=kwargs.pop("pretrained_path", None),
        )
        wrapper.load_kwargs.update(kwargs)
        wrapper.load_model()
        return wrapper

    @classmethod
    def from_pretrained_no_processing(
        cls,
        model_name: str,
        **kwargs: Any,
    ) -> HuggingFaceModelWrapper:
        """Alias for ``from_pretrained``; SafeLens does not apply TL weight processing."""
        return cls.from_pretrained(model_name, **kwargs)

    def _resolve_pretrained_path(self) -> str:
        return self.pretrained_path or self.name

    def _pretrained_kwargs(self) -> dict[str, Any]:
        kwargs: dict[str, Any] = {}
        if self.revision is not None:
            kwargs["revision"] = self.revision
        if self.cache_dir is not None:
            kwargs["cache_dir"] = self.cache_dir
        return kwargs

    def _load_text_tokenizer(
        self,
        tokenizer_cls: Any,
        pretrained_path: str,
        pretrained_kwargs: dict[str, Any],
    ) -> Any | None:
        try:
            tokenizer = tokenizer_cls.from_pretrained(
                pretrained_path,
                trust_remote_code=self.trust_remote_code,
                **pretrained_kwargs,
                **self.tokenizer_kwargs,
            )
        except Exception as exc:
            self._tokenizer_load_error = exc
            return None
        self._tokenizer_load_error = None
        return tokenizer

    def add_hook(
        self,
        layer: LayerRef | Callable[[str], bool],
        hook_fn: HookFn | None = None,
        *,
        hook: HookFn | None = None,
        dir: str = "fwd",
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
    ) -> Any:
        """Register a TransformerLens-style hook on a component or module."""
        resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
        if dir == "fwd":
            return self._add_managed_hook(
                layer,
                resolved_hook,
                is_permanent=is_permanent,
                level=level,
                prepend=prepend,
            )
        if dir == "bwd":
            return self._add_managed_backward_hook(
                layer,
                resolved_hook,
                is_permanent=is_permanent,
                level=level,
                prepend=prepend,
            )
        raise ValueError(f"Invalid hook direction {dir!r}.")

    def add_perma_hook(
        self,
        layer: LayerRef,
        hook_fn: HookFn | None = None,
        *,
        hook: HookFn | None = None,
        dir: str = "fwd",
    ) -> Any:
        """Register a TransformerLens-style permanent hook."""
        return self.add_hook(layer, hook_fn, hook=hook, dir=dir, is_permanent=True)

    def _add_managed_hook(
        self,
        layer: LayerRef | Callable[[str], bool],
        hook_fn: HookFn,
        *,
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
    ) -> _ManagedWrapperHookHandle:
        if callable(layer) and not isinstance(layer, str):
            handles: list[Any] = []
            try:
                for expanded_layer, _hook_fn in self._expand_hook_specs(((layer, hook_fn),)):
                    handles.append(self._register_hook(expanded_layer, hook_fn, prepend=prepend))
            except Exception:
                _remove_wrapper_handles(handles)
                raise
            if not handles:
                return self._track_hook_handle(
                    _CompositeWrapperHookHandle(()),
                    is_permanent=is_permanent,
                    level=level,
                )
            handle: Any = _CompositeWrapperHookHandle(handles)
        else:
            handle = self._register_hook(layer, hook_fn, prepend=prepend)
        managed_handle = self._track_hook_handle(
            handle,
            is_permanent=is_permanent,
            level=level,
        )
        self._hooks.append(managed_handle)
        return managed_handle

    def _add_managed_backward_hook(
        self,
        layer: LayerRef | Callable[[str], bool],
        hook_fn: HookFn,
        *,
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
    ) -> _ManagedWrapperHookHandle:
        if callable(layer) and not isinstance(layer, str):
            handles: list[Any] = []
            try:
                for expanded_layer, _hook_fn in self._expand_hook_specs(((layer, hook_fn),)):
                    handles.append(
                        self._register_backward_hook(expanded_layer, hook_fn, prepend=prepend)
                    )
            except Exception:
                _remove_wrapper_handles(handles)
                raise
            if not handles:
                return self._track_hook_handle(
                    _CompositeWrapperHookHandle(()),
                    is_permanent=is_permanent,
                    level=level,
                )
            handle: Any = _CompositeWrapperHookHandle(handles)
        else:
            handle = self._register_backward_hook(layer, hook_fn, prepend=prepend)
        managed_handle = self._track_hook_handle(
            handle,
            is_permanent=is_permanent,
            level=level,
        )
        self._hooks.append(managed_handle)
        return managed_handle

    def __call__(
        self,
        batch: Any,
        *forward_args: Any,
        return_type: str | None | object = _DEFAULT_RETURN_TYPE,
        **kwargs: Any,
    ) -> Any:
        """Run the wrapped model directly, returning logits by default like TransformerLens."""
        call_kwargs = _merge_transformer_lens_forward_positionals(
            forward_args,
            kwargs,
            return_type=return_type,
            default_return_type="logits",
        )
        resolved_call_return_type = call_kwargs.pop("return_type", "logits")
        loss_per_token = bool(call_kwargs.pop("loss_per_token", False))
        forward_keys = {
            "prepend_bos",
            "padding_side",
            "truncate",
            "start_at_layer",
            "tokens",
            "shortformer_pos_embed",
            "attention_mask",
            "stop_at_layer",
            "past_kv_cache",
        }
        forward_kwargs = {
            key: call_kwargs.pop(key) for key in list(call_kwargs) if key in forward_keys
        }
        if forward_kwargs:
            if call_kwargs:
                model_input = _merge_extra_model_kwargs(batch, call_kwargs)
            else:
                model_input = batch
            return self.forward(
                model_input,
                return_type=resolved_call_return_type,
                loss_per_token=loss_per_token,
                **forward_kwargs,
            )
        if call_kwargs:
            model_input = _merge_extra_model_kwargs(batch, call_kwargs)
        else:
            model_input = batch
        return self._run_model_forward(
            model_input,
            return_type=resolved_call_return_type,
            loss_per_token=loss_per_token,
        )

    def forward(
        self,
        input: Any,
        return_type: str | None = "logits",
        loss_per_token: bool = False,
        prepend_bos: bool | None = None,
        padding_side: str | None = None,
        truncate: bool | None = None,
        start_at_layer: int | None = None,
        tokens: Any | None = None,
        shortformer_pos_embed: Any | None = None,
        attention_mask: Any | None = None,
        stop_at_layer: int | None = None,
        past_kv_cache: Any | None = None,
    ) -> Any:
        """TransformerLens-style explicit forward method."""
        if start_at_layer is not None or stop_at_layer is not None:
            return self._run_partial_layer_forward(
                input,
                return_type=return_type,
                loss_per_token=loss_per_token,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
                truncate=truncate,
                start_at_layer=start_at_layer,
                tokens=tokens,
                shortformer_pos_embed=shortformer_pos_embed,
                attention_mask=attention_mask,
                stop_at_layer=stop_at_layer,
                past_kv_cache=past_kv_cache,
            )
        if shortformer_pos_embed is not None:
            raise NotImplementedError(
                "SafeLens' Transformers wrapper does not synthesize "
                "shortformer positional embeddings."
            )
        batch = input
        extra_kwargs: dict[str, Any] = {}
        if prepend_bos is not None:
            extra_kwargs["prepend_bos"] = prepend_bos
        if padding_side is not None:
            extra_kwargs["padding_side"] = padding_side
        if truncate is not None:
            extra_kwargs["truncate"] = truncate
        if attention_mask is not None:
            extra_kwargs["attention_mask"] = attention_mask
        if tokens is not None:
            extra_kwargs["tokens"] = tokens
        if past_kv_cache is not None:
            extra_kwargs["past_kv_cache"] = past_kv_cache
        if extra_kwargs:
            batch = _merge_extra_model_kwargs(batch, extra_kwargs)
        return self._run_model_forward(
            batch,
            return_type=return_type,
            loss_per_token=loss_per_token,
        )

    def _run_partial_layer_forward(
        self,
        input: Any,
        *,
        return_type: str | None,
        loss_per_token: bool,
        prepend_bos: bool | None,
        padding_side: str | None,
        truncate: bool | None,
        start_at_layer: int | None,
        tokens: Any | None,
        shortformer_pos_embed: Any | None,
        attention_mask: Any | None,
        stop_at_layer: int | None,
        past_kv_cache: Any | None,
    ) -> Any:
        """Run a TransformerLens-style layer slice over decoder blocks."""
        model = self._require_model()
        if self._uses_decoder_text_input_semantics() is False:
            raise NotImplementedError(
                "SafeLens partial-layer forward currently supports decoder-only "
                "TransformerLens-style models."
            )
        model_cache, sync_back = _past_kv_cache_to_transformers_cache(past_kv_cache)

        if start_at_layer is None:
            residual, tokens, shortformer_pos_embed, attention_mask = self.input_to_embed(
                input,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
                truncate=truncate,
                attention_mask=attention_mask,
                past_kv_cache=past_kv_cache,
            )
            start_at_layer = 0
        else:
            _assert_residual_stream_input(input)
            residual = input
            if tokens is not None:
                tokens = _coerce_token_model_input(
                    _ensure_token_batch_dim(tokens), device=self.device
                )
            if attention_mask is not None:
                attention_mask = _coerce_token_model_input(
                    _ensure_token_batch_dim(attention_mask),
                    device=self.device,
                )

        blocks = _decoder_layer_modules(model)
        indexed_blocks = list(enumerate(blocks))
        pos_offset = _past_kv_position_offset_for_partial_forward(
            past_kv_cache,
            attention_mask,
        )
        position_ids = (
            _position_ids_for_forward(
                tokens,
                attention_mask=attention_mask,
                pos_offset=pos_offset,
                device=self.device,
            )
            if tokens is not None
            else None
        )
        cache_position = (
            _cache_position_for_forward(tokens, pos_offset=pos_offset, device=self.device)
            if tokens is not None
            else None
        )

        with (
            _grad_context(enabled=self._has_active_backward_hooks()),
            _temporary_eager_attention(
                model,
                enabled=self._run_requires_output_attentions or self._attention_hook_count > 0,
            ),
        ):
            try:
                for layer_index, block in indexed_blocks[start_at_layer:stop_at_layer]:
                    residual = _call_decoder_block(
                        block,
                        residual,
                        layer_index=layer_index,
                        attention_mask=attention_mask,
                        position_ids=position_ids,
                        cache_position=cache_position,
                        past_key_values=model_cache,
                        past_kv_cache_entry=_past_kv_cache_entry_at(past_kv_cache, layer_index),
                        shortformer_pos_embed=shortformer_pos_embed,
                        output_attentions=(
                            self._run_requires_output_attentions or self._attention_hook_count > 0
                        ),
                    )
            finally:
                self._run_requires_output_attentions = False

        if sync_back is not None:
            sync_back()
        if stop_at_layer is not None:
            return residual

        final_norm = _final_norm_module(model)
        if final_norm is not None and callable(final_norm):
            residual = final_norm(residual)
        if return_type is None:
            return None
        resolved_return_type = _normalize_return_type(return_type)
        logits = _unembed_residual(model, residual)
        return _format_model_output(
            {"logits": logits},
            resolved_return_type,
            model_inputs=_loss_model_inputs_for_partial_forward(
                tokens,
                attention_mask,
                return_type=resolved_return_type,
            ),
            loss_per_token=loss_per_token,
        )

    def _bridge_past_kv_cache(
        self,
        model_inputs: dict[str, Any],
        past_kv_cache: Any,
    ) -> tuple[dict[str, Any], Callable[[Any], None] | None]:
        model_cache, sync_back = _past_kv_cache_to_transformers_cache(past_kv_cache)
        if model_cache is None:
            return model_inputs, None
        bridged_inputs = dict(model_inputs)
        bridged_inputs["past_key_values"] = model_cache
        bridged_inputs["use_cache"] = True
        if "attention_mask" in bridged_inputs:
            append_attention_mask = getattr(past_kv_cache, "append_attention_mask", None)
            if callable(append_attention_mask):
                bridged_inputs["attention_mask"] = append_attention_mask(
                    bridged_inputs["attention_mask"]
                )
            else:
                bridged_inputs["attention_mask"] = _extend_attention_mask_for_past_cache(
                    bridged_inputs["attention_mask"],
                    past_kv_cache,
                )
        return bridged_inputs, lambda output: _sync_past_kv_cache_from_model_output(
            past_kv_cache,
            model_cache,
            output,
            sync_back,
        )

    def loss_fn(
        self,
        logits: Any,
        tokens: Any,
        attention_mask: Any | None = None,
        per_token: bool = False,
    ) -> Any:
        """Compute TransformerLens-style next-token cross-entropy loss."""
        return _extract_or_compute_loss(
            {"logits": logits},
            _loss_model_inputs(tokens, attention_mask),
            logits=logits,
            loss_per_token=per_token,
        )

    def _register_hook(self, layer: LayerRef, hook_fn: HookFn, *, prepend: bool = False) -> Any:
        model = self._require_model()
        layer = self._resolve_hook_layer_ref(model, layer, for_cache=False)
        self.check_hooks_to_add(None, str(layer), hook_fn, dir="fwd", prepend=prepend)
        top_level_handle = self._try_register_top_level_hook(
            model,
            layer,
            hook_fn,
            prepend=prepend,
        )
        if top_level_handle is not None:
            return top_level_handle
        component_handle = self._try_register_component_hook(
            model,
            layer,
            hook_fn,
            prepend=prepend,
        )
        if component_handle is not None:
            return component_handle
        module = self._resolve_layer(model, layer)
        return _register_module_forward_hook(
            module,
            lambda mod, inputs, output: hook_fn(mod, inputs, output),
            prepend=prepend,
        )

    def _register_backward_hook(
        self,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        model = self._require_model()
        layer = self._resolve_hook_layer_ref(model, layer, for_cache=False)
        self.check_hooks_to_add(None, str(layer), hook_fn, dir="bwd", prepend=prepend)
        top_level_handle = self._try_register_top_level_backward_hook(
            model,
            layer,
            hook_fn,
            prepend=prepend,
        )
        if top_level_handle is not None:
            return top_level_handle
        component_handle = self._try_register_component_backward_hook(
            model,
            layer,
            hook_fn,
            prepend=prepend,
        )
        if component_handle is not None:
            return component_handle
        module = self._resolve_layer(model, layer)
        return _register_raw_activation_backward_hook(
            module,
            hook_fn,
            hook_context=_RawHookContext(
                activation_name_for_layer(layer) if isinstance(layer, int) else str(layer)
            ),
            use_input=False,
            prepend=prepend,
        )

    def run_with_cache(
        self,
        batch: Any,
        *model_args: Any,
        layers: Sequence[LayerRef] | LayerRef | None = None,
        names_filter: NamesFilter = None,
        return_cache_object: bool | object = _DEFAULT_RETURN_CACHE_OBJECT,
        remove_batch_dim: bool = False,
        detach: bool = True,
        clone: bool = False,
        device: Any = None,
        pos_slice: Any = None,
        cache_all: bool | object = _DEFAULT_CACHE_ALL,
        return_type: str | None | object = _DEFAULT_RETURN_TYPE,
        loss_per_token: bool | object = _DEFAULT_LOSS_PER_TOKEN,
        incl_bwd: bool = False,
        reset_hooks_end: bool = True,
        clear_contexts: bool = False,
        **forward_kwargs: Any,
    ) -> tuple[Any, dict[str, Any] | ActivationCache]:
        model = self._require_model()
        cache = ActivationCache(model=self, has_batch_dim=not remove_batch_dim)
        temp_handles: list[Any] = []
        install_complete = False
        layers, forward_args = _split_run_with_cache_positionals(model_args, layers=layers)
        forward_options = _merge_transformer_lens_forward_positionals(
            forward_args,
            forward_kwargs,
            return_type=return_type,
            loss_per_token=loss_per_token,
        )
        forward_return_type = forward_options.pop("return_type", _DEFAULT_RETURN_TYPE)
        forward_loss_per_token = bool(forward_options.pop("loss_per_token", False))
        resolved_return_type = _resolve_return_type(batch, forward_return_type)
        resolved_cache_all = _resolve_cache_all(batch, layers, names_filter, cache_all)
        resolved_return_cache_object = _resolve_return_cache_object(batch, return_cache_object)

        try:
            self.is_caching = True
            for layer in self._cache_layers(
                model,
                layers,
                names_filter,
                cache_all=resolved_cache_all,
            ):
                temp_handles.append(
                    self._register_cache_hook(
                        model,
                        layer,
                        cache,
                        detach=detach,
                        clone=clone,
                        device=device,
                        pos_slice=pos_slice,
                        remove_batch_dim=remove_batch_dim,
                    )
                )
                if incl_bwd:
                    cache_name = self._cache_name_for_layer(model, layer, for_cache=True)
                    temp_handles.append(
                        self._register_backward_hook(
                            layer,
                            make_cache_hook(
                                cache,
                                f"{cache_name}_grad",
                                detach=detach,
                                clone=clone,
                                device=device,
                                pos_slice=pos_slice,
                                remove_batch_dim=remove_batch_dim,
                            ),
                            prepend=False,
                        )
                    )

            install_complete = True
            if forward_options:
                with self._temporary_backward_hook_context(enabled=incl_bwd):
                    output = self.forward(
                        batch,
                        return_type=resolved_return_type,
                        loss_per_token=forward_loss_per_token,
                        **forward_options,
                    )
            else:
                output = self._run_model_forward(
                    batch,
                    return_type=resolved_return_type,
                    loss_per_token=forward_loss_per_token,
                    enable_grad=incl_bwd,
                )
            if incl_bwd:
                _backward_scalar_output(output)
        finally:
            if reset_hooks_end or not install_complete:
                _remove_wrapper_handles(temp_handles)
                if clear_contexts:
                    self._clear_hook_contexts_for_handles(temp_handles)
                    self.clear_contexts()
            else:
                _keep_wrapper_handles(self, temp_handles, is_cache=True)
            self.is_caching = _wrapper_has_active_permanent_cache_hooks(self._hooks)

        return output, _format_cache_result(
            cache,
            model=self,
            return_cache_object=resolved_return_cache_object,
            remove_batch_dim=False,
        )

    def add_caching_hooks(
        self,
        names_filter: NamesFilter = None,
        *,
        layers: Sequence[LayerRef] | None = None,
        cache_all: bool | object = _DEFAULT_CACHE_ALL,
        incl_bwd: bool = False,
        detach: bool = True,
        clone: bool = False,
        device: Any = None,
        pos_slice: Any = None,
        remove_batch_dim: bool = False,
        cache: ActivationCache | dict[str, Any] | None = None,
    ) -> ActivationCache:
        """Install persistent forward cache hooks and return the live cache."""
        model = self._require_model()
        previous_is_caching = self.is_caching
        activation_cache = _coerce_activation_cache(
            cache,
            model=self,
            has_batch_dim=not remove_batch_dim,
        )
        previous_cache_model = activation_cache.model
        previous_has_batch_dim = activation_cache.has_batch_dim
        activation_cache.model = self
        if remove_batch_dim:
            activation_cache.has_batch_dim = False
        resolved_cache_all = True if cache_all is _DEFAULT_CACHE_ALL else bool(cache_all)
        installed_handles: list[_ManagedWrapperHookHandle] = []
        try:
            for layer in self._cache_layers(
                model,
                layers,
                names_filter,
                cache_all=resolved_cache_all,
            ):
                handle = self._register_cache_hook(
                    model,
                    layer,
                    activation_cache,
                    detach=detach,
                    clone=clone,
                    device=device,
                    pos_slice=pos_slice,
                    remove_batch_dim=remove_batch_dim,
                    is_permanent=True,
                )
                managed_handle = self._track_hook_handle(
                    handle,
                    is_permanent=True,
                    is_cache=True,
                )
                self._hooks.append(managed_handle)
                installed_handles.append(managed_handle)
                if incl_bwd:
                    cache_name = self._cache_name_for_layer(model, layer, for_cache=True)
                    bwd_handle = self._register_backward_hook(
                        layer,
                        make_cache_hook(
                            activation_cache,
                            f"{cache_name}_grad",
                            detach=detach,
                            clone=clone,
                            device=device,
                            pos_slice=pos_slice,
                            remove_batch_dim=remove_batch_dim,
                        ),
                        prepend=False,
                    )
                    managed_bwd_handle = self._track_hook_handle(
                        bwd_handle,
                        is_permanent=True,
                        is_cache=True,
                    )
                    self._hooks.append(managed_bwd_handle)
                    installed_handles.append(managed_bwd_handle)
        except Exception:
            for handle in reversed(installed_handles):
                handle.remove()
            self.is_caching = previous_is_caching
            activation_cache.model = previous_cache_model
            activation_cache.has_batch_dim = previous_has_batch_dim
            raise
        self.is_caching = True
        return activation_cache

    def cache_all(
        self,
        cache: ActivationCache | dict[str, Any] | None = None,
        *,
        incl_bwd: bool = False,
        detach: bool = True,
        clone: bool = False,
        device: Any = None,
        pos_slice: Any = None,
        remove_batch_dim: bool = False,
    ) -> ActivationCache:
        """Permanently cache all supported component hooks until hooks are reset."""
        return self.add_caching_hooks(
            None,
            cache_all=True,
            incl_bwd=incl_bwd,
            detach=detach,
            clone=clone,
            device=device,
            pos_slice=pos_slice,
            remove_batch_dim=remove_batch_dim,
            cache=cache,
        )

    def cache_some(
        self,
        cache_or_names_filter: ActivationCache | dict[str, Any] | NamesFilter,
        names_filter: NamesFilter = None,
        *,
        incl_bwd: bool = False,
        detach: bool = True,
        clone: bool = False,
        device: Any = None,
        pos_slice: Any = None,
        remove_batch_dim: bool = False,
        cache: ActivationCache | dict[str, Any] | None = None,
    ) -> ActivationCache:
        """Permanently cache hook names matching a names filter until hooks are reset."""
        if _looks_like_external_cache(cache_or_names_filter):
            if cache is not None:
                raise TypeError("Pass external cache either positionally or by keyword, not both.")
            cache = cast(ActivationCache | dict[str, Any], cache_or_names_filter)
        else:
            if names_filter is not None:
                raise TypeError("Pass only one names filter.")
            names_filter = cast(NamesFilter, cache_or_names_filter)
        if names_filter is None:
            raise TypeError("cache_some() missing required argument: 'names_filter' or 'names'")
        return self.add_caching_hooks(
            names_filter,
            incl_bwd=incl_bwd,
            detach=detach,
            clone=clone,
            device=device,
            pos_slice=pos_slice,
            remove_batch_dim=remove_batch_dim,
            cache=cache,
        )

    def run_with_hooks(
        self,
        batch: Any,
        *model_args: Any,
        fwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
        bwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
        prepend: bool = False,
        reset_hooks_end: bool = True,
        clear_contexts: bool = False,
        return_type: str | None | object = _DEFAULT_RETURN_TYPE,
        loss_per_token: bool | object = _DEFAULT_LOSS_PER_TOKEN,
        **forward_kwargs: Any,
    ) -> Any:
        """Run one forward pass with temporary hooks, mirroring TransformerLens."""
        bwd_hook_specs = list(bwd_hooks)
        handles: list[Any] = []
        forward_options = _merge_transformer_lens_forward_positionals(
            model_args,
            forward_kwargs,
            return_type=return_type,
            loss_per_token=loss_per_token,
        )
        forward_return_type = forward_options.pop("return_type", _DEFAULT_RETURN_TYPE)
        forward_loss_per_token = bool(forward_options.pop("loss_per_token", False))
        try:
            for layer, hook_fn in self._expand_hook_specs(fwd_hooks):
                handles.append(self._add_managed_hook(layer, hook_fn, prepend=prepend))
            for layer, hook_fn in self._expand_hook_specs(bwd_hook_specs):
                handles.append(self._add_managed_backward_hook(layer, hook_fn, prepend=prepend))
            if forward_options:
                resolved_return_type = _resolve_return_type(batch, forward_return_type)
                with self._temporary_backward_hook_context(enabled=bool(bwd_hook_specs)):
                    return self.forward(
                        batch,
                        return_type=resolved_return_type,
                        loss_per_token=forward_loss_per_token,
                        **forward_options,
                    )
            return self._run_model_forward(
                batch,
                return_type=forward_return_type,
                loss_per_token=forward_loss_per_token,
                enable_grad=bool(bwd_hook_specs),
            )
        finally:
            if reset_hooks_end:
                _remove_wrapper_handles(handles)
                if clear_contexts:
                    self._clear_hook_contexts_for_handles(handles)
                    self.clear_contexts()

    @contextmanager
    def hooks(
        self,
        fwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
        bwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
        *,
        prepend: bool = False,
        reset_hooks_end: bool = True,
        clear_contexts: bool = False,
    ) -> Any:
        """Temporarily register hooks around arbitrary wrapper calls."""
        handles: list[Any] = []
        try:
            for layer, hook_fn in self._expand_hook_specs(fwd_hooks):
                handles.append(self._add_managed_hook(layer, hook_fn, prepend=prepend))
            for layer, hook_fn in self._expand_hook_specs(bwd_hooks):
                handles.append(self._add_managed_backward_hook(layer, hook_fn, prepend=prepend))
            yield self
        finally:
            if reset_hooks_end:
                _remove_wrapper_handles(handles)
                if clear_contexts:
                    self._clear_hook_contexts_for_handles(handles)
                    self.clear_contexts()

    def _run_model_forward(
        self,
        batch: Any,
        *,
        return_type: str | None | object = _DEFAULT_RETURN_TYPE,
        loss_per_token: bool = False,
        enable_grad: bool = False,
    ) -> Any:
        """Run the wrapped model without adding temporary cache hooks."""
        model = self._require_model()
        resolved_return_type = _resolve_return_type(batch, return_type)
        model_inputs = self._prepare_model_inputs(batch)
        past_kv_cache = model_inputs.pop("past_kv_cache", None)
        cache_bridge = None
        if past_kv_cache is not None:
            model_inputs, cache_bridge = self._bridge_past_kv_cache(model_inputs, past_kv_cache)
        try:
            with (
                _grad_context(enabled=enable_grad or self._has_active_backward_hooks()),
                _temporary_eager_attention(
                    model,
                    enabled=self._run_requires_output_attentions or self._attention_hook_count > 0,
                ),
            ):
                raw_output = model(**model_inputs)
                formatted_output = _format_model_output(
                    raw_output,
                    resolved_return_type,
                    model_inputs=model_inputs,
                    loss_per_token=loss_per_token,
                )
                if cache_bridge is not None:
                    cache_bridge(raw_output)
                return formatted_output
        finally:
            self._run_requires_output_attentions = False

    def _has_active_backward_hooks(self) -> bool:
        return self._temporary_backward_hook_depth > 0 or any(
            getattr(handle, "is_backward", False) for handle in self._hooks
        )

    @contextmanager
    def _temporary_backward_hook_context(self, *, enabled: bool) -> Any:
        if not enabled:
            yield
            return
        self._temporary_backward_hook_depth += 1
        try:
            yield
        finally:
            self._temporary_backward_hook_depth = max(
                0,
                self._temporary_backward_hook_depth - 1,
            )

    def _cache_name_for_layer(
        self,
        model: Any,
        layer: LayerRef,
        *,
        for_cache: bool | None,
    ) -> str:
        resolved_layer = self._resolve_hook_layer_ref(model, layer, for_cache=for_cache)
        top_level_name = _canonical_top_level_hook_name(resolved_layer)
        if top_level_name is not None:
            return top_level_name
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        component_ref = adapter.parse_component_ref(resolved_layer)
        if component_ref is not None:
            if isinstance(resolved_layer, tuple):
                return component_ref.transformer_lens_name
            return (
                activation_name_for_layer(resolved_layer)
                if isinstance(resolved_layer, int)
                else str(resolved_layer)
            )
        return (
            activation_name_for_layer(resolved_layer)
            if isinstance(resolved_layer, int)
            else str(resolved_layer)
        )

    def _uses_encoder_decoder_generation_semantics(self) -> bool:
        model = getattr(self, "model", None)
        config = _core_model_config(_config_attr(model, "config"))
        return bool(_config_attr(config, "is_encoder_decoder", False))

    def _generation_should_append_input_embeds(self) -> bool:
        return not self._uses_encoder_decoder_generation_semantics()

    def generate(self, prompt: Any = "", **generation_kwargs: Any) -> Any:
        model = self._require_model()

        import torch

        prepend_bos = _resolve_default_prepend_bos(
            self,
            generation_kwargs.pop("prepend_bos", None),
        )
        input_is_text = isinstance(prompt, str) or _is_text_batch(prompt)
        input_is_embeds = _looks_like_input_embeds(prompt)
        input_ids: Any | None = None
        input_embeds: Any | None = None
        default_padding_side = "left" if _is_text_batch(prompt) else None
        padding_side = generation_kwargs.pop("padding_side", default_padding_side)
        truncate = bool(generation_kwargs.pop("truncate", True))
        return_type = generation_kwargs.pop("return_type", "input")
        _normalize_transformer_lens_generation_kwargs(generation_kwargs, model=model)
        normalized_return_type = _normalize_generate_return_type(
            return_type,
            prompt=prompt,
            input_is_text=input_is_text,
            input_is_embeds=input_is_embeds,
        )
        if self.tokenizer is None and (input_is_text or normalized_return_type == "str"):
            detail = _tokenizer_error_detail(self._tokenizer_load_error)
            raise RuntimeError(
                f"Tokenizer is not loaded, so text generation is unavailable. {detail}"
            )
        if input_is_text:
            input_ids = self.to_tokens(
                prompt,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
                truncate=truncate,
            )
        elif _looks_like_token_ids(prompt):
            input_ids = _ensure_token_batch_dim(prompt)
            if not isinstance(input_ids, torch.Tensor):
                input_ids = torch.as_tensor(input_ids, dtype=torch.long)
            if self.device is not None:
                to_fn = getattr(input_ids, "to", None)
                if callable(to_fn):
                    input_ids = to_fn(self.device)
        elif input_is_embeds:
            input_ids = None
            input_embeds = prompt
            if not isinstance(input_embeds, torch.Tensor):
                input_embeds = torch.as_tensor(input_embeds)
            if not torch.is_floating_point(input_embeds):
                raise TypeError(
                    "generate embedding inputs must be floating point tensors shaped "
                    "[batch, pos, hidden]."
                )
            if self.device is not None:
                input_embeds = input_embeds.to(self.device)
        else:
            raise TypeError(
                "generate input must be a text string, list of strings, token ids shaped "
                "[pos] or [batch, pos], or embeddings shaped [batch, pos, hidden]."
            )
        inputs = {"inputs_embeds": input_embeds} if input_is_embeds else {"input_ids": input_ids}
        if not input_is_embeds:
            attention_mask = _generation_attention_mask_from_tokens(
                input_ids,
                tokenizer=self.tokenizer,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
            )
            if attention_mask is not None:
                inputs["attention_mask"] = attention_mask
        with torch.no_grad():
            generated_output = model.generate(**inputs, **generation_kwargs)
        output_ids = _generated_sequences(generated_output)
        if normalized_return_type == "model_output":
            return _normalize_generated_model_output(generated_output, output_ids)
        if normalized_return_type == "tokens":
            return output_ids
        if normalized_return_type == "embeds":
            output_embeds = _tokens_to_input_embeddings(model, output_ids)
            if input_is_embeds and self._generation_should_append_input_embeds():
                return _append_sequence_values(input_embeds, output_embeds)
            return output_embeds
        return _decode_generated_sequences(
            self.tokenizer,
            output_ids,
            force_batch=_is_text_batch(prompt),
        )

    def generate_stream(
        self,
        prompt: Any = "",
        *,
        max_new_tokens: int = 10,
        max_tokens_per_yield: int = 25,
        **generation_kwargs: Any,
    ) -> Iterable[Any]:
        model = self._require_model()
        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative.")
        if max_tokens_per_yield <= 0:
            raise ValueError("max_tokens_per_yield must be positive.")

        import torch

        prepend_bos = _resolve_default_prepend_bos(
            self,
            generation_kwargs.pop("prepend_bos", None),
        )
        input_is_text = isinstance(prompt, str)
        if _is_text_batch(prompt):
            raise TypeError(
                "generate_stream supports a single text string or token ids; pass token ids "
                "for batched streaming inputs."
            )
        input_is_embeds = _looks_like_input_embeds(prompt)
        padding_side = generation_kwargs.pop("padding_side", None)
        truncate = bool(generation_kwargs.pop("truncate", True))
        return_type = generation_kwargs.pop("return_type", "input")
        _normalize_transformer_lens_generation_kwargs(generation_kwargs, model=model)
        normalized_return_type = _normalize_generate_return_type(
            return_type,
            prompt=prompt,
            input_is_text=input_is_text,
            input_is_embeds=input_is_embeds,
            token_return_name="tensor",
        )
        if normalized_return_type == "model_output":
            raise ValueError("generate_stream does not support return_type='model_output'.")
        if self.tokenizer is None and (input_is_text or normalized_return_type == "str"):
            detail = _tokenizer_error_detail(self._tokenizer_load_error)
            raise RuntimeError(
                f"Tokenizer is not loaded, so streaming text generation is unavailable. {detail}"
            )

        if input_is_text:
            tokens = self.to_tokens(
                prompt,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
                truncate=truncate,
            )
        elif _looks_like_token_ids(prompt):
            tokens = _ensure_token_batch_dim(prompt)
            if not isinstance(tokens, torch.Tensor):
                tokens = torch.as_tensor(tokens, dtype=torch.long)
            if self.device is not None:
                tokens = tokens.to(self.device)
        elif input_is_embeds:
            input_embeds = prompt if isinstance(prompt, torch.Tensor) else torch.as_tensor(prompt)
            if not torch.is_floating_point(input_embeds):
                raise TypeError(
                    "generate_stream embedding inputs must be floating point tensors shaped "
                    "[batch, pos, hidden]."
                )
            if self.device is not None:
                input_embeds = input_embeds.to(self.device)
            yield from self._generate_stream_from_embeds(
                model,
                input_embeds,
                max_new_tokens=max_new_tokens,
                max_tokens_per_yield=max_tokens_per_yield,
                return_type=normalized_return_type,
                generation_kwargs=generation_kwargs,
            )
            return
        else:
            raise TypeError(
                "generate_stream input must be a text string, token ids shaped [pos] or "
                "[batch, pos], or embeddings shaped [batch, pos, hidden]."
            )

        yield from self._generate_stream_from_tokens(
            model,
            tokens,
            max_new_tokens=max_new_tokens,
            max_tokens_per_yield=max_tokens_per_yield,
            return_type=normalized_return_type,
            generation_kwargs=generation_kwargs,
            text_output_is_single=input_is_text,
            padding_side=padding_side,
        )

    def _generate_stream_from_tokens(
        self,
        model: Any,
        tokens: Any,
        *,
        max_new_tokens: int,
        max_tokens_per_yield: int,
        return_type: str,
        generation_kwargs: dict[str, Any],
        text_output_is_single: bool,
        padding_side: str | None,
    ) -> Iterable[Any]:
        import torch

        if self._uses_encoder_decoder_generation_semantics():
            yield from self._generate_encoder_decoder_stream_from_tokens(
                tokens,
                max_new_tokens=max_new_tokens,
                max_tokens_per_yield=max_tokens_per_yield,
                return_type=return_type,
                generation_kwargs=generation_kwargs,
                text_output_is_single=text_output_is_single,
                padding_side=padding_side,
            )
            return

        current_tokens = tokens
        accumulated: list[Any] = []
        with torch.no_grad():
            for _index in range(max_new_tokens):
                attention_mask = _generation_attention_mask_from_tokens(
                    current_tokens,
                    tokenizer=self.tokenizer,
                    prepend_bos=False,
                    padding_side=padding_side,
                )
                step_kwargs = dict(generation_kwargs)
                if attention_mask is not None:
                    step_kwargs["attention_mask"] = attention_mask
                step_output = model.generate(
                    input_ids=current_tokens,
                    max_new_tokens=1,
                    **step_kwargs,
                )
                step_sequences = _generated_sequences(step_output)
                new_tokens = _slice_generated_suffix(step_sequences, current_tokens)
                accumulated.append(new_tokens)
                current_tokens = _append_sequence_values(current_tokens, new_tokens)
                if len(accumulated) >= max_tokens_per_yield:
                    yield self._format_stream_chunk(
                        accumulated,
                        return_type=return_type,
                        text_output_is_single=text_output_is_single,
                    )
                    accumulated = []
        if accumulated:
            yield self._format_stream_chunk(
                accumulated,
                return_type=return_type,
                text_output_is_single=text_output_is_single,
            )

    def _generate_stream_from_embeds(
        self,
        model: Any,
        input_embeds: Any,
        *,
        max_new_tokens: int,
        max_tokens_per_yield: int,
        return_type: str,
        generation_kwargs: dict[str, Any],
    ) -> Iterable[Any]:
        import torch

        if self._uses_encoder_decoder_generation_semantics():
            yield from self._generate_encoder_decoder_stream_from_embeds(
                input_embeds,
                max_new_tokens=max_new_tokens,
                max_tokens_per_yield=max_tokens_per_yield,
                return_type=return_type,
                generation_kwargs=generation_kwargs,
            )
            return

        current_embeds = input_embeds
        accumulated: list[Any] = []
        with torch.no_grad():
            for _index in range(max_new_tokens):
                step_output = model.generate(
                    inputs_embeds=current_embeds,
                    max_new_tokens=1,
                    **generation_kwargs,
                )
                step_sequences = _generated_sequences(step_output)
                new_tokens = _normalize_generated_new_tokens(step_sequences)
                accumulated.append(new_tokens)
                new_embeds = _tokens_to_input_embeddings(model, new_tokens)
                current_embeds = _append_sequence_values(current_embeds, new_embeds)
                if len(accumulated) >= max_tokens_per_yield:
                    yield self._format_stream_chunk(
                        accumulated,
                        return_type=return_type,
                        text_output_is_single=False,
                    )
                    accumulated = []
        if accumulated:
            yield self._format_stream_chunk(
                accumulated,
                return_type=return_type,
                text_output_is_single=False,
            )

    def _generate_encoder_decoder_stream_from_tokens(
        self,
        tokens: Any,
        *,
        max_new_tokens: int,
        max_tokens_per_yield: int,
        return_type: str,
        generation_kwargs: dict[str, Any],
        text_output_is_single: bool,
        padding_side: str | None,
    ) -> Iterable[Any]:
        if max_new_tokens == 0:
            return
        output_tokens = self.generate(
            tokens,
            max_new_tokens=max_new_tokens,
            return_type="tokens",
            prepend_bos=False,
            padding_side=padding_side,
            **generation_kwargs,
        )
        for chunk in _iter_generated_token_chunks(output_tokens, max_tokens_per_yield):
            yield self._format_stream_chunk(
                [chunk],
                return_type=return_type,
                text_output_is_single=text_output_is_single,
            )

    def _generate_encoder_decoder_stream_from_embeds(
        self,
        input_embeds: Any,
        *,
        max_new_tokens: int,
        max_tokens_per_yield: int,
        return_type: str,
        generation_kwargs: dict[str, Any],
    ) -> Iterable[Any]:
        if max_new_tokens == 0:
            return
        output_tokens = self.generate(
            input_embeds,
            max_new_tokens=max_new_tokens,
            return_type="tokens",
            **generation_kwargs,
        )
        for chunk in _iter_generated_token_chunks(output_tokens, max_tokens_per_yield):
            yield self._format_stream_chunk(
                [chunk],
                return_type=return_type,
                text_output_is_single=False,
            )

    def _format_stream_chunk(
        self,
        chunks: Sequence[Any],
        *,
        return_type: str,
        text_output_is_single: bool,
    ) -> Any:
        combined_tokens = _concat_token_chunks(chunks)
        if return_type == "str":
            return _decode_generated_sequences(
                self.tokenizer,
                combined_tokens,
                force_batch=not text_output_is_single,
            )
        if return_type in {"tokens", "tensor"}:
            return combined_tokens
        if return_type == "embeds":
            return _tokens_to_input_embeddings(self._require_model(), combined_tokens)
        raise ValueError(
            "generate_stream return_type must be 'input', 'str', 'tokens', or 'embeds'."
        )

    def to_tokens(
        self,
        text: str | Sequence[str],
        *,
        prepend_bos: bool | None = None,
        padding_side: str | None = None,
        move_to_device: bool = True,
        truncate: bool = True,
    ) -> Any:
        """Tokenize text into a tensor, mirroring TransformerLens' convenience method."""
        tokenizer = self._require_tokenizer_for_text("tokenization")
        prepend_bos = _resolve_default_prepend_bos(self, prepend_bos)
        token_kwargs: dict[str, Any] = {
            "return_tensors": "pt",
            "add_special_tokens": False,
            "padding": not isinstance(text, str),
        }
        with (
            _temporary_tokenizer_padding_side(tokenizer, padding_side),
            _temporary_tokenizer_pad_token(tokenizer, enabled=bool(token_kwargs["padding"])),
        ):
            effective_pad_token_id = _tokenizer_effective_pad_token_id(tokenizer)
            effective_padding_side = str(getattr(tokenizer, "padding_side", "right"))
            if truncate:
                n_ctx = _tokenization_context_length(self.model, tokenizer)
                token_kwargs["truncation"] = True
                if n_ctx is not None:
                    max_length = int(n_ctx) - (1 if prepend_bos else 0)
                    if max_length > 0:
                        token_kwargs["max_length"] = max_length
            tokenized = _call_tokenizer_with_supported_kwargs(
                tokenizer,
                text,
                token_kwargs,
            )
        tokens = tokenized["input_ids"] if isinstance(tokenized, dict) else tokenized.input_ids
        if prepend_bos:
            tokens = _prepend_bos_token(
                tokens,
                tokenizer,
                pad_token_id=effective_pad_token_id,
                padding_side=effective_padding_side,
            )
        if move_to_device and self.device is not None:
            tokens = tokens.to(self.device)
        return tokens

    def to_string(
        self,
        tokens: Any,
        *,
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: bool = False,
    ) -> str | list[str]:
        """Decode token ids into text."""
        tokenizer = self._require_tokenizer_for_text("decoding")
        shape = _shape_of_token_ids(tokens)
        if shape is not None and len(shape) > 2:
            raise ValueError(f"Invalid token shape for decoding: {shape!r}.")
        if shape is not None and len(shape) == 2:
            batch_decode = getattr(tokenizer, "batch_decode", None)
            if callable(batch_decode):
                return list(
                    _call_decode_with_supported_kwargs(
                        batch_decode,
                        tokens,
                        {
                            "skip_special_tokens": skip_special_tokens,
                            "clean_up_tokenization_spaces": clean_up_tokenization_spaces,
                        },
                    )
                )
            return [
                str(
                    _call_decode_with_supported_kwargs(
                        tokenizer.decode,
                        row,
                        {
                            "skip_special_tokens": skip_special_tokens,
                            "clean_up_tokenization_spaces": clean_up_tokenization_spaces,
                        },
                    )
                )
                for row in tokens
            ]
        if isinstance(tokens, int):
            tokens = [tokens]
        return str(
            _call_decode_with_supported_kwargs(
                tokenizer.decode,
                tokens,
                {
                    "skip_special_tokens": skip_special_tokens,
                    "clean_up_tokenization_spaces": clean_up_tokenization_spaces,
                },
            )
        )

    def to_str_tokens(
        self,
        text_or_tokens: str | Any,
        *,
        prepend_bos: bool | None = None,
        padding_side: str | None = None,
    ) -> list[str] | list[list[str]]:
        """Return per-token strings for text or token ids."""
        tokenizer = self._require_tokenizer_for_text("token string conversion")
        resolved_prepend_bos = _resolve_default_prepend_bos(self, prepend_bos)
        if (
            isinstance(text_or_tokens, Sequence)
            and not isinstance(text_or_tokens, str | bytes)
            and text_or_tokens
            and isinstance(
                text_or_tokens[0],
                Sequence | str,
            )
        ):
            return cast(
                list[list[str]],
                [
                    self.to_str_tokens(
                        item,
                        prepend_bos=resolved_prepend_bos,
                        padding_side=padding_side,
                    )
                    for item in text_or_tokens
                ],
            )
        tokens = (
            self.to_tokens(
                text_or_tokens,
                prepend_bos=resolved_prepend_bos,
                padding_side=padding_side,
            )
            if isinstance(text_or_tokens, str)
            else text_or_tokens
        )
        shape = getattr(tokens, "shape", None)
        if shape is not None:
            shape_tuple = tuple(int(dim) for dim in shape)
            if len(shape_tuple) == 2 and shape_tuple[0] == 1:
                tokens = tokens[0]
            elif len(shape_tuple) > 1:
                raise ValueError(
                    f"Invalid token shape for token string conversion: {shape_tuple!r}."
                )
        token_list = _single_token_list(tokens)
        batch_decode = getattr(tokenizer, "batch_decode", None)
        if callable(batch_decode):
            return [
                str(token)
                for token in _call_decode_with_supported_kwargs(
                    batch_decode,
                    [[token] for token in token_list],
                    {"clean_up_tokenization_spaces": False},
                )
            ]
        convert = getattr(tokenizer, "convert_ids_to_tokens", None)
        if callable(convert):
            converted = convert(token_list)
            if isinstance(converted, str):
                return [converted]
            if isinstance(converted, Sequence) and not isinstance(converted, bytes):
                return [str(token) for token in converted]
        return [str(tokenizer.decode([token])) for token in token_list]

    def to_single_token(self, text: str) -> int:
        """Return the single token id for text or raise when it tokenizes to multiple ids."""
        tokens = self.to_tokens(text, prepend_bos=False)
        shape = getattr(tokens, "shape", None)
        token_values = tokens.reshape(-1).tolist() if shape is not None else list(tokens)
        if len(token_values) != 1:
            raise ValueError(
                f"Expected {text!r} to tokenize to a single token, got {token_values}."
            )
        return int(token_values[0])

    def to_single_str_token(self, token: int) -> str:
        """Return the string for a single token id."""
        if not isinstance(token, int):
            raise TypeError(f"Expected an integer token id, got {type(token)!r}.")
        tokens = self.to_str_tokens([token])
        if len(tokens) != 1 or isinstance(tokens[0], list):
            raise ValueError(f"Expected token id {token!r} to decode to one string token.")
        return str(tokens[0])

    def get_token_position(
        self,
        single_token: str | int | Any,
        text_or_tokens: str | Any,
        *,
        mode: str = "first",
        prepend_bos: bool | None = None,
        padding_side: str | None = None,
    ) -> int:
        """Return the first or last position of one token in a prompt or token sequence."""
        tokens = (
            self.to_tokens(
                text_or_tokens,
                prepend_bos=prepend_bos,
                padding_side=padding_side,
            )
            if isinstance(text_or_tokens, str)
            else text_or_tokens
        )
        token_values = _flatten_single_token_sequence(tokens)
        if isinstance(single_token, str):
            token_id = self.to_single_token(single_token)
        else:
            item = getattr(single_token, "item", None)
            token_id = int(cast(Any, item() if callable(item) else single_token))
        positions = [index for index, value in enumerate(token_values) if int(value) == token_id]
        if not positions:
            raise ValueError("The token does not occur in the prompt.")
        if mode == "first":
            return positions[0]
        if mode == "last":
            return positions[-1]
        raise ValueError(f"mode must be 'first' or 'last', not {mode!r}.")

    def tokens_to_residual_directions(self, tokens: Any) -> Any:
        """Map token ids to unembedding residual directions."""
        weight = self._output_embedding_weight()
        if isinstance(tokens, str):
            tokens = self.to_single_token(tokens)
        elif isinstance(tokens, int):
            pass
        else:
            numel = getattr(tokens, "numel", None)
            item = getattr(tokens, "item", None)
            if callable(numel) and callable(item):
                try:
                    if int(cast(Any, numel())) == 1:
                        tokens = int(cast(Any, item()))
                except Exception:
                    pass
        try:
            if isinstance(weight, Sequence) and not isinstance(weight, str | bytes):
                return _gather_sequence_residual_directions(weight, tokens)
            tokens = _coerce_tokens_for_weight_index(weight, tokens)
            return weight[tokens]
        except Exception as exc:
            raise RuntimeError(
                "Could not index residual directions with the provided tokens."
            ) from exc

    def remove_hooks(self) -> None:
        self.reset_hooks(including_permanent=True)

    def reset_hooks(
        self,
        *,
        clear_contexts: bool = True,
        direction: Any = None,
        dir: Any = None,
        including_permanent: bool = False,
        level: int | None = None,
    ) -> None:
        """Remove wrapper-managed hooks using TransformerLens reset semantics."""
        hook_direction = _normalize_hook_direction(direction or dir or "both")
        removed_handles: list[_ManagedWrapperHookHandle] = []
        for handle in reversed(list(self._hooks)):
            if handle.is_permanent and not including_permanent:
                continue
            if level is not None and handle.level != level:
                continue
            if not _managed_handle_matches_direction(handle, hook_direction):
                continue
            removed_handles.append(handle)
            handle.remove()
        if clear_contexts:
            self._clear_hook_contexts_for_handles(removed_handles)
            self.clear_contexts()
        self.is_caching = _wrapper_has_active_permanent_cache_hooks(self._hooks)

    def clear_contexts(self) -> None:
        """Clear mutable context dictionaries on component hook objects."""
        self._clear_hook_contexts_for_handles(list(self._hooks))

    @staticmethod
    def _clear_hook_contexts_for_handles(handles: Iterable[Any]) -> None:
        for handle in handles:
            for hook_context in getattr(handle, "hook_contexts", ()):
                clear = getattr(getattr(hook_context, "ctx", None), "clear", None)
                if callable(clear):
                    clear()

    def _track_hook_handle(
        self,
        handle: Any,
        *,
        is_permanent: bool = False,
        level: int | None = None,
        is_cache: bool = False,
    ) -> _ManagedWrapperHookHandle:
        managed_handle: _ManagedWrapperHookHandle

        def untrack() -> None:
            if managed_handle in self._hooks:
                self._hooks.remove(managed_handle)

        managed_handle = _ManagedWrapperHookHandle(
            handle,
            untrack,
            is_permanent=is_permanent,
            level=level,
            is_cache=is_cache,
        )
        return managed_handle

    def _require_model(self) -> Any:
        if self.model is None:
            self.load_model()
        return self.model

    def _require_tokenizer_for_text(self, operation: str) -> Any:
        if self.tokenizer is None:
            detail = _tokenizer_error_detail(self._tokenizer_load_error)
            raise RuntimeError(f"Tokenizer is not loaded, so {operation} is unavailable. {detail}")
        return self.tokenizer

    def _output_embedding_weight(self) -> Any:
        embeddings = self._output_embeddings()
        weight = getattr(embeddings, "weight", None)
        model = self._require_model()
        if weight is None:
            weight = getattr(model, "W_U", None)
            if weight is not None:
                return transpose_2d_weight(weight)
        if weight is None:
            raise RuntimeError(
                "Could not find output embedding weights for residual direction lookup."
            )
        return weight

    def _output_embeddings(self) -> Any:
        model = self._require_model()
        get_output_embeddings = getattr(model, "get_output_embeddings", None)
        return get_output_embeddings() if callable(get_output_embeddings) else None

    def _stack_attention_weights(self, component: str) -> Any:
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        n_layers = _infer_model_layers(model)
        if n_layers <= 0:
            raise RuntimeError(f"Could not infer layer count for W_{component.upper()}.")
        weights = [
            adapter.get_attention_weight(model, component, layer) for layer in range(n_layers)
        ]
        return _stack_tensor_like(weights)

    def _stack_attention_biases(self, component: str) -> Any:
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        n_layers = _infer_model_layers(model)
        if n_layers <= 0:
            raise RuntimeError(f"Could not infer layer count for b_{component.upper()}.")
        biases = [adapter.get_attention_bias(model, component, layer) for layer in range(n_layers)]
        return _stack_tensor_like(biases)

    def _stack_mlp_weights(self, component: str) -> Any:
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        n_layers = _infer_model_layers(model)
        if n_layers <= 0:
            raise RuntimeError(f"Could not infer layer count for W_{component}.")
        weights = [adapter.get_mlp_weight(model, component, layer) for layer in range(n_layers)]
        return _stack_tensor_like(weights)

    def _stack_mlp_biases(self, component: str) -> Any:
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        n_layers = _infer_model_layers(model)
        if n_layers <= 0:
            raise RuntimeError(f"Could not infer layer count for b_{component}.")
        biases = [adapter.get_mlp_bias(model, component, layer) for layer in range(n_layers)]
        return _stack_tensor_like(biases)

    def _prepare_model_inputs(self, batch: Any) -> dict[str, Any]:
        batch = _normalize_model_batch(batch)
        model_kwargs = _model_kwargs_without_tokenization(batch)
        if "input_ids" in batch:
            model_inputs = {
                key: value
                for key, value in batch.items()
                if key not in {"model_kwargs", "prepend_bos", "padding_side", "truncate"}
            }
            model_inputs.update(model_kwargs)
            model_inputs["input_ids"] = _coerce_token_model_input(
                _ensure_token_batch_dim(model_inputs["input_ids"]),
                device=self.device,
            )
            _coerce_token_mask_fields(model_inputs, device=self.device)
            return self._with_attention_flags(model_inputs)
        for token_key in ("tokens", "token_ids"):
            if token_key in batch:
                model_inputs = {
                    key: value
                    for key, value in batch.items()
                    if key
                    not in {
                        "tokens",
                        "token_ids",
                        "model_kwargs",
                        "prepend_bos",
                        "padding_side",
                        "truncate",
                    }
                }
                model_inputs.update(model_kwargs)
                model_inputs["input_ids"] = _coerce_token_model_input(
                    _ensure_token_batch_dim(batch[token_key]),
                    device=self.device,
                )
                _coerce_token_mask_fields(model_inputs, device=self.device)
                return self._with_attention_flags(model_inputs)
        if self.tokenizer is not None and "input_ids" not in batch:
            text = _text_or_prompt_value(batch)
            if text is not None:
                if self._uses_decoder_text_input_semantics():
                    prepared = _prepare_text_inputs_with_to_tokens(self, batch)
                    if prepared is not None:
                        return self._with_attention_flags(prepared)
                tokenized = _tokenize_text_batch(self.tokenizer, text)
                if self.device is not None:
                    tokenized = tokenized.to(self.device)
                return self._with_attention_flags(dict(tokenized))
        if self.tokenizer is None and "input_ids" not in batch:
            text = _text_or_prompt_value(batch)
            if text is not None:
                detail = _tokenizer_error_detail(self._tokenizer_load_error)
                raise ValueError(
                    "This model did not load a tokenizer, so text batches cannot be "
                    f"tokenized. Provide `input_ids` or `inputs_embeds` directly. {detail}"
                )
        return self._with_attention_flags(dict(batch))

    def _try_register_component_hook(
        self,
        model: Any,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any | None:
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        if adapter.parse_component_ref(layer) is None:
            return None
        requires_output_attentions = adapter.requires_output_attentions(layer)
        handle = adapter.register_component_hook(model, layer, hook_fn, prepend=prepend)
        if requires_output_attentions:
            self._attention_hook_count += 1
            return _TrackedAttentionHandle(handle, self._release_attention_hook)
        return handle

    def _try_register_component_backward_hook(
        self,
        model: Any,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any | None:
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        component_ref = adapter.parse_component_ref(layer)
        if component_ref is None:
            return None
        spec = adapter._spec_for_ref(component_ref, for_cache=False)
        module = adapter.get_component(model, component_ref)
        if spec.value in {"attention_pattern", "attention_scores"}:
            raise NotImplementedError(
                "Backward hooks for attention pattern/scores require custom attention "
                "softmax instrumentation and are not yet supported."
            )
        if spec.component == "result":
            raise NotImplementedError(
                "Backward hooks for derived attention result activations are not yet supported."
            )
        hook_context = ComponentHookContext(component_ref)
        if spec.mode == "forward_input":
            input_hook = _make_component_input_backward_registration_hook(
                hook_fn,
                component_ref,
                adapter.name,
                spec,
                model,
                hook_context,
            )
            return _BackwardHookRegistrationHandle(
                _register_module_forward_pre_hook(module, input_hook, prepend=prepend),
                (hook_context,),
            )
        output_hook = _make_component_output_backward_registration_hook(
            hook_fn,
            component_ref,
            adapter.name,
            spec,
            model,
            hook_context,
        )
        return _BackwardHookRegistrationHandle(
            _register_module_forward_hook(module, output_hook, prepend=prepend),
            (hook_context,),
        )

    def _try_register_top_level_hook(
        self,
        model: Any,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any | None:
        hook_name = _canonical_top_level_hook_name(layer)
        if hook_name is None:
            return None
        module = _resolve_top_level_hook_module(model, hook_name)
        if module is None:
            return None
        hook_context = _TopLevelHookContext(hook_name)

        def hook(_module: Any, _inputs: Any, output: Any) -> Any:
            activation = (
                _final_norm_scale_from_hook(_module, _inputs, output)
                if hook_name == "ln_final.hook_scale"
                else output
            )
            patched = _call_top_level_hook(
                hook_fn,
                activation=activation,
                hook_name=hook_name,
                hook_context=hook_context,
            )
            if hook_name == "ln_final.hook_scale":
                if patched is None:
                    return None
                return _replace_final_norm_output_from_scale(_module, _inputs, output, patched)
            return None if patched is None else patched

        return _HookHandleWithContexts(
            _register_module_forward_hook(module, hook, prepend=prepend),
            (hook_context,),
        )

    def _try_register_top_level_backward_hook(
        self,
        model: Any,
        layer: LayerRef,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any | None:
        hook_name = _canonical_top_level_hook_name(layer)
        if hook_name is None:
            return None
        if hook_name == "ln_final.hook_scale":
            raise NotImplementedError(
                "Backward hooks for derived normalization scale activations are not yet supported."
            )
        module = _resolve_top_level_hook_module(model, hook_name)
        if module is None:
            return None
        hook_context = _TopLevelHookContext(hook_name)
        hook = _make_raw_output_backward_registration_hook(
            hook_fn,
            hook_context,
            hook_name=hook_name,
        )
        return _BackwardHookRegistrationHandle(
            _register_module_forward_hook(module, hook, prepend=prepend),
            (hook_context,),
        )

    def _resolve_hook_layer_ref(
        self,
        model: Any,
        layer: LayerRef,
        *,
        for_cache: bool | None,
    ) -> LayerRef:
        if not isinstance(layer, str):
            return layer
        top_level_name = _canonical_top_level_hook_name(layer)
        if top_level_name is not None and _top_level_hook_is_resolvable(model, top_level_name):
            return top_level_name
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        if adapter.parse_component_ref(layer) is not None:
            return layer
        matched = _filter_hook_names(
            _candidate_hook_names(model, adapter, for_cache=for_cache),
            layer,
            adapter=adapter,
        )
        if len(matched) == 1:
            return matched[0]
        return layer

    def _try_register_component_cache_hook(
        self,
        model: Any,
        layer: LayerRef,
        cache: ActivationCache,
        *,
        detach: bool,
        clone: bool,
        device: Any,
        pos_slice: Any,
        remove_batch_dim: bool,
        is_permanent: bool = False,
    ) -> Any | None:
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        component_ref = adapter.parse_component_ref(layer)
        if component_ref is None:
            return None
        if isinstance(layer, tuple):
            cache_name = component_ref.transformer_lens_name
        else:
            cache_name = activation_name_for_layer(layer) if isinstance(layer, int) else str(layer)
        if _cache_hook_requires_runtime_flag(cache_name):
            self.check_hooks_to_add(None, cache_name, None)

        handle = adapter.register_component_hook_for_mode(
            model,
            layer,
            make_cache_hook(
                cache,
                cache_name,
                detach=detach,
                clone=clone,
                device=device,
                pos_slice=pos_slice,
                remove_batch_dim=remove_batch_dim,
            ),
            for_cache=True,
        )
        requires_output_attentions = adapter.requires_output_attentions(layer)
        if requires_output_attentions and is_permanent:
            self._attention_hook_count += 1
            return _TrackedAttentionHandle(handle, self._release_attention_hook)
        if requires_output_attentions:
            self._run_requires_output_attentions = True
        return handle

    def _try_register_top_level_cache_hook(
        self,
        model: Any,
        layer: LayerRef,
        cache: ActivationCache,
        *,
        detach: bool,
        clone: bool,
        device: Any,
        pos_slice: Any,
        remove_batch_dim: bool,
    ) -> Any | None:
        hook_name = _canonical_top_level_hook_name(layer)
        if hook_name is None:
            return None
        module = _resolve_top_level_hook_module(model, hook_name)
        if module is None:
            return None
        if hook_name == "ln_final.hook_scale":
            return _register_module_forward_hook(
                module,
                make_final_norm_scale_cache_hook(
                    cache,
                    hook_name,
                    detach=detach,
                    clone=clone,
                    device=device,
                    pos_slice=pos_slice,
                    remove_batch_dim=remove_batch_dim,
                ),
                prepend=False,
            )
        return _register_module_forward_hook(
            module,
            make_cache_hook(
                cache,
                hook_name,
                detach=detach,
                clone=clone,
                device=device,
                pos_slice=pos_slice,
                remove_batch_dim=remove_batch_dim,
            ),
            prepend=False,
        )

    def _register_cache_hook(
        self,
        model: Any,
        layer: LayerRef,
        cache: ActivationCache,
        *,
        detach: bool,
        clone: bool,
        device: Any,
        pos_slice: Any,
        remove_batch_dim: bool,
        is_permanent: bool = False,
    ) -> Any:
        layer = self._resolve_hook_layer_ref(model, layer, for_cache=True)
        top_level_handle = self._try_register_top_level_cache_hook(
            model,
            layer,
            cache,
            detach=detach,
            clone=clone,
            device=device,
            pos_slice=pos_slice,
            remove_batch_dim=remove_batch_dim,
        )
        if top_level_handle is not None:
            return top_level_handle
        component_handle = self._try_register_component_cache_hook(
            model,
            layer,
            cache,
            detach=detach,
            clone=clone,
            device=device,
            pos_slice=pos_slice,
            remove_batch_dim=remove_batch_dim,
            is_permanent=is_permanent,
        )
        if component_handle is not None:
            return component_handle
        module = self._resolve_layer(model, layer)
        cache_name = activation_name_for_layer(layer)
        return _register_module_forward_hook(
            module,
            make_cache_hook(
                cache,
                cache_name,
                detach=detach,
                clone=clone,
                device=device,
                pos_slice=pos_slice,
                remove_batch_dim=remove_batch_dim,
            ),
            prepend=False,
        )

    def _with_attention_flags(self, model_inputs: dict[str, Any]) -> dict[str, Any]:
        if self._attention_hook_count > 0 or self._run_requires_output_attentions:
            model_inputs.setdefault("output_attentions", True)
        return model_inputs

    def _release_attention_hook(self) -> None:
        self._attention_hook_count = max(0, self._attention_hook_count - 1)

    def _expand_hook_specs(
        self,
        hook_specs: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]],
    ) -> list[tuple[LayerRef, HookFn]]:
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        names = _candidate_hook_names(model, adapter, for_cache=False)
        expanded: list[tuple[LayerRef, HookFn]] = []
        for layer_or_filter, hook_fn in hook_specs:
            if callable(layer_or_filter) and not isinstance(layer_or_filter, str):
                matched = _filter_hook_names(names, layer_or_filter, adapter=adapter)
                expanded.extend((name, hook_fn) for name in matched)
                continue
            if isinstance(layer_or_filter, str):
                matched = _filter_hook_names(names, layer_or_filter, adapter=adapter)
                if matched:
                    expanded.extend((name, hook_fn) for name in matched)
                    continue
            expanded.append((layer_or_filter, hook_fn))
        return expanded

    def _cache_layers(
        self,
        model: Any,
        layers: Sequence[LayerRef] | None,
        names_filter: NamesFilter,
        *,
        cache_all: bool = False,
    ) -> list[LayerRef]:
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        normalized_layers = _normalize_cache_layers_arg(layers)
        if normalized_layers is not None:
            selected = normalized_layers
        elif names_filter is not None:
            selected = cast(
                list[LayerRef],
                _filter_hook_names(
                    _candidate_hook_names(model, adapter, for_cache=True),
                    names_filter,
                    adapter=adapter,
                ),
            )
        elif cache_all:
            selected = cast(list[LayerRef], _default_cache_hook_names(model, adapter))
        else:
            selected = []
        return selected

    @staticmethod
    def _resolve_layer(model: Any, layer: LayerRef) -> Any:
        if isinstance(layer, str):
            modules = dict(model.named_modules())
            if layer not in modules:
                examples = ", ".join(list(modules)[:8])
                raise KeyError(
                    f"Unknown module or hook name {layer!r}. Use an integer layer index, "
                    "a module name from model.named_modules(), or a component hook name "
                    "supported by the selected model adapter. "
                    f"First available module names: {examples}"
                )
            return modules[layer]

        for path in (
            "model.layers",
            "model.language_model.layers",
            "transformer.h",
            "gpt_neox.layers",
        ):
            target = model
            try:
                for part in path.split("."):
                    target = getattr(target, part)
                return target[layer]
            except (AttributeError, IndexError, TypeError):
                continue
        raise KeyError(
            f"Could not resolve layer index {layer} for model {type(model).__name__}. "
            "Known decoder-layer paths tried: model.layers, model.language_model.layers, "
            "transformer.h, gpt_neox.layers."
        )

OV property

Return the TransformerLens-style OV circuit as a factored matrix.

QK property

Return the TransformerLens-style QK circuit as a factored matrix.

W_E property

Return token embedding weights shaped [vocab, d_model].

W_E_pos property

Return concatenated token and positional embeddings.

W_K property

Return key weights shaped [layer, head, d_model, d_head].

W_O property

Return output weights shaped [layer, head, d_head, d_model].

W_Q property

Return query weights shaped [layer, head, d_model, d_head].

W_U property

Return an unembedding matrix shaped [d_model, vocab] when available.

W_U_S property

Return the singular values of the unembedding matrix.

W_U_U property

Return the left singular vectors of the unembedding matrix.

W_U_V property

Return the right singular vectors of the unembedding matrix.

W_V property

Return value weights shaped [layer, head, d_model, d_head].

W_gate property

Return gated-MLP gate weights shaped [layer, d_model, d_mlp].

W_in property

Return MLP input weights shaped [layer, d_model, d_mlp].

W_out property

Return MLP output weights shaped [layer, d_mlp, d_model].

W_pos property

Return positional embedding weights when available.

b_K property

Return key biases shaped [layer, head, d_head].

b_O property

Return attention output biases shaped [layer, d_model].

b_Q property

Return query biases shaped [layer, head, d_head].

b_U property

Return unembedding bias shaped [vocab] when available, else zeros.

b_V property

Return value biases shaped [layer, head, d_head].

b_in property

Return MLP input biases shaped [layer, d_mlp].

b_out property

Return MLP output biases shaped [layer, d_model].

cfg property

Return a TransformerLens-style normalized config view.

n_params_total property

Return the wrapped model's total parameter count.

__call__(batch, *forward_args, return_type=_DEFAULT_RETURN_TYPE, **kwargs)

Run the wrapped model directly, returning logits by default like TransformerLens.

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

def __call__(
    self,
    batch: Any,
    *forward_args: Any,
    return_type: str | None | object = _DEFAULT_RETURN_TYPE,
    **kwargs: Any,
) -> Any:
    """Run the wrapped model directly, returning logits by default like TransformerLens."""
    call_kwargs = _merge_transformer_lens_forward_positionals(
        forward_args,
        kwargs,
        return_type=return_type,
        default_return_type="logits",
    )
    resolved_call_return_type = call_kwargs.pop("return_type", "logits")
    loss_per_token = bool(call_kwargs.pop("loss_per_token", False))
    forward_keys = {
        "prepend_bos",
        "padding_side",
        "truncate",
        "start_at_layer",
        "tokens",
        "shortformer_pos_embed",
        "attention_mask",
        "stop_at_layer",
        "past_kv_cache",
    }
    forward_kwargs = {
        key: call_kwargs.pop(key) for key in list(call_kwargs) if key in forward_keys
    }
    if forward_kwargs:
        if call_kwargs:
            model_input = _merge_extra_model_kwargs(batch, call_kwargs)
        else:
            model_input = batch
        return self.forward(
            model_input,
            return_type=resolved_call_return_type,
            loss_per_token=loss_per_token,
            **forward_kwargs,
        )
    if call_kwargs:
        model_input = _merge_extra_model_kwargs(batch, call_kwargs)
    else:
        model_input = batch
    return self._run_model_forward(
        model_input,
        return_type=resolved_call_return_type,
        loss_per_token=loss_per_token,
    )

accumulated_bias(layer, mlp_input=False, include_mlp_biases=True)

Return accumulated attention/MLP output biases before a layer.

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

def accumulated_bias(
    self,
    layer: int,
    mlp_input: bool = False,
    include_mlp_biases: bool = True,
) -> Any:
    """Return accumulated attention/MLP output biases before a layer."""
    b_o = self.b_O
    n_layers = _stack_first_dim(b_o)
    if layer < 0 or layer > n_layers:
        raise ValueError(f"layer must be between 0 and {n_layers}, got {layer}.")

    accumulated = zeros_like_last_dim(b_o, axis=-1)
    b_out = self.b_out if include_mlp_biases else None
    for layer_index in range(layer):
        accumulated = _add_tensor_like_values(accumulated, b_o[layer_index])
        if b_out is not None:
            accumulated = _add_tensor_like_values(accumulated, b_out[layer_index])
    if mlp_input:
        assert layer < n_layers, "Cannot include attn_bias from beyond the final layer"
        accumulated = _add_tensor_like_values(accumulated, b_o[layer])
    return accumulated

add_caching_hooks(names_filter=None, *, layers=None, cache_all=_DEFAULT_CACHE_ALL, incl_bwd=False, detach=True, clone=False, device=None, pos_slice=None, remove_batch_dim=False, cache=None)

Install persistent forward cache hooks and return the live cache.

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

def add_caching_hooks(
    self,
    names_filter: NamesFilter = None,
    *,
    layers: Sequence[LayerRef] | None = None,
    cache_all: bool | object = _DEFAULT_CACHE_ALL,
    incl_bwd: bool = False,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
    pos_slice: Any = None,
    remove_batch_dim: bool = False,
    cache: ActivationCache | dict[str, Any] | None = None,
) -> ActivationCache:
    """Install persistent forward cache hooks and return the live cache."""
    model = self._require_model()
    previous_is_caching = self.is_caching
    activation_cache = _coerce_activation_cache(
        cache,
        model=self,
        has_batch_dim=not remove_batch_dim,
    )
    previous_cache_model = activation_cache.model
    previous_has_batch_dim = activation_cache.has_batch_dim
    activation_cache.model = self
    if remove_batch_dim:
        activation_cache.has_batch_dim = False
    resolved_cache_all = True if cache_all is _DEFAULT_CACHE_ALL else bool(cache_all)
    installed_handles: list[_ManagedWrapperHookHandle] = []
    try:
        for layer in self._cache_layers(
            model,
            layers,
            names_filter,
            cache_all=resolved_cache_all,
        ):
            handle = self._register_cache_hook(
                model,
                layer,
                activation_cache,
                detach=detach,
                clone=clone,
                device=device,
                pos_slice=pos_slice,
                remove_batch_dim=remove_batch_dim,
                is_permanent=True,
            )
            managed_handle = self._track_hook_handle(
                handle,
                is_permanent=True,
                is_cache=True,
            )
            self._hooks.append(managed_handle)
            installed_handles.append(managed_handle)
            if incl_bwd:
                cache_name = self._cache_name_for_layer(model, layer, for_cache=True)
                bwd_handle = self._register_backward_hook(
                    layer,
                    make_cache_hook(
                        activation_cache,
                        f"{cache_name}_grad",
                        detach=detach,
                        clone=clone,
                        device=device,
                        pos_slice=pos_slice,
                        remove_batch_dim=remove_batch_dim,
                    ),
                    prepend=False,
                )
                managed_bwd_handle = self._track_hook_handle(
                    bwd_handle,
                    is_permanent=True,
                    is_cache=True,
                )
                self._hooks.append(managed_bwd_handle)
                installed_handles.append(managed_bwd_handle)
    except Exception:
        for handle in reversed(installed_handles):
            handle.remove()
        self.is_caching = previous_is_caching
        activation_cache.model = previous_cache_model
        activation_cache.has_batch_dim = previous_has_batch_dim
        raise
    self.is_caching = True
    return activation_cache

add_hook(layer, hook_fn=None, *, hook=None, dir='fwd', is_permanent=False, level=None, prepend=False)

Register a TransformerLens-style hook on a component or module.

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

def add_hook(
    self,
    layer: LayerRef | Callable[[str], bool],
    hook_fn: HookFn | None = None,
    *,
    hook: HookFn | None = None,
    dir: str = "fwd",
    is_permanent: bool = False,
    level: int | None = None,
    prepend: bool = False,
) -> Any:
    """Register a TransformerLens-style hook on a component or module."""
    resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
    if dir == "fwd":
        return self._add_managed_hook(
            layer,
            resolved_hook,
            is_permanent=is_permanent,
            level=level,
            prepend=prepend,
        )
    if dir == "bwd":
        return self._add_managed_backward_hook(
            layer,
            resolved_hook,
            is_permanent=is_permanent,
            level=level,
            prepend=prepend,
        )
    raise ValueError(f"Invalid hook direction {dir!r}.")

add_perma_hook(layer, hook_fn=None, *, hook=None, dir='fwd')

Register a TransformerLens-style permanent hook.

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

def add_perma_hook(
    self,
    layer: LayerRef,
    hook_fn: HookFn | None = None,
    *,
    hook: HookFn | None = None,
    dir: str = "fwd",
) -> Any:
    """Register a TransformerLens-style permanent hook."""
    return self.add_hook(layer, hook_fn, hook=hook, dir=dir, is_permanent=True)

all_composition_scores(mode)

Return all TransformerLens-style head composition scores.

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

def all_composition_scores(self, mode: str) -> Any:
    """Return all TransformerLens-style head composition scores."""
    left = self.OV
    if mode == "Q":
        right = self.QK
    elif mode == "K":
        right = self.QK.T
    elif mode == "V":
        right = self.OV
    else:
        raise ValueError(f"mode must be one of ['Q', 'K', 'V'] not {mode}")

    scores = composition_scores(left, right, broadcast_dims=True)
    return _mask_composition_scores_to_future_layers(scores)

all_head_labels()

Return TransformerLens-style labels for all attention heads.

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

def all_head_labels(self) -> list[str]:
    """Return TransformerLens-style labels for all attention heads."""
    weight_shape = shape_of(self.W_Q)
    if len(weight_shape) < 2:
        raise ValueError(f"Could not infer layer/head counts from W_Q shape {weight_shape}.")
    n_layers, n_heads = int(weight_shape[0]), int(weight_shape[1])
    return [f"L{layer}H{head}" for layer in range(n_layers) for head in range(n_heads)]

cache_all(cache=None, *, incl_bwd=False, detach=True, clone=False, device=None, pos_slice=None, remove_batch_dim=False)

Permanently cache all supported component hooks until hooks are reset.

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

def cache_all(
    self,
    cache: ActivationCache | dict[str, Any] | None = None,
    *,
    incl_bwd: bool = False,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
    pos_slice: Any = None,
    remove_batch_dim: bool = False,
) -> ActivationCache:
    """Permanently cache all supported component hooks until hooks are reset."""
    return self.add_caching_hooks(
        None,
        cache_all=True,
        incl_bwd=incl_bwd,
        detach=detach,
        clone=clone,
        device=device,
        pos_slice=pos_slice,
        remove_batch_dim=remove_batch_dim,
        cache=cache,
    )

cache_some(cache_or_names_filter, names_filter=None, *, incl_bwd=False, detach=True, clone=False, device=None, pos_slice=None, remove_batch_dim=False, cache=None)

Permanently cache hook names matching a names filter until hooks are reset.

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

def cache_some(
    self,
    cache_or_names_filter: ActivationCache | dict[str, Any] | NamesFilter,
    names_filter: NamesFilter = None,
    *,
    incl_bwd: bool = False,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
    pos_slice: Any = None,
    remove_batch_dim: bool = False,
    cache: ActivationCache | dict[str, Any] | None = None,
) -> ActivationCache:
    """Permanently cache hook names matching a names filter until hooks are reset."""
    if _looks_like_external_cache(cache_or_names_filter):
        if cache is not None:
            raise TypeError("Pass external cache either positionally or by keyword, not both.")
        cache = cast(ActivationCache | dict[str, Any], cache_or_names_filter)
    else:
        if names_filter is not None:
            raise TypeError("Pass only one names filter.")
        names_filter = cast(NamesFilter, cache_or_names_filter)
    if names_filter is None:
        raise TypeError("cache_some() missing required argument: 'names_filter' or 'names'")
    return self.add_caching_hooks(
        names_filter,
        incl_bwd=incl_bwd,
        detach=detach,
        clone=clone,
        device=device,
        pos_slice=pos_slice,
        remove_batch_dim=remove_batch_dim,
        cache=cache,
    )

center_unembed(*args, **kwargs)

Center unembedding directions, matching TransformerLens weight processing.

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

def center_unembed(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
    """Center unembedding directions, matching TransformerLens weight processing."""
    _ = args, kwargs
    model = self._require_model()
    native_weight = getattr(model, "W_U", None)
    if native_weight is not None:
        centered = _center_unembed_weight(native_weight, weight_layout="d_model_vocab")
        try:
            model.W_U = centered
        except Exception as exc:
            raise RuntimeError("Could not update native TransformerLens W_U.") from exc
        native_bias = getattr(model, "b_U", None)
        if native_bias is not None:
            try:
                model.b_U = _center_bias_like(native_bias)
            except Exception as exc:
                raise RuntimeError("Could not update native TransformerLens b_U.") from exc
        return self

    embeddings = self._output_embeddings()
    weight = getattr(embeddings, "weight", None)
    if weight is None:
        raise RuntimeError("Could not find unembedding weights to center.")
    centered = _center_unembed_weight(weight, weight_layout="vocab_d_model")
    updated = False
    try:
        embeddings.weight = centered
        updated = True
    except Exception as exc:
        data = getattr(weight, "data", None)
        if data is None:
            raise RuntimeError("Could not update output embedding weight.") from exc
        data.copy_(centered)
        updated = True
    if updated and hasattr(model, "_weight"):
        try:
            model._weight = centered
        except Exception:
            pass
    bias = getattr(embeddings, "bias", None)
    if bias is not None:
        centered_bias = _center_bias_like(bias)
        _set_module_bias(embeddings, centered_bias)
        if hasattr(model, "_bias"):
            try:
                model._bias = centered_bias
            except Exception:
                pass
    return self

center_writing_weights(*args, **kwargs)

Center weights and biases that write directly to the residual stream.

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

def center_writing_weights(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
    """Center weights and biases that write directly to the residual stream."""
    _ = args, kwargs
    model = self._require_model()
    if _is_olmo2_post_norm_model(model):
        return self
    config = getattr(model, "config", None)
    _center_module_weight(_input_embeddings_module(model), axis=-1)
    try:
        positional_embeddings = _positional_embeddings_module(model)
    except KeyError:
        positional_embeddings = None
    if positional_embeddings is not None:
        _center_module_weight(positional_embeddings, axis=-1)

    native_w_o = getattr(model, "W_O", None)
    if native_w_o is not None:
        model.W_O = _center_residual_stream_weight(native_w_o, d_model=self.cfg.d_model)
    native_b_o = getattr(model, "b_O", None)
    if native_b_o is not None:
        model.b_O = _center_residual_stream_bias(native_b_o)
    native_w_out = getattr(model, "W_out", None)
    if native_w_out is not None:
        model.W_out = _center_residual_stream_weight(native_w_out, d_model=self.cfg.d_model)
    native_b_out = getattr(model, "b_out", None)
    if native_b_out is not None:
        model.b_out = _center_residual_stream_bias(native_b_out)

    adapter = architecture_adapter_for_model(model, model_name=self.name)
    for layer in range(_infer_model_layers(model)):
        try:
            attention_ref = adapter.parse_component_ref(
                transformer_lens_component_name("z", layer)
            )
            if attention_ref is not None:
                attention_module = adapter.get_component(model, attention_ref)
                _center_attention_output_module(
                    attention_module,
                    architecture=adapter.name,
                )
        except (KeyError, NotImplementedError, ValueError):
            pass
        if _config_attr(config, "attn_only", False):
            continue
        try:
            mlp_ref = adapter.parse_component_ref(
                transformer_lens_component_name("post", layer)
            )
            if mlp_ref is not None:
                mlp_module = adapter.get_component(model, mlp_ref)
                _center_mlp_output_module(mlp_module, d_model=self.cfg.d_model)
        except (KeyError, NotImplementedError, ValueError):
            pass
    return self

check_hooks_to_add(hook_point, hook_point_name, hook, dir='fwd', is_permanent=False, prepend=False)

Validate TransformerLens runtime-hook flags before adding hooks.

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

def check_hooks_to_add(
    self,
    hook_point: Any,
    hook_point_name: str,
    hook: Any,
    dir: str = "fwd",
    is_permanent: bool = False,
    prepend: bool = False,
) -> None:
    """Validate TransformerLens runtime-hook flags before adding hooks."""
    _ = hook_point, hook, dir, is_permanent, prepend
    if hook_point_name.endswith("attn.hook_result") or hook_point_name.endswith("hook_result"):
        assert (
            self.cfg.use_attn_result
        ), f"Cannot add hook {hook_point_name} if use_attn_result_hook is False"
    if hook_point_name.endswith(("hook_q_input", "hook_k_input", "hook_v_input")):
        assert (
            self.cfg.use_split_qkv_input
        ), f"Cannot add hook {hook_point_name} if use_split_qkv_input is False"
    if hook_point_name.endswith("mlp_in"):
        assert (
            self.cfg.use_hook_mlp_in
        ), f"Cannot add hook {hook_point_name} if use_hook_mlp_in is False"
    if hook_point_name.endswith("attn_in"):
        assert (
            self.cfg.use_attn_in
        ), f"Cannot add hook {hook_point_name} if use_attn_in is False"

clear_contexts()

Clear mutable context dictionaries on component hook objects.

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

def clear_contexts(self) -> None:
    """Clear mutable context dictionaries on component hook objects."""
    self._clear_hook_contexts_for_handles(list(self._hooks))

fill_missing_keys(*args, **kwargs)

Unsupported HookedTransformer state-dict helper.

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

def fill_missing_keys(self, *args: Any, **kwargs: Any) -> None:
    """Unsupported HookedTransformer state-dict helper."""
    _ = args, kwargs
    raise NotImplementedError(
        "SafeLens' dependency-free wrapper does not mutate HookedTransformer-format "
        "state dictionaries."
    )

fold_layer_norm(*args, **kwargs)

Fold LayerNorm/RMSNorm affine parameters into reader weights in place.

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

def fold_layer_norm(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
    """Fold LayerNorm/RMSNorm affine parameters into reader weights in place."""
    _ = args
    fold_biases_override = kwargs.pop("fold_biases", None)
    center_weights_override = kwargs.pop("center_weights", None)
    if kwargs:
        unexpected = ", ".join(sorted(kwargs))
        raise TypeError(f"Unexpected fold_layer_norm keyword argument(s): {unexpected}.")
    model = self._require_model()
    normalization_type = self.cfg.normalization_type or "LN"
    is_rms_norm = normalization_type in {"RMS", "RMSPre"}
    fold_biases = (
        (not is_rms_norm) if fold_biases_override is None else bool(fold_biases_override)
    )
    center_weights = (
        (not is_rms_norm) if center_weights_override is None else bool(center_weights_override)
    )
    _fold_layer_norm_weights_in_model(
        model,
        model_name=self.name,
        fold_biases=fold_biases,
        center_weights=center_weights,
        rmsnorm_uses_offset=self.cfg.rmsnorm_uses_offset,
    )
    return self

fold_value_biases(*args, **kwargs)

Fold attention value biases into output biases and clear value biases.

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

def fold_value_biases(self, *args: Any, **kwargs: Any) -> HuggingFaceModelWrapper:
    """Fold attention value biases into output biases and clear value biases."""
    _ = args, kwargs
    model = self._require_model()
    native_b_v = getattr(model, "b_V", None)
    native_w_o = getattr(model, "W_O", None)
    if native_b_v is not None and native_w_o is not None:
        native_b_o = getattr(model, "b_O", None)
        folded_bias, zero_b_v = _fold_value_bias_tensor_like(
            native_b_v,
            native_w_o,
            native_b_o,
            target_heads=self.cfg.n_heads,
        )
        model.b_O = folded_bias
        model.b_V = zero_b_v

    adapter = architecture_adapter_for_model(model, model_name=self.name)
    for layer in range(_infer_model_layers(model)):
        try:
            value_ref = adapter.parse_component_ref(transformer_lens_component_name("v", layer))
            output_ref = adapter.parse_component_ref(
                transformer_lens_component_name("z", layer)
            )
            if value_ref is None or output_ref is None:
                continue
            value_spec = adapter._spec_for_ref(value_ref, for_cache=True)
            value_module = adapter.get_component(model, value_ref)
            output_module = adapter.get_component(model, output_ref)
            _fold_value_bias_modules(
                model,
                value_module=value_module,
                output_module=output_module,
                value_spec=value_spec,
                architecture=adapter.name,
            )
        except (KeyError, NotImplementedError, ValueError):
            pass
    return self

forward(input, return_type='logits', loss_per_token=False, prepend_bos=None, padding_side=None, truncate=None, start_at_layer=None, tokens=None, shortformer_pos_embed=None, attention_mask=None, stop_at_layer=None, past_kv_cache=None)

TransformerLens-style explicit forward method.

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

def forward(
    self,
    input: Any,
    return_type: str | None = "logits",
    loss_per_token: bool = False,
    prepend_bos: bool | None = None,
    padding_side: str | None = None,
    truncate: bool | None = None,
    start_at_layer: int | None = None,
    tokens: Any | None = None,
    shortformer_pos_embed: Any | None = None,
    attention_mask: Any | None = None,
    stop_at_layer: int | None = None,
    past_kv_cache: Any | None = None,
) -> Any:
    """TransformerLens-style explicit forward method."""
    if start_at_layer is not None or stop_at_layer is not None:
        return self._run_partial_layer_forward(
            input,
            return_type=return_type,
            loss_per_token=loss_per_token,
            prepend_bos=prepend_bos,
            padding_side=padding_side,
            truncate=truncate,
            start_at_layer=start_at_layer,
            tokens=tokens,
            shortformer_pos_embed=shortformer_pos_embed,
            attention_mask=attention_mask,
            stop_at_layer=stop_at_layer,
            past_kv_cache=past_kv_cache,
        )
    if shortformer_pos_embed is not None:
        raise NotImplementedError(
            "SafeLens' Transformers wrapper does not synthesize "
            "shortformer positional embeddings."
        )
    batch = input
    extra_kwargs: dict[str, Any] = {}
    if prepend_bos is not None:
        extra_kwargs["prepend_bos"] = prepend_bos
    if padding_side is not None:
        extra_kwargs["padding_side"] = padding_side
    if truncate is not None:
        extra_kwargs["truncate"] = truncate
    if attention_mask is not None:
        extra_kwargs["attention_mask"] = attention_mask
    if tokens is not None:
        extra_kwargs["tokens"] = tokens
    if past_kv_cache is not None:
        extra_kwargs["past_kv_cache"] = past_kv_cache
    if extra_kwargs:
        batch = _merge_extra_model_kwargs(batch, extra_kwargs)
    return self._run_model_forward(
        batch,
        return_type=return_type,
        loss_per_token=loss_per_token,
    )

from_pretrained(model_name, **kwargs) classmethod

Build and load a wrapper from a pretrained Transformers model id/path.

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

@classmethod
def from_pretrained(cls, model_name: str, **kwargs: Any) -> HuggingFaceModelWrapper:
    """Build and load a wrapper from a pretrained Transformers model id/path."""
    wrapper = cls(
        name=model_name,
        dtype=str(kwargs.pop("dtype", "float32")),
        device=kwargs.pop("device", None),
        revision=kwargs.pop("revision", None),
        cache_dir=kwargs.pop("cache_dir", None),
        trust_remote_code=bool(kwargs.pop("trust_remote_code", False)),
        load_kwargs=dict(kwargs.pop("load_kwargs", {})),
        tokenizer_kwargs=dict(kwargs.pop("tokenizer_kwargs", {})),
        pretrained_path=kwargs.pop("pretrained_path", None),
    )
    wrapper.load_kwargs.update(kwargs)
    wrapper.load_model()
    return wrapper

from_pretrained_no_processing(model_name, **kwargs) classmethod

Alias for from_pretrained; SafeLens does not apply TL weight processing.

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

@classmethod
def from_pretrained_no_processing(
    cls,
    model_name: str,
    **kwargs: Any,
) -> HuggingFaceModelWrapper:
    """Alias for ``from_pretrained``; SafeLens does not apply TL weight processing."""
    return cls.from_pretrained(model_name, **kwargs)

get_pos_offset(past_kv_cache, batch_size)

Return the positional offset implied by a TransformerLens KV cache.

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

def get_pos_offset(self, past_kv_cache: Any, batch_size: int) -> int:
    """Return the positional offset implied by a TransformerLens KV cache."""
    if past_kv_cache is None:
        return 0
    cached_batch_size = _past_kv_cache_batch_size(past_kv_cache)
    if cached_batch_size is not None:
        assert int(cached_batch_size) == int(batch_size)
    return _past_kv_cache_length(past_kv_cache)

get_residual(embed, pos_offset, prepend_bos=None, attention_mask=None, tokens=None, return_shortformer_pos_embed=True, device=None)

Convert token embeddings into the first residual stream.

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

def get_residual(
    self,
    embed: Any,
    pos_offset: int,
    prepend_bos: bool | None = None,
    attention_mask: Any | None = None,
    tokens: Any | None = None,
    return_shortformer_pos_embed: bool = True,
    device: Any = None,
) -> Any:
    """Convert token embeddings into the first residual stream."""
    model = self._require_model()
    resolved_device = self.device if device is None else device
    if tokens is None:
        tokens = _ones_token_batch_like_embedding(embed, device=resolved_device)
    else:
        tokens = _coerce_token_model_input(
            _ensure_token_batch_dim(tokens), device=resolved_device
        )

    position_type = _infer_positional_embedding_type(model)
    if position_type == "standard":
        pos_module = _positional_embeddings_module(model)
        if pos_module is None:
            residual = embed
            shortformer_pos_embed = None
        else:
            pos_embed = _call_position_embedding_module(
                pos_module,
                tokens,
                pos_offset=pos_offset,
                attention_mask=attention_mask,
                device=resolved_device,
            )
            residual = _add_tensor_like_values(embed, pos_embed)
            shortformer_pos_embed = None
    elif position_type == "shortformer":
        pos_module = _positional_embeddings_module(model)
        if pos_module is None:
            raise NotImplementedError(
                "SafeLens cannot compute shortformer positional embeddings without "
                "a resolvable positional embedding module."
            )
        shortformer_pos_embed = _call_position_embedding_module(
            pos_module,
            tokens,
            pos_offset=pos_offset,
            attention_mask=attention_mask,
            device=resolved_device,
        )
        residual = embed
    elif position_type in {"rotary", "alibi", None}:
        residual = embed
        shortformer_pos_embed = None
    else:
        raise ValueError(f"Invalid positional_embedding_type {position_type!r}.")

    if return_shortformer_pos_embed:
        return residual, shortformer_pos_embed
    return residual

get_token_position(single_token, text_or_tokens, *, mode='first', prepend_bos=None, padding_side=None)

Return the first or last position of one token in a prompt or token sequence.

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

def get_token_position(
    self,
    single_token: str | int | Any,
    text_or_tokens: str | Any,
    *,
    mode: str = "first",
    prepend_bos: bool | None = None,
    padding_side: str | None = None,
) -> int:
    """Return the first or last position of one token in a prompt or token sequence."""
    tokens = (
        self.to_tokens(
            text_or_tokens,
            prepend_bos=prepend_bos,
            padding_side=padding_side,
        )
        if isinstance(text_or_tokens, str)
        else text_or_tokens
    )
    token_values = _flatten_single_token_sequence(tokens)
    if isinstance(single_token, str):
        token_id = self.to_single_token(single_token)
    else:
        item = getattr(single_token, "item", None)
        token_id = int(cast(Any, item() if callable(item) else single_token))
    positions = [index for index, value in enumerate(token_values) if int(value) == token_id]
    if not positions:
        raise ValueError("The token does not occur in the prompt.")
    if mode == "first":
        return positions[0]
    if mode == "last":
        return positions[-1]
    raise ValueError(f"mode must be 'first' or 'last', not {mode!r}.")

hooks(fwd_hooks=(), bwd_hooks=(), *, prepend=False, reset_hooks_end=True, clear_contexts=False)

Temporarily register hooks around arbitrary wrapper calls.

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

@contextmanager
def hooks(
    self,
    fwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
    bwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
    *,
    prepend: bool = False,
    reset_hooks_end: bool = True,
    clear_contexts: bool = False,
) -> Any:
    """Temporarily register hooks around arbitrary wrapper calls."""
    handles: list[Any] = []
    try:
        for layer, hook_fn in self._expand_hook_specs(fwd_hooks):
            handles.append(self._add_managed_hook(layer, hook_fn, prepend=prepend))
        for layer, hook_fn in self._expand_hook_specs(bwd_hooks):
            handles.append(self._add_managed_backward_hook(layer, hook_fn, prepend=prepend))
        yield self
    finally:
        if reset_hooks_end:
            _remove_wrapper_handles(handles)
            if clear_contexts:
                self._clear_hook_contexts_for_handles(handles)
                self.clear_contexts()

init_weights()

Unsupported HookedTransformer weight-initialization helper.

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

def init_weights(self) -> None:
    """Unsupported HookedTransformer weight-initialization helper."""
    raise NotImplementedError(
        "SafeLens wraps pretrained Transformers modules and does not initialize "
        "HookedTransformer-format weights."
    )

input_to_embed(input, prepend_bos=None, padding_side=None, truncate=None, attention_mask=None, past_kv_cache=None)

Convert token/text input to TL-style (residual, tokens, pos_embed, mask).

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

def input_to_embed(
    self,
    input: Any,
    prepend_bos: bool | None = None,
    padding_side: str | None = None,
    truncate: bool | None = None,
    attention_mask: Any | None = None,
    past_kv_cache: Any | None = None,
) -> tuple[Any, Any, Any | None, Any | None]:
    """Convert token/text input to TL-style `(residual, tokens, pos_embed, mask)`."""
    model = self._require_model()
    if isinstance(input, Mapping):
        if prepend_bos is None and "prepend_bos" in input:
            prepend_bos = input["prepend_bos"]
        if padding_side is None and "padding_side" in input:
            padding_side = input["padding_side"]
        if truncate is None and "truncate" in input:
            truncate = input["truncate"]
        if attention_mask is None and "attention_mask" in input:
            attention_mask = input["attention_mask"]
        if "input_ids" in input:
            input = input["input_ids"]
        elif "tokens" in input:
            input = input["tokens"]
        elif "token_ids" in input:
            input = input["token_ids"]
        else:
            text = _text_or_prompt_value(input)
            if text is not None:
                input = text
    if isinstance(input, str) or _is_text_batch(input):
        resolved_truncate = True if truncate is None else bool(truncate)
        text_input = cast(str | Sequence[str], input)
        tokens = self.to_tokens(
            text_input,
            prepend_bos=prepend_bos,
            padding_side=padding_side,
            truncate=resolved_truncate,
        )
    else:
        tokens = input
    tokens = _coerce_token_model_input(_ensure_token_batch_dim(tokens), device=self.device)

    if attention_mask is None:
        effective_padding_side = str(
            padding_side or getattr(self.tokenizer, "padding_side", "right")
        )
        needs_mask = effective_padding_side == "left" or past_kv_cache is not None
        if needs_mask and self.tokenizer is not None:
            resolved_prepend_bos = _resolve_default_prepend_bos(self, prepend_bos)
            attention_mask = _attention_mask_from_tokens(
                tokens,
                _tokenizer_effective_pad_token_id(self.tokenizer),
                prepend_bos=resolved_prepend_bos,
                padding_side=effective_padding_side,
                bos_token_id=getattr(self.tokenizer, "bos_token_id", None),
            )
    if attention_mask is not None:
        attention_mask = _coerce_token_model_input(
            _ensure_token_batch_dim(attention_mask),
            device=self.device,
        )
        if _shape_of_token_ids(attention_mask) != _shape_of_token_ids(tokens):
            raise AssertionError(
                f"Attention mask shape {_shape_of_token_ids(attention_mask)!r} "
                f"does not match tokens shape {_shape_of_token_ids(tokens)!r}."
            )
        append_attention_mask = getattr(past_kv_cache, "append_attention_mask", None)
        if callable(append_attention_mask):
            attention_mask = append_attention_mask(attention_mask)
        elif past_kv_cache is not None:
            attention_mask = _extend_attention_mask_for_past_cache(
                attention_mask, past_kv_cache
            )

    pos_offset = self.get_pos_offset(past_kv_cache, _token_batch_size(tokens))
    embed_module = _input_embeddings_module(model)
    if embed_module is None:
        raise RuntimeError("Could not resolve an input embedding module.")
    embed = embed_module(tokens)
    residual, shortformer_pos_embed = self.get_residual(
        embed,
        pos_offset,
        prepend_bos=prepend_bos,
        attention_mask=attention_mask,
        tokens=tokens,
        return_shortformer_pos_embed=True,
        device=self.device,
    )
    return residual, tokens, shortformer_pos_embed, attention_mask

load_and_process_state_dict(*args, **kwargs)

Unsupported HookedTransformer weight-processing helper.

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

def load_and_process_state_dict(self, *args: Any, **kwargs: Any) -> None:
    """Unsupported HookedTransformer weight-processing helper."""
    _ = args, kwargs
    raise NotImplementedError(
        "SafeLens' dependency-free wrapper does not load or process "
        "HookedTransformer-format state dictionaries."
    )

load_sample_training_dataset(*args, **kwargs)

Store an empty sample dataset placeholder for TL notebook compatibility.

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

def load_sample_training_dataset(self, *args: Any, **kwargs: Any) -> None:
    """Store an empty sample dataset placeholder for TL notebook compatibility."""
    _ = args, kwargs
    self.dataset: list[Any] = []

loss_fn(logits, tokens, attention_mask=None, per_token=False)

Compute TransformerLens-style next-token cross-entropy loss.

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

def loss_fn(
    self,
    logits: Any,
    tokens: Any,
    attention_mask: Any | None = None,
    per_token: bool = False,
) -> Any:
    """Compute TransformerLens-style next-token cross-entropy loss."""
    return _extract_or_compute_loss(
        {"logits": logits},
        _loss_model_inputs(tokens, attention_mask),
        logits=logits,
        loss_per_token=per_token,
    )

move_model_modules_to_device()

Compatibility no-op unless a wrapper device has been set.

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

def move_model_modules_to_device(self) -> HuggingFaceModelWrapper:
    """Compatibility no-op unless a wrapper device has been set."""
    if self.device is not None:
        return self.to(self.device)
    return self

named_parameters(*args, **kwargs)

Proxy named parameter iteration to the wrapped model when available.

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

def named_parameters(self, *args: Any, **kwargs: Any) -> Any:
    """Proxy named parameter iteration to the wrapped model when available."""
    model = self._require_model()
    named_parameters = getattr(model, "named_parameters", None)
    if callable(named_parameters):
        return named_parameters(*args, **kwargs)
    return iter(())

parameters(*args, **kwargs)

Proxy parameter iteration to the wrapped model when available.

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

def parameters(self, *args: Any, **kwargs: Any) -> Any:
    """Proxy parameter iteration to the wrapped model when available."""
    model = self._require_model()
    parameters = getattr(model, "parameters", None)
    if callable(parameters):
        return parameters(*args, **kwargs)
    return iter(())

process_weights_(fold_ln=True, center_writing_weights=True, center_unembed=True, fold_value_biases=True, refactor_factored_attn_matrices=False, *args, **kwargs)

Run supported in-place TransformerLens-style weight processing passes.

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

def process_weights_(
    self,
    fold_ln: bool = True,
    center_writing_weights: bool = True,
    center_unembed: bool = True,
    fold_value_biases: bool = True,
    refactor_factored_attn_matrices: bool = False,
    *args: Any,
    **kwargs: Any,
) -> HuggingFaceModelWrapper:
    """Run supported in-place TransformerLens-style weight processing passes."""
    _ = args
    if kwargs:
        unexpected = ", ".join(sorted(kwargs))
        raise TypeError(f"Unexpected process_weights_ keyword argument(s): {unexpected}.")
    model = self._require_model()
    has_shortformer_positional_embeddings = (
        _infer_positional_embedding_type(model) == "shortformer"
    )
    has_olmo2_post_norm = _is_olmo2_post_norm_model(model)
    center_writing_weights_by_default = _can_center_writing_weights_by_default(model)
    should_center_writing_weights = (
        center_writing_weights
        and not has_shortformer_positional_embeddings
        and not has_olmo2_post_norm
        and center_writing_weights_by_default
    )
    if fold_ln and not has_shortformer_positional_embeddings and not has_olmo2_post_norm:
        self.fold_layer_norm()
    if should_center_writing_weights:
        self.center_writing_weights()
    if (
        center_unembed
        and not has_shortformer_positional_embeddings
        and not _has_output_logits_soft_cap(model)
        and self._transformer_lens_model_kind() not in {"encoder", "audio_encoder"}
    ):
        self.center_unembed()
    if fold_value_biases:
        self.fold_value_biases()
        if should_center_writing_weights:
            self._center_attention_output_biases_after_value_folding()
    if refactor_factored_attn_matrices:
        self.refactor_factored_attn_matrices()
    return self

refactor_factored_attn_matrices(*args, **kwargs)

Refactor native attention QK/OV matrices into SVD-based factorizations.

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

def refactor_factored_attn_matrices(
    self,
    *args: Any,
    **kwargs: Any,
) -> HuggingFaceModelWrapper:
    """Refactor native attention QK/OV matrices into SVD-based factorizations."""
    _ = args, kwargs
    model = self._require_model()
    if _uses_rotary_embeddings(model):
        raise AssertionError(
            "You can't refactor the QK circuit when using rotary embeddings "
            "(the QK matrix depends on query/key position)."
        )
    refactored = False
    native_w_q = getattr(model, "W_Q", None)
    native_w_k = getattr(model, "W_K", None)
    if native_w_q is not None and native_w_k is not None:
        native_b_q = getattr(model, "b_Q", None)
        native_b_k = getattr(model, "b_K", None)
        refactored_w_q, refactored_b_q, refactored_w_k, refactored_b_k = _refactor_qk_matrices(
            native_w_q,
            native_b_q,
            native_w_k,
            native_b_k,
        )
        model.W_Q = refactored_w_q
        model.W_K = refactored_w_k
        if native_b_q is not None:
            model.b_Q = refactored_b_q
        if native_b_k is not None:
            model.b_K = refactored_b_k
        refactored = True
    native_w_v = getattr(model, "W_V", None)
    native_w_o = getattr(model, "W_O", None)
    if native_w_v is not None and native_w_o is not None:
        native_b_v = getattr(model, "b_V", None)
        if native_b_v is not None:
            native_b_o = getattr(model, "b_O", None)
            folded_bias, zero_b_v = _fold_value_bias_tensor_like(
                native_b_v,
                native_w_o,
                native_b_o,
                target_heads=self.cfg.n_heads,
            )
            model.b_O = folded_bias
            model.b_V = zero_b_v
        refactored_w_v, refactored_w_o = _refactor_ov_matrices(native_w_v, native_w_o)
        model.W_V = refactored_w_v
        model.W_O = refactored_w_o
        refactored = True
    adapter = architecture_adapter_for_model(model, model_name=self.name)
    for layer in range(_infer_model_layers(model)):
        if _refactor_split_attention_layer(model, adapter=adapter, layer=layer):
            refactored = True
        elif _refactor_joint_qkv_attention_layer(model, adapter=adapter, layer=layer):
            refactored = True
    if not refactored:
        raise NotImplementedError(
            "refactor_factored_attn_matrices currently requires native TL-shaped "
            "W_Q/W_K or W_V/W_O attributes, HF split q/k/v/o projection modules, "
            "or joint QKV projection modules with matching query/key/value head counts. "
            "Grouped-query refactor writeback for this experimental TL pass is not "
            "implemented yet."
        )
    return self

reset_hooks(*, clear_contexts=True, direction=None, dir=None, including_permanent=False, level=None)

Remove wrapper-managed hooks using TransformerLens reset semantics.

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

def reset_hooks(
    self,
    *,
    clear_contexts: bool = True,
    direction: Any = None,
    dir: Any = None,
    including_permanent: bool = False,
    level: int | None = None,
) -> None:
    """Remove wrapper-managed hooks using TransformerLens reset semantics."""
    hook_direction = _normalize_hook_direction(direction or dir or "both")
    removed_handles: list[_ManagedWrapperHookHandle] = []
    for handle in reversed(list(self._hooks)):
        if handle.is_permanent and not including_permanent:
            continue
        if level is not None and handle.level != level:
            continue
        if not _managed_handle_matches_direction(handle, hook_direction):
            continue
        removed_handles.append(handle)
        handle.remove()
    if clear_contexts:
        self._clear_hook_contexts_for_handles(removed_handles)
        self.clear_contexts()
    self.is_caching = _wrapper_has_active_permanent_cache_hooks(self._hooks)

run_with_hooks(batch, *model_args, fwd_hooks=(), bwd_hooks=(), prepend=False, reset_hooks_end=True, clear_contexts=False, return_type=_DEFAULT_RETURN_TYPE, loss_per_token=_DEFAULT_LOSS_PER_TOKEN, **forward_kwargs)

Run one forward pass with temporary hooks, mirroring TransformerLens.

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

def run_with_hooks(
    self,
    batch: Any,
    *model_args: Any,
    fwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
    bwd_hooks: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]] = (),
    prepend: bool = False,
    reset_hooks_end: bool = True,
    clear_contexts: bool = False,
    return_type: str | None | object = _DEFAULT_RETURN_TYPE,
    loss_per_token: bool | object = _DEFAULT_LOSS_PER_TOKEN,
    **forward_kwargs: Any,
) -> Any:
    """Run one forward pass with temporary hooks, mirroring TransformerLens."""
    bwd_hook_specs = list(bwd_hooks)
    handles: list[Any] = []
    forward_options = _merge_transformer_lens_forward_positionals(
        model_args,
        forward_kwargs,
        return_type=return_type,
        loss_per_token=loss_per_token,
    )
    forward_return_type = forward_options.pop("return_type", _DEFAULT_RETURN_TYPE)
    forward_loss_per_token = bool(forward_options.pop("loss_per_token", False))
    try:
        for layer, hook_fn in self._expand_hook_specs(fwd_hooks):
            handles.append(self._add_managed_hook(layer, hook_fn, prepend=prepend))
        for layer, hook_fn in self._expand_hook_specs(bwd_hook_specs):
            handles.append(self._add_managed_backward_hook(layer, hook_fn, prepend=prepend))
        if forward_options:
            resolved_return_type = _resolve_return_type(batch, forward_return_type)
            with self._temporary_backward_hook_context(enabled=bool(bwd_hook_specs)):
                return self.forward(
                    batch,
                    return_type=resolved_return_type,
                    loss_per_token=forward_loss_per_token,
                    **forward_options,
                )
        return self._run_model_forward(
            batch,
            return_type=forward_return_type,
            loss_per_token=forward_loss_per_token,
            enable_grad=bool(bwd_hook_specs),
        )
    finally:
        if reset_hooks_end:
            _remove_wrapper_handles(handles)
            if clear_contexts:
                self._clear_hook_contexts_for_handles(handles)
                self.clear_contexts()

sample_datapoint(*args, **kwargs)

Return one datapoint from a previously loaded sample dataset.

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

def sample_datapoint(self, *args: Any, **kwargs: Any) -> Any:
    """Return one datapoint from a previously loaded sample dataset."""
    _ = args, kwargs
    dataset = getattr(self, "dataset", None)
    if not dataset:
        raise ValueError("No sample training dataset is loaded.")
    try:
        import random

        return random.choice(dataset)
    except Exception:
        return dataset[0]

set_tokenizer(tokenizer, default_padding_side=None)

Set the tokenizer used by TransformerLens-style token helpers.

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

def set_tokenizer(self, tokenizer: Any, default_padding_side: str | None = None) -> None:
    """Set the tokenizer used by TransformerLens-style token helpers."""
    if default_padding_side not in {"right", "left", None}:
        raise AssertionError(
            f"padding_side must be 'right', 'left' or None, got {default_padding_side!r}"
        )
    self.tokenizer = tokenizer
    if default_padding_side is not None:
        _set_attr_if_possible(tokenizer, "padding_side", default_padding_side)
    elif getattr(tokenizer, "padding_side", None) is None:
        _set_attr_if_possible(tokenizer, "padding_side", "right")

    eos_token = getattr(tokenizer, "eos_token", None)
    if eos_token is None:
        eos_token = "<|endoftext|>"
        _set_attr_if_possible(tokenizer, "eos_token", eos_token)
    if getattr(tokenizer, "pad_token", None) is None:
        _set_attr_if_possible(tokenizer, "pad_token", eos_token)
    if getattr(tokenizer, "bos_token", None) is None:
        _set_attr_if_possible(tokenizer, "bos_token", eos_token)
    eos_token_id = getattr(tokenizer, "eos_token_id", None)
    if getattr(tokenizer, "pad_token_id", None) is None and eos_token_id is not None:
        _set_attr_if_possible(tokenizer, "pad_token_id", eos_token_id)
    if getattr(tokenizer, "bos_token_id", None) is None and eos_token_id is not None:
        _set_attr_if_possible(tokenizer, "bos_token_id", eos_token_id)
    self.tokenizer_prepends_bos = _tokenizer_prepends_bos(tokenizer)

set_ungroup_grouped_query_attention(ungroup_grouped_query_attention)

Mirror TransformerLens' grouped-query attention compatibility switch.

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

def set_ungroup_grouped_query_attention(self, ungroup_grouped_query_attention: bool) -> None:
    """Mirror TransformerLens' grouped-query attention compatibility switch."""
    self._set_transformer_lens_runtime_flag(
        "ungroup_grouped_query_attention",
        ungroup_grouped_query_attention,
    )

set_use_attn_in(use_attn_in)

Mirror TransformerLens' runtime switch for attention input hooks.

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

def set_use_attn_in(self, use_attn_in: bool) -> None:
    """Mirror TransformerLens' runtime switch for attention input hooks."""
    if bool(use_attn_in) and self._has_grouped_query_attention():
        raise AssertionError(
            "Cannot use attn_in hooks when key/value heads are grouped; "
            "SafeLens exposes the config switch but does not synthesize missing HF hooks."
        )
    self._set_transformer_lens_runtime_flag("use_attn_in", use_attn_in)

set_use_attn_result(use_attn_result)

Mirror TransformerLens' runtime switch for per-head attention results.

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

def set_use_attn_result(self, use_attn_result: bool) -> None:
    """Mirror TransformerLens' runtime switch for per-head attention results."""
    self._set_transformer_lens_runtime_flag("use_attn_result", use_attn_result)

set_use_hook_mlp_in(use_hook_mlp_in)

Mirror TransformerLens' runtime switch for MLP input hooks.

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

def set_use_hook_mlp_in(self, use_hook_mlp_in: bool) -> None:
    """Mirror TransformerLens' runtime switch for MLP input hooks."""
    if bool(use_hook_mlp_in) and self.cfg.attn_only:
        raise AssertionError("Cannot use hook_mlp_in on attention-only models.")
    self._set_transformer_lens_runtime_flag("use_hook_mlp_in", use_hook_mlp_in)

set_use_split_qkv_input(use_split_qkv_input)

Mirror TransformerLens' runtime switch for split Q/K/V input hooks.

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

def set_use_split_qkv_input(self, use_split_qkv_input: bool) -> None:
    """Mirror TransformerLens' runtime switch for split Q/K/V input hooks."""
    self._set_transformer_lens_runtime_flag("use_split_qkv_input", use_split_qkv_input)

tl_parameters()

Return a TransformerLens-style parameter mapping for analysis helpers.

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

def tl_parameters(self) -> dict[str, Any]:
    """Return a TransformerLens-style parameter mapping for analysis helpers."""
    parameters: dict[str, Any] = {}
    try:
        parameters["embed.W_E"] = self.W_E
    except KeyError:
        pass
    try:
        parameters["unembed.W_U"] = self.W_U
        parameters["unembed.b_U"] = self.b_U
    except RuntimeError:
        pass
    try:
        parameters["pos_embed.W_pos"] = self.W_pos
    except KeyError:
        pass

    cfg = self.cfg
    n_layers = int(cfg.n_layers or 0)
    stacked: dict[str, Any] = {}
    parameter_getters: tuple[tuple[str, Callable[[], Any]], ...] = (
        ("blocks.{layer}.attn.W_Q", lambda: self.W_Q),
        ("blocks.{layer}.attn.W_K", lambda: self.W_K),
        ("blocks.{layer}.attn.W_V", lambda: self.W_V),
        ("blocks.{layer}.attn.W_O", lambda: self.W_O),
        ("blocks.{layer}.attn.b_Q", lambda: self.b_Q),
        ("blocks.{layer}.attn.b_K", lambda: self.b_K),
        ("blocks.{layer}.attn.b_V", lambda: self.b_V),
        ("blocks.{layer}.attn.b_O", lambda: self.b_O),
    )
    for template, getter in parameter_getters:
        try:
            stacked[template] = getter()
        except (KeyError, RuntimeError, ValueError):
            pass
    try:
        stacked.update(
            {
                "blocks.{layer}.mlp.W_in": self.W_in,
                "blocks.{layer}.mlp.W_out": self.W_out,
                "blocks.{layer}.mlp.b_in": self.b_in,
                "blocks.{layer}.mlp.b_out": self.b_out,
            }
        )
    except (KeyError, RuntimeError, ValueError):
        pass
    try:
        stacked["blocks.{layer}.mlp.W_gate"] = self.W_gate
    except (KeyError, RuntimeError, ValueError):
        pass
    for template, value in stacked.items():
        for layer_index in range(min(n_layers, _stack_first_dim(value))):
            parameters[template.format(layer=layer_index)] = value[layer_index]
    try:
        parameters.update(self._layer_norm_parameters())
    except (KeyError, RuntimeError, ValueError, NotImplementedError):
        pass
    return parameters

to(device_or_dtype, print_details=True)

Move the underlying model or change dtype, returning self like PyTorch/TL.

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

def to(self, device_or_dtype: Any, print_details: bool = True) -> HuggingFaceModelWrapper:
    """Move the underlying model or change dtype, returning ``self`` like PyTorch/TL."""
    _ = print_details
    target = _coerce_to_target(device_or_dtype)
    to_fn = getattr(self.model, "to", None)
    if callable(to_fn):
        to_fn(target)
    _update_wrapper_device_dtype(self, device_or_dtype, target)
    return self

to_single_str_token(token)

Return the string for a single token id.

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

def to_single_str_token(self, token: int) -> str:
    """Return the string for a single token id."""
    if not isinstance(token, int):
        raise TypeError(f"Expected an integer token id, got {type(token)!r}.")
    tokens = self.to_str_tokens([token])
    if len(tokens) != 1 or isinstance(tokens[0], list):
        raise ValueError(f"Expected token id {token!r} to decode to one string token.")
    return str(tokens[0])

to_single_token(text)

Return the single token id for text or raise when it tokenizes to multiple ids.

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

def to_single_token(self, text: str) -> int:
    """Return the single token id for text or raise when it tokenizes to multiple ids."""
    tokens = self.to_tokens(text, prepend_bos=False)
    shape = getattr(tokens, "shape", None)
    token_values = tokens.reshape(-1).tolist() if shape is not None else list(tokens)
    if len(token_values) != 1:
        raise ValueError(
            f"Expected {text!r} to tokenize to a single token, got {token_values}."
        )
    return int(token_values[0])

to_str_tokens(text_or_tokens, *, prepend_bos=None, padding_side=None)

Return per-token strings for text or token ids.

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

def to_str_tokens(
    self,
    text_or_tokens: str | Any,
    *,
    prepend_bos: bool | None = None,
    padding_side: str | None = None,
) -> list[str] | list[list[str]]:
    """Return per-token strings for text or token ids."""
    tokenizer = self._require_tokenizer_for_text("token string conversion")
    resolved_prepend_bos = _resolve_default_prepend_bos(self, prepend_bos)
    if (
        isinstance(text_or_tokens, Sequence)
        and not isinstance(text_or_tokens, str | bytes)
        and text_or_tokens
        and isinstance(
            text_or_tokens[0],
            Sequence | str,
        )
    ):
        return cast(
            list[list[str]],
            [
                self.to_str_tokens(
                    item,
                    prepend_bos=resolved_prepend_bos,
                    padding_side=padding_side,
                )
                for item in text_or_tokens
            ],
        )
    tokens = (
        self.to_tokens(
            text_or_tokens,
            prepend_bos=resolved_prepend_bos,
            padding_side=padding_side,
        )
        if isinstance(text_or_tokens, str)
        else text_or_tokens
    )
    shape = getattr(tokens, "shape", None)
    if shape is not None:
        shape_tuple = tuple(int(dim) for dim in shape)
        if len(shape_tuple) == 2 and shape_tuple[0] == 1:
            tokens = tokens[0]
        elif len(shape_tuple) > 1:
            raise ValueError(
                f"Invalid token shape for token string conversion: {shape_tuple!r}."
            )
    token_list = _single_token_list(tokens)
    batch_decode = getattr(tokenizer, "batch_decode", None)
    if callable(batch_decode):
        return [
            str(token)
            for token in _call_decode_with_supported_kwargs(
                batch_decode,
                [[token] for token in token_list],
                {"clean_up_tokenization_spaces": False},
            )
        ]
    convert = getattr(tokenizer, "convert_ids_to_tokens", None)
    if callable(convert):
        converted = convert(token_list)
        if isinstance(converted, str):
            return [converted]
        if isinstance(converted, Sequence) and not isinstance(converted, bytes):
            return [str(token) for token in converted]
    return [str(tokenizer.decode([token])) for token in token_list]

to_string(tokens, *, skip_special_tokens=False, clean_up_tokenization_spaces=False)

Decode token ids into text.

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

def to_string(
    self,
    tokens: Any,
    *,
    skip_special_tokens: bool = False,
    clean_up_tokenization_spaces: bool = False,
) -> str | list[str]:
    """Decode token ids into text."""
    tokenizer = self._require_tokenizer_for_text("decoding")
    shape = _shape_of_token_ids(tokens)
    if shape is not None and len(shape) > 2:
        raise ValueError(f"Invalid token shape for decoding: {shape!r}.")
    if shape is not None and len(shape) == 2:
        batch_decode = getattr(tokenizer, "batch_decode", None)
        if callable(batch_decode):
            return list(
                _call_decode_with_supported_kwargs(
                    batch_decode,
                    tokens,
                    {
                        "skip_special_tokens": skip_special_tokens,
                        "clean_up_tokenization_spaces": clean_up_tokenization_spaces,
                    },
                )
            )
        return [
            str(
                _call_decode_with_supported_kwargs(
                    tokenizer.decode,
                    row,
                    {
                        "skip_special_tokens": skip_special_tokens,
                        "clean_up_tokenization_spaces": clean_up_tokenization_spaces,
                    },
                )
            )
            for row in tokens
        ]
    if isinstance(tokens, int):
        tokens = [tokens]
    return str(
        _call_decode_with_supported_kwargs(
            tokenizer.decode,
            tokens,
            {
                "skip_special_tokens": skip_special_tokens,
                "clean_up_tokenization_spaces": clean_up_tokenization_spaces,
            },
        )
    )

to_tokens(text, *, prepend_bos=None, padding_side=None, move_to_device=True, truncate=True)

Tokenize text into a tensor, mirroring TransformerLens' convenience method.

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

def to_tokens(
    self,
    text: str | Sequence[str],
    *,
    prepend_bos: bool | None = None,
    padding_side: str | None = None,
    move_to_device: bool = True,
    truncate: bool = True,
) -> Any:
    """Tokenize text into a tensor, mirroring TransformerLens' convenience method."""
    tokenizer = self._require_tokenizer_for_text("tokenization")
    prepend_bos = _resolve_default_prepend_bos(self, prepend_bos)
    token_kwargs: dict[str, Any] = {
        "return_tensors": "pt",
        "add_special_tokens": False,
        "padding": not isinstance(text, str),
    }
    with (
        _temporary_tokenizer_padding_side(tokenizer, padding_side),
        _temporary_tokenizer_pad_token(tokenizer, enabled=bool(token_kwargs["padding"])),
    ):
        effective_pad_token_id = _tokenizer_effective_pad_token_id(tokenizer)
        effective_padding_side = str(getattr(tokenizer, "padding_side", "right"))
        if truncate:
            n_ctx = _tokenization_context_length(self.model, tokenizer)
            token_kwargs["truncation"] = True
            if n_ctx is not None:
                max_length = int(n_ctx) - (1 if prepend_bos else 0)
                if max_length > 0:
                    token_kwargs["max_length"] = max_length
        tokenized = _call_tokenizer_with_supported_kwargs(
            tokenizer,
            text,
            token_kwargs,
        )
    tokens = tokenized["input_ids"] if isinstance(tokenized, dict) else tokenized.input_ids
    if prepend_bos:
        tokens = _prepend_bos_token(
            tokens,
            tokenizer,
            pad_token_id=effective_pad_token_id,
            padding_side=effective_padding_side,
        )
    if move_to_device and self.device is not None:
        tokens = tokens.to(self.device)
    return tokens

tokens_to_residual_directions(tokens)

Map token ids to unembedding residual directions.

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

def tokens_to_residual_directions(self, tokens: Any) -> Any:
    """Map token ids to unembedding residual directions."""
    weight = self._output_embedding_weight()
    if isinstance(tokens, str):
        tokens = self.to_single_token(tokens)
    elif isinstance(tokens, int):
        pass
    else:
        numel = getattr(tokens, "numel", None)
        item = getattr(tokens, "item", None)
        if callable(numel) and callable(item):
            try:
                if int(cast(Any, numel())) == 1:
                    tokens = int(cast(Any, item()))
            except Exception:
                pass
    try:
        if isinstance(weight, Sequence) and not isinstance(weight, str | bytes):
            return _gather_sequence_residual_directions(weight, tokens)
        tokens = _coerce_tokens_for_weight_index(weight, tokens)
        return weight[tokens]
    except Exception as exc:
        raise RuntimeError(
            "Could not index residual directions with the provided tokens."
        ) from exc

LocalModelWrapper

Bases: HuggingFaceModelWrapper

Transformers-compatible local directory wrapper with no provider download.

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

class LocalModelWrapper(HuggingFaceModelWrapper):
    """Transformers-compatible local directory wrapper with no provider download."""

ModelScopeModelWrapper

Bases: HuggingFaceModelWrapper

ModelScope-backed wrapper that downloads a snapshot, then loads it with Transformers.

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

class ModelScopeModelWrapper(HuggingFaceModelWrapper):
    """ModelScope-backed wrapper that downloads a snapshot, then loads it with Transformers."""

    def __init__(
        self,
        name: str,
        dtype: str = "float32",
        device: str | None = None,
        revision: str | None = None,
        cache_dir: str | None = None,
        local_dir: str | None = None,
        trust_remote_code: bool = False,
        load_kwargs: dict[str, Any] | None = None,
        tokenizer_kwargs: dict[str, Any] | None = None,
        modelscope_kwargs: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(
            name=name,
            dtype=dtype,
            device=device,
            revision=revision,
            cache_dir=cache_dir,
            trust_remote_code=trust_remote_code,
            load_kwargs=load_kwargs,
            tokenizer_kwargs=tokenizer_kwargs,
        )
        self.local_dir = local_dir
        self.modelscope_kwargs = modelscope_kwargs or {}

    def _resolve_pretrained_path(self) -> str:
        try:
            from modelscope import snapshot_download
        except ImportError as exc:
            raise ImportError(
                "ModelScopeModelWrapper requires ModelScope dependencies. "
                "Install them with `pip install -e '.[modelscope]'`."
            ) from exc

        kwargs = dict(self.modelscope_kwargs)
        if self.revision is not None:
            kwargs["revision"] = self.revision
        if self.cache_dir is not None:
            kwargs["cache_dir"] = self.cache_dir
        if self.local_dir is not None:
            kwargs["local_dir"] = self.local_dir
        return str(snapshot_download(model_id=self.name, **kwargs))

    def _pretrained_kwargs(self) -> dict[str, Any]:
        return {}

Qwen3DenseModelWrapper

Bases: HuggingFaceModelWrapper

Qwen3 dense wrapper exposing SafeLens component hook names.

Supported model family: Qwen3 dense language models up to 35B parameters (0.6B, 1.7B, 4B, 8B, 14B, and 32B). MoE variants such as 30B-A3B and non-dense/VL/Coder variants are intentionally rejected.

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

class Qwen3DenseModelWrapper(HuggingFaceModelWrapper):
    """Qwen3 dense wrapper exposing SafeLens component hook names.

    Supported model family: Qwen3 dense language models up to 35B parameters
    (`0.6B`, `1.7B`, `4B`, `8B`, `14B`, and `32B`). MoE variants such as
    `30B-A3B` and non-dense/VL/Coder variants are intentionally rejected.
    """

    def load_model(self) -> Any:
        validate_qwen3_dense_model_name(self.name)
        model = super().load_model()
        self._validate_loaded_qwen3_dense_model(model)
        return model

    def add_hook(
        self,
        layer: LayerRef | Callable[[str], bool],
        hook_fn: HookFn | None = None,
        *,
        hook: HookFn | None = None,
        dir: str = "fwd",
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
    ) -> Any:
        """Register a TransformerLens-style hook on a Qwen3 component or module."""
        resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
        if dir == "fwd":
            return self._add_managed_hook(
                layer,
                resolved_hook,
                is_permanent=is_permanent,
                level=level,
                prepend=prepend,
            )
        if dir == "bwd":
            return self._add_managed_backward_hook(
                layer,
                resolved_hook,
                is_permanent=is_permanent,
                level=level,
                prepend=prepend,
            )
        raise ValueError(f"Invalid hook direction {dir!r}.")

    def _register_hook(self, layer: LayerRef, hook_fn: HookFn, *, prepend: bool = False) -> Any:
        layer = self._resolve_hook_layer_ref(
            self._require_model(),
            layer,
            for_cache=False,
        )
        self.check_hooks_to_add(None, str(layer), hook_fn, dir="fwd", prepend=prepend)
        component_ref = parse_qwen3_component_ref(layer)
        if component_ref is None:
            return super()._register_hook(layer, hook_fn, prepend=prepend)
        layer_index, component = component_ref
        return self._register_qwen3_component_hook(
            layer_index,
            component,
            hook_fn,
            prepend=prepend,
        )

    def run_with_cache(
        self,
        batch: Any,
        *model_args: Any,
        layers: Sequence[LayerRef] | LayerRef | None = None,
        names_filter: NamesFilter = None,
        return_cache_object: bool | object = _DEFAULT_RETURN_CACHE_OBJECT,
        remove_batch_dim: bool = False,
        detach: bool = True,
        clone: bool = False,
        device: Any = None,
        pos_slice: Any = None,
        cache_all: bool | object = _DEFAULT_CACHE_ALL,
        return_type: str | None | object = _DEFAULT_RETURN_TYPE,
        loss_per_token: bool | object = _DEFAULT_LOSS_PER_TOKEN,
        incl_bwd: bool = False,
        reset_hooks_end: bool = True,
        clear_contexts: bool = False,
        **forward_kwargs: Any,
    ) -> tuple[Any, dict[str, Any] | ActivationCache]:
        model = self._require_model()
        cache = ActivationCache(model=self, has_batch_dim=not remove_batch_dim)
        temp_handles: list[Any] = []
        install_complete = False
        layers, forward_args = _split_run_with_cache_positionals(model_args, layers=layers)
        forward_options = _merge_transformer_lens_forward_positionals(
            forward_args,
            forward_kwargs,
            return_type=return_type,
            loss_per_token=loss_per_token,
        )
        forward_return_type = forward_options.pop("return_type", _DEFAULT_RETURN_TYPE)
        forward_loss_per_token = bool(forward_options.pop("loss_per_token", False))
        resolved_return_type = _resolve_return_type(batch, forward_return_type)
        resolved_cache_all = _resolve_cache_all(batch, layers, names_filter, cache_all)
        resolved_return_cache_object = _resolve_return_cache_object(batch, return_cache_object)

        try:
            self.is_caching = True
            for layer in self._cache_layers(
                model,
                layers,
                names_filter,
                cache_all=resolved_cache_all,
            ):
                temp_handles.append(
                    self._register_cache_hook(
                        model,
                        layer,
                        cache,
                        detach=detach,
                        clone=clone,
                        device=device,
                        pos_slice=pos_slice,
                        remove_batch_dim=remove_batch_dim,
                    )
                )
                if incl_bwd:
                    cache_name = self._cache_name_for_layer(model, layer, for_cache=True)
                    temp_handles.append(
                        self._register_backward_hook(
                            layer,
                            make_cache_hook(
                                cache,
                                f"{cache_name}_grad",
                                detach=detach,
                                clone=clone,
                                device=device,
                                pos_slice=pos_slice,
                                remove_batch_dim=remove_batch_dim,
                            ),
                            prepend=False,
                        )
                    )

            install_complete = True
            if forward_options:
                with self._temporary_backward_hook_context(enabled=incl_bwd):
                    output = self.forward(
                        batch,
                        return_type=resolved_return_type,
                        loss_per_token=forward_loss_per_token,
                        **forward_options,
                    )
            else:
                output = self._run_model_forward(
                    batch,
                    return_type=resolved_return_type,
                    loss_per_token=forward_loss_per_token,
                    enable_grad=incl_bwd,
                )
            if incl_bwd:
                _backward_scalar_output(output)
        finally:
            if reset_hooks_end or not install_complete:
                _remove_wrapper_handles(temp_handles)
                if clear_contexts:
                    self._clear_hook_contexts_for_handles(temp_handles)
                    self.clear_contexts()
            else:
                _keep_wrapper_handles(self, temp_handles, is_cache=True)
            self.is_caching = _wrapper_has_active_permanent_cache_hooks(self._hooks)

        return output, _format_cache_result(
            cache,
            model=self,
            return_cache_object=resolved_return_cache_object,
            remove_batch_dim=False,
        )

    def _register_cache_hook(
        self,
        model: Any,
        layer: LayerRef,
        cache: ActivationCache,
        *,
        detach: bool,
        clone: bool,
        device: Any,
        pos_slice: Any,
        remove_batch_dim: bool,
        is_permanent: bool = False,
    ) -> Any:
        layer = self._resolve_hook_layer_ref(model, layer, for_cache=True)
        component_ref = parse_qwen3_component_ref(layer)
        if component_ref is None:
            return super()._register_cache_hook(
                model,
                layer,
                cache,
                detach=detach,
                clone=clone,
                device=device,
                pos_slice=pos_slice,
                remove_batch_dim=remove_batch_dim,
                is_permanent=is_permanent,
            )

        layer_index, component = component_ref
        cache_name = str(layer)
        cache_hook = make_cache_hook(
            cache,
            cache_name,
            detach=detach,
            clone=clone,
            device=device,
            pos_slice=pos_slice,
            remove_batch_dim=remove_batch_dim,
        )
        if (
            component in _QWEN3_ATTENTION_COMPONENTS
            or component in _QWEN3_CACHE_ONLY_COMPONENTS
            or component == "result"
        ):
            component_handle = self._try_register_component_cache_hook(
                model,
                layer,
                cache,
                detach=detach,
                clone=clone,
                device=device,
                pos_slice=pos_slice,
                remove_batch_dim=remove_batch_dim,
                is_permanent=is_permanent,
            )
            if component_handle is not None:
                return component_handle
            raise KeyError(f"Could not resolve Qwen3 attention component {layer!r}.")
        return self._register_qwen3_component_hook(layer_index, component, cache_hook)

    def _expand_hook_specs(
        self,
        hook_specs: Iterable[tuple[LayerRef | Callable[[str], bool], HookFn]],
    ) -> list[tuple[LayerRef, HookFn]]:
        specs = list(hook_specs)
        if not specs:
            return []
        model = self._require_model()
        adapter = architecture_adapter_for_model(model, model_name=self.name)
        names = _candidate_hook_names(model, adapter, for_cache=False)
        expanded: list[tuple[LayerRef, HookFn]] = []
        for layer_or_filter, hook_fn in specs:
            if callable(layer_or_filter) and not isinstance(layer_or_filter, str):
                matched = _filter_hook_names(names, layer_or_filter, adapter=adapter)
                expanded.extend((name, hook_fn) for name in matched)
                continue
            if isinstance(layer_or_filter, str):
                matched = _filter_hook_names(names, layer_or_filter, adapter=adapter)
                if matched:
                    expanded.extend((name, hook_fn) for name in matched)
                    continue
                expanded.append((layer_or_filter, hook_fn))
            else:
                expanded.append((layer_or_filter, hook_fn))
        return expanded

    def _register_qwen3_component_hook(
        self,
        layer_index: int,
        component: str,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        if component in _QWEN3_ATTENTION_COMPONENTS or component == "result":
            handle = self._try_register_component_hook(
                self._require_model(),
                f"layer_{layer_index}.{component}",
                hook_fn,
                prepend=prepend,
            )
            if handle is None:
                raise KeyError(f"Could not resolve Qwen3 attention component {component!r}.")
            return handle
        if (
            component not in _QWEN3_PATCHABLE_COMPONENTS
            and component not in _QWEN3_EXPLICIT_COMPONENTS
        ):
            supported = ", ".join(qwen3_supported_hook_components(include_attention=True))
            examples = ", ".join(_QWEN3_COMPONENT_EXAMPLES[:4])
            raise KeyError(
                f"Unsupported Qwen3 dense component {component!r}. "
                f"Supported components: {supported}. Example hook names: {examples}."
            )

        qwen_layer = self._qwen3_layer(layer_index)
        if component == "resid_pre":
            return self._register_input_hook(
                qwen_layer,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "resid_mid":
            return self._register_input_hook(
                qwen_layer.post_attention_layernorm,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "resid_post":
            return self._register_first_output_hook(
                qwen_layer,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "attn_out":
            return self._register_first_output_hook(
                qwen_layer.self_attn,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "mlp_out":
            return self._register_tensor_output_hook(
                qwen_layer.mlp,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "pre":
            return self._register_tensor_output_hook(
                qwen_layer.mlp.gate_proj,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "pre_linear":
            return self._register_tensor_output_hook(
                qwen_layer.mlp.up_proj,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component == "post":
            return self._register_input_hook(
                qwen_layer.mlp.down_proj,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        if component in {"q", "k", "v"}:
            projection = getattr(qwen_layer.self_attn, f"{component}_proj")
            return self._register_head_projection_hook(
                projection,
                layer_index,
                component,
                hook_fn,
                prepend=prepend,
            )
        return self._register_z_hook(
            qwen_layer.self_attn.o_proj,
            layer_index,
            hook_fn,
            prepend=prepend,
        )

    def _register_input_hook(
        self,
        module: Any,
        layer_index: int,
        component: str,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        hook_context = _qwen3_hook_context(layer_index, component)

        def pre_hook(_module: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...] | None:
            if not inputs:
                return None
            patched = _call_qwen3_component_hook(
                hook_fn,
                activation=inputs[0],
                layer=layer_index,
                component=component,
                hook_context=hook_context,
            )
            if patched is None:
                return None
            return (patched, *inputs[1:])

        return _register_module_forward_pre_hook(module, pre_hook, prepend=prepend)

    def _register_first_output_hook(
        self,
        module: Any,
        layer_index: int,
        component: str,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        hook_context = _qwen3_hook_context(layer_index, component)

        def forward_hook(_module: Any, _inputs: Any, output: Any) -> Any:
            activation = _first_output(output)
            patched = _call_qwen3_component_hook(
                hook_fn,
                activation=activation,
                layer=layer_index,
                component=component,
                hook_context=hook_context,
            )
            if patched is None:
                return None
            return _replace_first_output(output, patched)

        return _register_module_forward_hook(module, forward_hook, prepend=prepend)

    def _register_tensor_output_hook(
        self,
        module: Any,
        layer_index: int,
        component: str,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        hook_context = _qwen3_hook_context(layer_index, component)

        def forward_hook(_module: Any, _inputs: Any, output: Any) -> Any:
            patched = _call_qwen3_component_hook(
                hook_fn,
                activation=output,
                layer=layer_index,
                component=component,
                hook_context=hook_context,
            )
            return None if patched is None else patched

        return _register_module_forward_hook(module, forward_hook, prepend=prepend)

    def _register_head_projection_hook(
        self,
        module: Any,
        layer_index: int,
        component: str,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        hook_context = _qwen3_hook_context(layer_index, component)

        def forward_hook(_module: Any, _inputs: Any, output: Any) -> Any:
            n_heads = self._heads_for_component(component)
            activation = _split_qwen3_heads(output, n_heads)
            patched = _call_qwen3_component_hook(
                hook_fn,
                activation=activation,
                layer=layer_index,
                component=component,
                hook_context=hook_context,
            )
            if patched is None:
                return None
            return _merge_qwen3_heads(patched, output)

        return _register_module_forward_hook(module, forward_hook, prepend=prepend)

    def _register_z_hook(
        self,
        module: Any,
        layer_index: int,
        hook_fn: HookFn,
        *,
        prepend: bool = False,
    ) -> Any:
        hook_context = _qwen3_hook_context(layer_index, "z")

        def pre_hook(_module: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...] | None:
            if not inputs:
                return None
            activation = _split_qwen3_heads(inputs[0], self._heads_for_component("z"))
            patched = _call_qwen3_component_hook(
                hook_fn,
                activation=activation,
                layer=layer_index,
                component="z",
                hook_context=hook_context,
            )
            if patched is None:
                return None
            return (_merge_qwen3_heads(patched, inputs[0]), *inputs[1:])

        return _register_module_forward_pre_hook(module, pre_hook, prepend=prepend)

    def _qwen3_layer(self, layer_index: int) -> Any:
        layers = _get_qwen3_layers(self._require_model())
        try:
            return layers[layer_index]
        except IndexError as exc:
            raise KeyError(f"Unknown Qwen3 dense layer index {layer_index}.") from exc

    def _heads_for_component(self, component: str) -> int:
        config = getattr(self._require_model(), "config", None)
        if component in {"k", "v"}:
            n_key_value_heads = _config_attr(config, "num_key_value_heads")
            if n_key_value_heads is not None:
                return int(n_key_value_heads)
        n_heads = _config_attr(config, "num_attention_heads")
        if n_heads is None:
            raise ValueError("Qwen3 config does not expose num_attention_heads.")
        return int(n_heads)

    @staticmethod
    def _validate_loaded_qwen3_dense_model(model: Any) -> None:
        config = getattr(model, "config", None)
        model_type = str(_config_attr(config, "model_type", "")).lower()
        if model_type not in {"qwen3", ""}:
            raise ValueError(f"Expected a Qwen3 dense model, got model_type={model_type!r}.")
        if _config_attr(config, "num_experts") is not None:
            raise ValueError("Qwen3 MoE models are not supported by Qwen3DenseModelWrapper.")

add_hook(layer, hook_fn=None, *, hook=None, dir='fwd', is_permanent=False, level=None, prepend=False)

Register a TransformerLens-style hook on a Qwen3 component or module.

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

def add_hook(
    self,
    layer: LayerRef | Callable[[str], bool],
    hook_fn: HookFn | None = None,
    *,
    hook: HookFn | None = None,
    dir: str = "fwd",
    is_permanent: bool = False,
    level: int | None = None,
    prepend: bool = False,
) -> Any:
    """Register a TransformerLens-style hook on a Qwen3 component or module."""
    resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
    if dir == "fwd":
        return self._add_managed_hook(
            layer,
            resolved_hook,
            is_permanent=is_permanent,
            level=level,
            prepend=prepend,
        )
    if dir == "bwd":
        return self._add_managed_backward_hook(
            layer,
            resolved_hook,
            is_permanent=is_permanent,
            level=level,
            prepend=prepend,
        )
    raise ValueError(f"Invalid hook direction {dir!r}.")

TransformerLensCompatibleModelWrapper

Bases: HuggingFaceModelWrapper

Independent Transformers wrapper for TransformerLens-compatible model IDs.

The compatibility table mirrors TransformerLens' public support matrix, but this class never imports or delegates to TransformerLens. Decoder, encoder-decoder, encoder, and audio-encoder families are loaded through the closest Transformers auto class.

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

class TransformerLensCompatibleModelWrapper(HuggingFaceModelWrapper):
    """Independent Transformers wrapper for TransformerLens-compatible model IDs.

    The compatibility table mirrors TransformerLens' public support matrix, but
    this class never imports or delegates to TransformerLens. Decoder,
    encoder-decoder, encoder, and audio-encoder families are loaded through the
    closest Transformers auto class.
    """

    def __init__(
        self,
        *args: Any,
        process_weights_kwargs: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self._process_weights_kwargs = (
            dict(process_weights_kwargs)
            if process_weights_kwargs is not None
            else _default_transformer_lens_process_kwargs()
        )
        self._weights_processed = False

    @classmethod
    def from_pretrained(
        cls,
        model_name: str,
        **kwargs: Any,
    ) -> TransformerLensCompatibleModelWrapper:
        """Build and load a dependency-free TransformerLens-compatible wrapper."""
        process_kwargs = _extract_transformer_lens_process_kwargs(kwargs)
        wrapper = cls(
            name=model_name,
            dtype=str(kwargs.pop("dtype", "float32")),
            device=kwargs.pop("device", None),
            revision=kwargs.pop("revision", None),
            cache_dir=kwargs.pop("cache_dir", None),
            trust_remote_code=bool(kwargs.pop("trust_remote_code", False)),
            load_kwargs=dict(kwargs.pop("load_kwargs", {})),
            tokenizer_kwargs=dict(kwargs.pop("tokenizer_kwargs", {})),
            pretrained_path=kwargs.pop("pretrained_path", None),
            process_weights_kwargs=process_kwargs,
        )
        wrapper.load_kwargs.update(_filter_transformer_lens_load_kwargs(kwargs))
        wrapper.load_model()
        return wrapper

    @classmethod
    def from_pretrained_no_processing(
        cls,
        model_name: str,
        **kwargs: Any,
    ) -> TransformerLensCompatibleModelWrapper:
        """Build a TransformerLens-compatible wrapper without TL weight processing."""
        kwargs.setdefault("fold_ln", False)
        kwargs.setdefault("center_writing_weights", False)
        kwargs.setdefault("center_unembed", False)
        kwargs.setdefault("fold_value_biases", False)
        kwargs.setdefault("refactor_factored_attn_matrices", False)
        return cls.from_pretrained(model_name, **kwargs)

    def load_model(self) -> Any:
        if not self._is_supported_transformer_lens_target():
            raise ValueError(
                f"Model {self.name!r} is not in SafeLens' vendored TransformerLens-compatible "
                "support table. Use source='huggingface' for generic Transformers loading, "
                "or source='local' for a local model directory."
            )
        pretrained_path = self._resolve_pretrained_path()
        self._raise_if_native_transformer_lens_checkpoint(pretrained_path)
        try:
            import torch
            from transformers import (
                AutoFeatureExtractor,
                AutoModel,
                AutoModelForCausalLM,
                AutoModelForSeq2SeqLM,
                AutoProcessor,
                AutoTokenizer,
            )
        except ImportError as exc:
            raise ImportError(
                "TransformerLensCompatibleModelWrapper requires SafeLens model "
                "dependencies. Install them with `pip install -e '.[models]'`."
            ) from exc

        dtype_map = {
            "float16": torch.float16,
            "bfloat16": torch.bfloat16,
            "float32": torch.float32,
            "auto": "auto",
        }
        torch_dtype = dtype_map.get(self.dtype, self.dtype)
        pretrained_kwargs = self._pretrained_kwargs()
        kind = self._transformer_lens_model_kind(
            pretrained_path=pretrained_path,
            probe_pretrained_path=True,
        )

        if kind == "audio_encoder":
            self.tokenizer = self._load_audio_processor(
                (AutoProcessor, AutoFeatureExtractor),
                pretrained_path,
                pretrained_kwargs,
            )
            self.model = AutoModel.from_pretrained(
                pretrained_path,
                dtype=torch_dtype,
                trust_remote_code=self.trust_remote_code,
                **pretrained_kwargs,
                **self.load_kwargs,
            )
        else:
            self.tokenizer = self._load_text_tokenizer(
                AutoTokenizer,
                pretrained_path,
                pretrained_kwargs,
            )
            model_cls: Any
            if kind == "encoder_decoder":
                model_cls = AutoModelForSeq2SeqLM
            elif kind == "encoder":
                model_cls = AutoModel
            else:
                model_cls = AutoModelForCausalLM
            self.model = model_cls.from_pretrained(
                pretrained_path,
                dtype=torch_dtype,
                trust_remote_code=self.trust_remote_code,
                **pretrained_kwargs,
                **self.load_kwargs,
            )

        if self.device is not None:
            self.model.to(self.device)
        self.model.eval()
        self._weights_processed = False
        self._process_loaded_weights()
        return self.model

    def _process_loaded_weights(self) -> None:
        if self._weights_processed:
            return
        self.process_weights_(**self._process_weights_kwargs)
        self._weights_processed = True

    def _transformer_lens_model_kind(
        self,
        *,
        pretrained_path: str | None = None,
        probe_pretrained_path: bool = False,
    ) -> str:
        model = getattr(self, "model", None)
        config = _core_model_config(_config_attr(model, "config"))
        model_type = _config_attr(config, "model_type")
        if isinstance(model_type, str) and model_type:
            return transformer_lens_model_kind(model_type)

        candidate = pretrained_path or self.pretrained_path or self.name
        if probe_pretrained_path and candidate:
            try:
                from transformers import AutoConfig
            except ImportError:
                pass
            else:
                try:
                    inferred_config = AutoConfig.from_pretrained(
                        candidate,
                        trust_remote_code=self.trust_remote_code,
                        **self._pretrained_kwargs(),
                    )
                except Exception:
                    pass
                else:
                    inferred_model_type = _config_attr(
                        _core_model_config(inferred_config),
                        "model_type",
                    )
                    if isinstance(inferred_model_type, str) and inferred_model_type:
                        return transformer_lens_model_kind(inferred_model_type)

        return transformer_lens_model_kind(candidate)

    def _resolve_pretrained_path(self) -> str:
        raw_path = self.pretrained_path or self.name
        if self._pretrained_path_is_local or _wrapper_looks_like_local_path(raw_path):
            return raw_path
        return resolve_transformer_lens_compatible_model_name(raw_path)

    def _raise_if_native_transformer_lens_checkpoint(self, pretrained_path: str) -> None:
        if (
            _wrapper_looks_like_local_path(self.name)
            or self._pretrained_path_is_local
            or self.pretrained_path is not None
            and _wrapper_looks_like_local_path(self.pretrained_path)
        ):
            return
        if not (
            is_transformer_lens_native_checkpoint(self.name)
            or is_transformer_lens_native_checkpoint(pretrained_path)
        ):
            return
        raise NotImplementedError(
            f"Model {self.name!r} resolves to TransformerLens-native checkpoint "
            f"{pretrained_path!r}. That repository stores HookedTransformer config "
            "and .pth weights rather than a HuggingFace Transformers model with a "
            "`model_type` config. SafeLens' transformer_lens source is dependency-free "
            "and currently loads only Transformers-compatible checkpoints; use a "
            "Transformers model id/local directory or convert the checkpoint to "
            "Transformers format before loading."
        )

    def _is_supported_transformer_lens_target(self) -> bool:
        if _wrapper_looks_like_local_path(self.name):
            return True
        if self._pretrained_path_is_local:
            return True
        if self.pretrained_path is not None and _wrapper_looks_like_local_path(
            self.pretrained_path
        ):
            return True
        return is_transformer_lens_supported_model_name(self.name)

    def _prepare_model_inputs(self, batch: Any) -> dict[str, Any]:
        batch = _normalize_model_batch(batch)
        kind = self._transformer_lens_model_kind()
        model_kwargs = _model_kwargs_without_tokenization(batch)
        if kind == "encoder_decoder" and "encoder_tokens" in batch and "decoder_tokens" in batch:
            return self._with_attention_flags(
                {
                    "input_ids": batch["encoder_tokens"],
                    "decoder_input_ids": batch["decoder_tokens"],
                    **model_kwargs,
                }
            )
        if kind == "encoder_decoder":
            prepared = super()._prepare_model_inputs(batch)
            prepared.update(model_kwargs)
            if not any(
                key in prepared for key in ("decoder_input_ids", "decoder_inputs_embeds", "labels")
            ):
                input_ids = prepared.get("input_ids")
                if input_ids is None:
                    raise ValueError(
                        "Encoder-decoder models require input_ids or explicit decoder inputs."
                    )
                config = getattr(self._require_model(), "config", None)
                decoder_start_token_id = batch.get(
                    "decoder_start_token_id",
                    _config_attr(config, "decoder_start_token_id"),
                )
                if decoder_start_token_id is None:
                    decoder_start_token_id = _config_attr(config, "pad_token_id")
                if decoder_start_token_id is None:
                    decoder_start_token_id = getattr(self.tokenizer, "pad_token_id", None)
                if decoder_start_token_id is None:
                    raise ValueError(
                        "Encoder-decoder models require decoder_input_ids when no "
                        "decoder_start_token_id or pad_token_id is available."
                    )
                prepared["decoder_input_ids"] = input_ids.new_full(
                    (input_ids.shape[0], 1),
                    int(decoder_start_token_id),
                )
            return self._with_attention_flags(prepared)
        if kind == "decoder":
            decoder_prepared = self._prepare_decoder_text_inputs(batch)
            if decoder_prepared is None:
                decoder_prepared = super()._prepare_model_inputs(batch)
            decoder_prepared.update(model_kwargs)
            return self._with_attention_flags(decoder_prepared)
        if kind != "audio_encoder":
            prepared = super()._prepare_model_inputs(batch)
            prepared.update(model_kwargs)
            return self._with_attention_flags(prepared)

        if self.tokenizer is None:
            return dict(batch)
        audio = _first_present(batch, ("audio", "wave", "raw_audio"))
        if audio is None:
            return model_kwargs
        processor_kwargs = dict(batch.get("processor_kwargs", {}))
        sampling_rate = batch.get("sampling_rate", 16000)
        processed = self.tokenizer(
            audio,
            sampling_rate=sampling_rate,
            return_tensors="pt",
            **processor_kwargs,
        )
        if self.device is not None:
            processed = processed.to(self.device)
        return self._with_attention_flags({**dict(processed), **model_kwargs})

    def _prepare_decoder_text_inputs(self, batch: Mapping[str, Any]) -> dict[str, Any] | None:
        return _prepare_text_inputs_with_to_tokens(self, batch)

    def _uses_decoder_text_input_semantics(self) -> bool:
        return self._transformer_lens_model_kind() == "decoder"

    def _uses_encoder_decoder_generation_semantics(self) -> bool:
        return self._transformer_lens_model_kind() == "encoder_decoder"

    def generate(self, prompt: Any = "", **generation_kwargs: Any) -> Any:
        kind = self._transformer_lens_model_kind()
        if kind in {"encoder", "audio_encoder"}:
            raise NotImplementedError(
                f"{kind} models do not expose autoregressive text generation through "
                "the independent SafeLens Transformers wrapper."
            )
        return super().generate(prompt, **generation_kwargs)

    def generate_stream(self, prompt: Any = "", **generation_kwargs: Any) -> Iterable[Any]:
        kind = self._transformer_lens_model_kind()
        if kind in {"encoder", "audio_encoder"}:
            raise NotImplementedError(
                f"{kind} models do not expose autoregressive text generation through "
                "the independent SafeLens Transformers wrapper."
            )
        return super().generate_stream(prompt, **generation_kwargs)

    def _load_audio_processor(
        self,
        processor_classes: tuple[Any, Any],
        pretrained_path: str,
        pretrained_kwargs: dict[str, Any],
    ) -> Any:
        last_error: Exception | None = None
        for processor_cls in processor_classes:
            try:
                return processor_cls.from_pretrained(
                    pretrained_path,
                    trust_remote_code=self.trust_remote_code,
                    **pretrained_kwargs,
                    **self.tokenizer_kwargs,
                )
            except Exception as exc:
                last_error = exc
        raise RuntimeError(
            f"Could not load an audio processor for {pretrained_path!r}."
        ) from last_error

from_pretrained(model_name, **kwargs) classmethod

Build and load a dependency-free TransformerLens-compatible wrapper.

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

@classmethod
def from_pretrained(
    cls,
    model_name: str,
    **kwargs: Any,
) -> TransformerLensCompatibleModelWrapper:
    """Build and load a dependency-free TransformerLens-compatible wrapper."""
    process_kwargs = _extract_transformer_lens_process_kwargs(kwargs)
    wrapper = cls(
        name=model_name,
        dtype=str(kwargs.pop("dtype", "float32")),
        device=kwargs.pop("device", None),
        revision=kwargs.pop("revision", None),
        cache_dir=kwargs.pop("cache_dir", None),
        trust_remote_code=bool(kwargs.pop("trust_remote_code", False)),
        load_kwargs=dict(kwargs.pop("load_kwargs", {})),
        tokenizer_kwargs=dict(kwargs.pop("tokenizer_kwargs", {})),
        pretrained_path=kwargs.pop("pretrained_path", None),
        process_weights_kwargs=process_kwargs,
    )
    wrapper.load_kwargs.update(_filter_transformer_lens_load_kwargs(kwargs))
    wrapper.load_model()
    return wrapper

from_pretrained_no_processing(model_name, **kwargs) classmethod

Build a TransformerLens-compatible wrapper without TL weight processing.

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

@classmethod
def from_pretrained_no_processing(
    cls,
    model_name: str,
    **kwargs: Any,
) -> TransformerLensCompatibleModelWrapper:
    """Build a TransformerLens-compatible wrapper without TL weight processing."""
    kwargs.setdefault("fold_ln", False)
    kwargs.setdefault("center_writing_weights", False)
    kwargs.setdefault("center_unembed", False)
    kwargs.setdefault("fold_value_biases", False)
    kwargs.setdefault("refactor_factored_attn_matrices", False)
    return cls.from_pretrained(model_name, **kwargs)

TransformerLensConfigView dataclass

Small read-only TransformerLens-style config view for wrapped HF models.

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

@dataclass(frozen=True)
class TransformerLensConfigView:
    """Small read-only TransformerLens-style config view for wrapped HF models."""

    model_name: str
    model_type: str | None = None
    n_layers: int | None = None
    n_heads: int | None = None
    n_key_value_heads: int | None = None
    d_model: int | None = None
    d_head: int | None = None
    d_vocab: int | None = None
    n_ctx: int | None = None
    d_mlp: int | None = None
    act_fn: str | None = None
    normalization_type: str | None = None
    positional_embedding_type: str | None = None
    device: str | None = None
    dtype: str | None = None
    original_architecture: str | None = None
    use_attn_result: bool = False
    use_split_qkv_input: bool = False
    use_hook_mlp_in: bool = False
    use_attn_in: bool = False
    ungroup_grouped_query_attention: bool = False
    attn_only: bool = False
    parallel_attn_mlp: bool = False
    rmsnorm_uses_offset: bool = False

    @property
    def n_params(self) -> None:
        return None

    def to_dict(self) -> dict[str, Any]:
        return dict(self.__dict__)

build_model_wrapper(config)

Build the configured model wrapper.

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

def build_model_wrapper(config: ModelLoadConfig) -> ModelWrapper:
    """Build the configured model wrapper."""
    register_builtin_model_adapters()
    if config.name.lower() in {"dummy", "mock", "none"}:
        return DummyModelWrapper(name=config.name)
    return get_model_adapter_registry().create(config)

is_supported_qwen3_dense_model_name(model_name)

Return whether a model name looks like a supported Qwen3 <=35B dense model.

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

def is_supported_qwen3_dense_model_name(model_name: str) -> bool:
    """Return whether a model name looks like a supported Qwen3 <=35B dense model."""
    lowered = model_name.lower()
    if "qwen3" not in lowered:
        return False
    if any(marker in lowered for marker in ("moe", "-a", "_a", "coder", "vl")):
        return False
    size_b = qwen3_dense_size_billion(model_name)
    return size_b is None or size_b <= _QWEN3_DENSE_MAX_PARAMS_B

make_final_norm_scale_cache_hook(cache, name, *, detach=True, clone=False, device=None, pos_slice=None, remove_batch_dim=False)

Create a cache hook for TransformerLens-style final norm scales.

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

def make_final_norm_scale_cache_hook(
    cache: ActivationCache,
    name: str,
    *,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
    pos_slice: Any = None,
    remove_batch_dim: bool = False,
) -> HookFn:
    """Create a cache hook for TransformerLens-style final norm scales."""
    normalized_pos_slice = _normalize_cache_pos_slice(pos_slice)

    def cache_hook(module: Any, inputs: tuple[Any, ...], output: Any) -> None:
        scale = _final_norm_scale_from_hook(module, inputs, output)
        if remove_batch_dim:
            scale = _remove_singleton_batch(scale)
            cache.has_batch_dim = False
        if normalized_pos_slice is not None:
            scale = _slice_tensor_like_dim(
                scale,
                normalized_pos_slice,
                dim=_cache_pos_dim_for_name(name, scale),
            )
        cache.store(name, scale, detach=detach, clone=clone, device=device)
        return None

    return cache_hook

parse_qwen3_component_ref(layer)

Parse SafeLens or TransformerLens-style Qwen3 component hook names.

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

def parse_qwen3_component_ref(layer: LayerRef) -> tuple[int, str] | None:
    """Parse SafeLens or TransformerLens-style Qwen3 component hook names."""
    if not isinstance(layer, str):
        return None

    safe_match = re.fullmatch(r"layer_(\d+)\.([a-z_]+)", layer)
    if safe_match is not None:
        return int(safe_match.group(1)), _normalize_qwen3_component(safe_match.group(2))

    block_match = re.fullmatch(r"blocks\.(\d+)\.(?:([a-z_]+)\.)?hook_([a-z_]+)", layer)
    if block_match is not None:
        layer_index = int(block_match.group(1))
        layer_type = block_match.group(2)
        component = _normalize_qwen3_component(block_match.group(3), layer_type=layer_type)
        return layer_index, component

    return None

qwen3_dense_size_billion(model_name)

Return the parsed Qwen3 model size in billions when present in the name.

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

def qwen3_dense_size_billion(model_name: str) -> float | None:
    """Return the parsed Qwen3 model size in billions when present in the name."""
    match = re.search(r"Qwen3[-_/](\d+(?:\.\d+)?)B", model_name, flags=re.IGNORECASE)
    if match is None:
        return None
    return float(match.group(1))

qwen3_hook_name_examples()

Return example SafeLens and TransformerLens-style Qwen3 hook names.

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

def qwen3_hook_name_examples() -> list[str]:
    """Return example SafeLens and TransformerLens-style Qwen3 hook names."""
    return list(_QWEN3_COMPONENT_EXAMPLES)

qwen3_supported_hook_components(*, include_attention=False, for_cache=False)

Return component names accepted by the Qwen3 dense adapter.

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

def qwen3_supported_hook_components(
    *,
    include_attention: bool = False,
    for_cache: bool = False,
) -> list[str]:
    """Return component names accepted by the Qwen3 dense adapter."""
    components = sorted(_QWEN3_PATCHABLE_COMPONENTS | _QWEN3_EXPLICIT_COMPONENTS)
    if include_attention:
        attention_components = set(_QWEN3_ATTENTION_COMPONENTS)
        if for_cache:
            attention_components.update(_QWEN3_CACHE_ONLY_COMPONENTS)
        components.extend(sorted(attention_components))
    return components

register_builtin_model_adapters()

Register SafeLens built-in model adapters once.

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

def register_builtin_model_adapters() -> None:
    """Register SafeLens built-in model adapters once."""
    global _MODEL_ADAPTERS_REGISTERED
    if _MODEL_ADAPTERS_REGISTERED:
        return

    registry = get_model_adapter_registry()
    registry.register(
        ModelAdapterSpec(
            name="dummy",
            display_name="Dummy",
            aliases=("mock", "none"),
            description="In-memory adapter for tests, CI, and architecture demos.",
            dependencies=(),
            model_name_patterns=("dummy", "mock", "none"),
            capabilities=ModelAdapterCapabilities(
                supported_hooks=("integer layer refs", "string layer refs"),
                supported_patches=("replace", "add"),
                supports_local_path=False,
                supports_remote_download=False,
                cache_policy="no external cache",
                notes=("Does not download or execute model code.",),
            ),
            build=_build_dummy_wrapper,
            inspect=_inspect_dummy_model,
            matches_model_name=lambda name: name.lower() in {"dummy", "mock", "none"},
            priority=100,
        )
    )
    registry.register(
        ModelAdapterSpec(
            name="qwen3_dense",
            display_name="Qwen3 Dense",
            aliases=("qwen3", "qwen3-dense"),
            description="Qwen3 dense <=35B adapter with component-level hooks.",
            dependencies=("torch>=2", "transformers>=5.8"),
            model_name_patterns=("Qwen/Qwen3-{0.6,1.7,4,8,14,32}B",),
            capabilities=ModelAdapterCapabilities(
                supported_hooks=tuple(qwen3_supported_hook_components(include_attention=True)),
                supported_patches=(
                    "resid_pre",
                    "resid_mid",
                    "resid_post",
                    "attn_out",
                    "mlp_out",
                    "pre",
                    "pre_linear",
                    "post",
                    "q",
                    "k",
                    "v",
                    "z",
                    "result",
                    "pattern",
                    "attn_scores",
                ),
                supports_attention_pattern=True,
                supports_attention_scores=True,
                supports_local_path=False,
                supports_remote_download=True,
                cache_policy="SafeLens cache_dir -> .cache/safelens/models/huggingface",
                notes=(
                    "Attention result hooks are implemented by deriving per-head "
                    "z @ W_O results and writing patched head-result deltas back "
                    "to the merged attention output.",
                    "Attention pattern and score hooks use eager softmax instrumentation; "
                    "flash or SDPA paths may need an eager attention implementation.",
                ),
            ),
            build=_build_qwen3_dense_wrapper,
            inspect=_inspect_qwen3_dense_model,
            matches_model_name=lambda name: (
                "qwen3" in name.lower() and not is_qwen_routed_moe_model_name(name)
            ),
            priority=200,
        )
    )
    registry.register(
        ModelAdapterSpec(
            name="transformer_lens",
            display_name="TransformerLens-Compatible Transformers",
            aliases=("transformerlens", "tl", "hooked_transformer"),
            description=(
                "Independent SafeLens adapter for model families mirrored from "
                "TransformerLens supported models and architecture bridge coverage."
            ),
            dependencies=("torch>=2", "transformers>=5.8"),
            model_name_patterns=(
                "gpt2",
                "EleutherAI/pythia-*",
                "meta-llama/*",
                "Qwen/Qwen*",
                "google/gemma-*",
                "google-bert/bert-*",
                "FacebookAI/roberta-*",
                "distilbert/distilbert-*",
                "google-t5/t5-*",
                "facebook/wav2vec2-*",
                "facebook/hubert-*",
                "state-spaces/mamba-*",
                "mistralai/Mamba-Codestral-*",
            ),
            capabilities=ModelAdapterCapabilities(
                supported_hooks=(
                    "integer layer refs",
                    "model.named_modules() names",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supported_patches=(
                    "module output replace",
                    "module output add",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supports_attention_pattern=True,
                supports_attention_scores=True,
                supports_local_path=True,
                supports_remote_download=True,
                cache_policy=(
                    "SafeLens cache_dir -> .cache/safelens/models/transformer_lens_compatible"
                ),
                notes=(
                    "No TransformerLens runtime dependency is used.",
                    "Decoder, encoder-decoder, encoder, and audio-encoder families "
                    "load through Transformers auto classes.",
                    "SafeLens architecture adapters map HF module paths to canonical "
                    "components for GPT-2, GPT-J, GPT-Neo, GPT-NeoX/Pythia, "
                    "BLOOM/Falcon, MPT, Phi, OPT, BERT/RoBERTa, DistilBERT, "
                    "T5, Wav2Vec2/Hubert, Mamba/Mamba2 SSMs, and LLaMA-like "
                    "decoder families.",
                    "Attention result hooks are implemented by deriving per-head "
                    "z @ W_O results and writing patched head-result deltas back "
                    "to the merged attention output.",
                    "Attention pattern and score hooks use eager softmax instrumentation; "
                    "flash or SDPA paths may need an eager attention implementation.",
                ),
            ),
            build=_build_transformer_lens_compatible_wrapper,
            inspect=_inspect_transformer_lens_compatible_model,
            matches_model_name=is_transformer_lens_supported_model_name,
            priority=90,
        )
    )
    registry.register(
        ModelAdapterSpec(
            name="huggingface",
            display_name="HuggingFace Transformers",
            aliases=("hf",),
            description="Generic Transformers causal language model adapter.",
            dependencies=("torch>=2", "transformers>=5.8"),
            model_name_patterns=("organization/model-name",),
            capabilities=ModelAdapterCapabilities(
                supported_hooks=(
                    "integer decoder layer refs",
                    "model.named_modules() names",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supported_patches=(
                    "module output replace",
                    "module output add",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supports_attention_pattern=True,
                supports_attention_scores=True,
                supports_local_path=False,
                supports_remote_download=True,
                cache_policy="SafeLens cache_dir -> .cache/safelens/models/huggingface",
                notes=(
                    "Component hooks use SafeLens architecture adapters when the "
                    "loaded Transformers architecture is recognized.",
                    "Attention result hooks are implemented by deriving per-head "
                    "z @ W_O results and writing patched head-result deltas back "
                    "to the merged attention output.",
                    "Attention pattern and score hooks use eager softmax instrumentation; "
                    "flash or SDPA paths may need an eager attention implementation.",
                ),
            ),
            build=_build_huggingface_wrapper,
            inspect=_inspect_huggingface_model,
            matches_model_name=lambda name: "/" in name and not name.startswith((".", "/", "~")),
            priority=10,
        )
    )
    registry.register(
        ModelAdapterSpec(
            name="modelscope",
            display_name="ModelScope",
            aliases=("ms",),
            description="ModelScope snapshot download followed by Transformers loading.",
            dependencies=("modelscope>=1.15", "torch>=2", "transformers>=5.8"),
            model_name_patterns=("namespace/model-name",),
            capabilities=ModelAdapterCapabilities(
                supported_hooks=(
                    "integer decoder layer refs",
                    "model.named_modules() names",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supported_patches=(
                    "module output replace",
                    "module output add",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supports_attention_pattern=True,
                supports_attention_scores=True,
                supports_local_path=False,
                supports_remote_download=True,
                cache_policy="SafeLens cache_dir -> .cache/safelens/models/modelscope",
                notes=(
                    "Use modelscope_kwargs for provider-specific snapshot filters.",
                    "Component hooks use SafeLens architecture adapters when the "
                    "loaded Transformers architecture is recognized.",
                    "Attention result hooks are implemented by deriving per-head "
                    "z @ W_O results and writing patched head-result deltas back "
                    "to the merged attention output.",
                    "Attention pattern and score hooks use eager softmax instrumentation; "
                    "flash or SDPA paths may need an eager attention implementation.",
                ),
            ),
            build=_build_modelscope_wrapper,
            inspect=_inspect_modelscope_model,
            matches_model_name=lambda _name: False,
            priority=5,
        )
    )
    registry.register(
        ModelAdapterSpec(
            name="local",
            display_name="Local Transformers Directory",
            aliases=(),
            description="Local Transformers-compatible model directory.",
            dependencies=("torch>=2", "transformers>=5.8"),
            model_name_patterns=("./models/local-causal-lm", "/abs/path/to/model"),
            capabilities=ModelAdapterCapabilities(
                supported_hooks=(
                    "integer decoder layer refs",
                    "model.named_modules() names",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supported_patches=(
                    "module output replace",
                    "module output add",
                    *supported_transformer_component_names(include_attention=True),
                ),
                supports_attention_pattern=True,
                supports_attention_scores=True,
                supports_local_path=True,
                supports_remote_download=False,
                cache_policy="No provider download; local_dir or name is used directly.",
                notes=(
                    "Keep local model paths and weights out of git.",
                    "Component hooks use SafeLens architecture adapters when the "
                    "loaded Transformers architecture is recognized.",
                    "Attention result hooks are implemented by deriving per-head "
                    "z @ W_O results and writing patched head-result deltas back "
                    "to the merged attention output.",
                    "Attention pattern and score hooks use eager softmax instrumentation; "
                    "flash or SDPA paths may need an eager attention implementation.",
                ),
            ),
            build=_build_local_wrapper,
            inspect=_inspect_local_model,
            matches_model_name=lambda name: name.startswith((".", "/", "~")),
            priority=50,
        )
    )
    _MODEL_ADAPTERS_REGISTERED = True

validate_qwen3_dense_model_name(model_name)

Reject known unsupported Qwen3 MoE or >35B model names before loading.

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

def validate_qwen3_dense_model_name(model_name: str) -> None:
    """Reject known unsupported Qwen3 MoE or >35B model names before loading."""
    lowered = model_name.lower()
    if "qwen3" not in lowered:
        return
    if any(marker in lowered for marker in ("moe", "-a", "_a", "coder", "vl")):
        raise ValueError(f"Unsupported Qwen3 non-dense model name: {model_name!r}.")
    size_b = qwen3_dense_size_billion(model_name)
    if size_b is not None and size_b > _QWEN3_DENSE_MAX_PARAMS_B:
        raise ValueError(
            f"Unsupported Qwen3 model size {size_b:g}B. "
            f"Only dense models <= {_QWEN3_DENSE_MAX_PARAMS_B:g}B are supported."
        )

validate_qwen3_hook_ref(layer)

Validate a static Qwen3 dense layer or component hook reference.

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

def validate_qwen3_hook_ref(layer: LayerRef) -> None:
    """Validate a static Qwen3 dense layer or component hook reference."""
    if isinstance(layer, int):
        if layer < 0:
            raise ValueError("Qwen3 layer indices must be non-negative integers.")
        return
    if not isinstance(layer, str):
        raise ValueError(
            f"Qwen3 hook references must be integers or strings, got {type(layer).__name__}."
        )

    component_ref = parse_qwen3_component_ref(layer)
    if component_ref is None:
        if layer.startswith("layer_") or layer.startswith("blocks."):
            examples = ", ".join(_QWEN3_COMPONENT_EXAMPLES[:6])
            raise ValueError(
                f"Invalid Qwen3 hook name {layer!r}. Expected SafeLens names such as "
                f"`layer_0.resid_pre` or TransformerLens-style names such as "
                f"`blocks.0.attn.hook_q`. Examples: {examples}."
            )
        return

    _layer_index, component = component_ref
    if component in _QWEN3_ATTENTION_COMPONENTS or component in _QWEN3_CACHE_ONLY_COMPONENTS:
        return
    if component not in _QWEN3_PATCHABLE_COMPONENTS and component not in _QWEN3_EXPLICIT_COMPONENTS:
        supported = ", ".join(qwen3_supported_hook_components(include_attention=True))
        examples = ", ".join(_QWEN3_COMPONENT_EXAMPLES[:6])
        raise ValueError(
            f"Unsupported Qwen3 hook component {component!r} in {layer!r}. "
            f"Supported components: {supported}. Examples: {examples}."
        )