Hooks And Activation Cache

SafeLens provides a lightweight hook layer inspired by TransformerLens, but without requiring TransformerLens as a dependency.

Design reference: TransformerLens hook points and ActivationCache.

Core utilities:

• HookPoint: dependency-free identity hook point with temporary/permanent hooks, context storage, prepend ordering, direction filters, and layer-name parsing.
• HookedRoot: root-level manager for named hook points, temporary hooks, and activation caching hooks.
• ActivationCache: dictionary-like activation store.
• make_cache_hook: creates a hook that captures an activation.
• temporary_hooks: registers hooks for one context and always removes them.
• run_with_hooks: runs a model with temporary hooks.
• HookedRoot.run_with_cache: runs a callable while temporarily caching named hook-point activations.
• HookedRoot.cache_all / cache_some: install persistent cache hooks until reset_hooks() removes them.
• cache_activations: runs a model and captures selected layer activations.
• activation_name_for_layer: standardizes cache names such as layer_0.
• get_act_name: TransformerLens-style shorthand names such as ("q", 2) to blocks.2.attn.hook_q, including layer-type aliases such as a, m, and b.
• safelens_act_name: SafeLens-style shorthand names such as layer_2.q.

Cache hooks follow TransformerLens' pos_slice conventions: head-vector activations such as hook_q, hook_k, hook_v, hook_z, and hook_result slice the [pos] axis before the head axis, while residual streams and attention patterns/scores slice the destination-position axis. ActivationCache decomposition helpers use the same negative-dimension position semantics, so position slicing still targets [pos] after remove_batch_dim.

Activation cache helpers copied in spirit from TransformerLens:

Helper	Purpose
cache_dict, has_embed, has_pos_embed	TransformerLens-compatible cache mapping and embed-presence attributes.
Tuple key lookup	Read cache[("resid_pre", 0)] or cache[("q", 2)].
keys_matching / select	Filter activations by names or predicates.
apply_to_values, detach, cpu, to	Transform all cached values; to mutates in place like TransformerLens.
remove_batch_dim	Remove a singleton batch dimension.
apply_slice_to_batch_dim	Slice all cached activations along batch.
stack_activation	Stack one activation type across layers.
accumulated_resid	Build a logit-lens residual stream stack.
decompose_resid	Split residual stream into embed, attention, and MLP terms.
stack_head_results	Stack per-head attention result vectors.
compute_head_results	Compute per-head result vectors from cached z and model W_O.
stack_neuron_results	Stack per-neuron MLP residual contributions when model.W_out is available, optionally projected onto output directions.
get_full_resid_decomposition	Stack head results, neuron/MLP terms, embeddings, optional model bias, and optional output-direction projections.
apply_ln_to_stack	Apply cached normalization scale with optional batch/position slicing, LN centering, or recomputed final LN for logit lens.
logit_attrs	Attribute residual components to token directions, token strings, or logit differences.

Example:

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

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

activation = cache["layer_0"]

Hook and activation-cache primitives inspired by TransformerLens.

ActivationCache

Bases: MutableMapping[Any, Any]

Dictionary-like activation cache with small tensor-friendly helpers.

Source code in src/SafeLens/core/hooks.py

class ActivationCache(MutableMapping[Any, Any]):
    """Dictionary-like activation cache with small tensor-friendly helpers."""

    def __init__(
        self,
        cache_dict: dict[str, Any] | None = None,
        model: Any = None,
        has_batch_dim: bool = True,
        canonicalize: bool = True,
    ) -> None:
        raw_cache = {} if cache_dict is None else cache_dict
        self._cache = _canonicalize_cache_dict(raw_cache) if canonicalize else raw_cache
        self.model = model
        self.has_batch_dim = has_batch_dim

    @property
    def cache_dict(self) -> dict[str, Any]:
        """TransformerLens-compatible view of the underlying activation mapping."""
        return self._cache

    @cache_dict.setter
    def cache_dict(self, value: dict[str, Any]) -> None:
        self._cache = _canonicalize_cache_dict(value)

    @property
    def has_embed(self) -> bool:
        """Return whether token embeddings are cached."""
        return "hook_embed" in self._cache

    @property
    def has_pos_embed(self) -> bool:
        """Return whether positional embeddings are cached."""
        return "hook_pos_embed" in self._cache

    def __getitem__(self, key: ActivationKey) -> Any:
        return self._cache[self.resolve_key(key)]

    def __setitem__(self, key: ActivationKey, value: Any) -> None:
        self._cache[self.storage_key(key)] = value

    def __delitem__(self, key: ActivationKey) -> None:
        del self._cache[self.resolve_key(key)]

    def __iter__(self) -> Iterator[str]:
        return iter(self._cache)

    def __len__(self) -> int:
        return len(self._cache)

    def __contains__(self, key: object) -> bool:
        if isinstance(key, str) or isinstance(key, tuple):
            try:
                self.resolve_key(key)
                return True
            except KeyError:
                return False
        return False

    def __repr__(self) -> str:
        batch = "with batch dim" if self.has_batch_dim else "without batch dim"
        return f"ActivationCache({len(self)} activations, {batch})"

    def keys(self) -> Any:
        """Return cached activation names, matching TransformerLens' mapping API."""
        return self._cache.keys()

    def values(self) -> Any:
        """Return cached activation values, matching TransformerLens' mapping API."""
        return self._cache.values()

    def items(self) -> Any:
        """Return cached activation items, matching TransformerLens' mapping API."""
        return self._cache.items()

    def resolve_key(self, key: ActivationKey) -> str:
        """Resolve exact, SafeLens-style, or TransformerLens-style activation keys."""
        candidates = activation_name_candidates(
            key,
            n_layers=self._infer_n_layers(),
            decoder_n_layers=self._infer_stack_n_layers("decoder"),
        )
        for candidate in candidates:
            if candidate in self._cache:
                return candidate
        raise KeyError(f"Unknown activation key {key!r}. Tried {candidates!r}.")

    def storage_key(self, key: ActivationKey) -> str:
        """Return the existing or canonical storage name for an activation key."""
        try:
            return self.resolve_key(key)
        except KeyError:
            if isinstance(key, tuple) and key:
                tuple_key = list(key)
                if len(tuple_key) == 1:
                    top_level_name = _TOP_LEVEL_ACT_NAMES.get(str(tuple_key[0]))
                    if top_level_name is not None:
                        return top_level_name
                layer = tuple_key[1] if len(tuple_key) >= 2 else None
                raw_name = _strip_hook_prefix(str(tuple_key[0]))
                name = _ACT_NAME_ALIASES.get(raw_name, raw_name)
                if layer == -1:
                    stack = _activation_key_stack_name(name)
                    n_layers = self._infer_stack_n_layers(stack)
                    if n_layers > 0:
                        layer = n_layers - 1
                layer_type = (
                    str(tuple_key[2]) if len(tuple_key) >= 3 and tuple_key[2] is not None else None
                )
                if layer_type is not None:
                    layer_type = _LAYER_TYPE_ALIASES.get(layer_type, layer_type)
                if layer_type and layer is not None:
                    return f"{activation_name_for_layer(layer)}.{layer_type}.{name}"
                return safelens_act_name(name, layer)
            if isinstance(key, str):
                canonical_key = _canonical_storage_key_for_string(key)
                if canonical_key != key:
                    return canonical_key
            candidates = activation_name_candidates(
                key,
                n_layers=self._infer_n_layers(),
                decoder_n_layers=self._infer_stack_n_layers("decoder"),
            )
            if candidates:
                return candidates[0]
            return str(key)

    def store(
        self,
        name: str,
        activation: Any,
        *,
        detach: bool = True,
        clone: bool = False,
        device: Any = None,
    ) -> None:
        """Store an activation, optionally detaching, cloning, or moving it."""
        self._cache[name] = prepare_activation_for_cache(
            activation,
            detach=detach,
            clone=clone,
            device=device,
        )

    def get_activation(self, name: str) -> Any:
        """Return one cached activation."""
        return self[name]

    def keys_matching(self, names_filter: NamesFilter) -> list[str]:
        """Return activation names matching a TransformerLens-style names filter."""
        return [name for name in self._cache if _cache_key_matches_filter(name, names_filter)]

    def select(self, names_filter: NamesFilter) -> ActivationCache:
        """Return a new cache containing only matching activation names."""
        return ActivationCache(
            {
                name: value
                for name, value in self._cache.items()
                if _cache_key_matches_filter(name, names_filter)
            },
            model=self.model,
            has_batch_dim=self.has_batch_dim,
        )

    def clone(self) -> ActivationCache:
        """Return a cloned copy when activations support `.clone()`, otherwise deep-copy values."""
        return ActivationCache(
            {name: clone_activation(value) for name, value in self._cache.items()},
            model=self.model,
            has_batch_dim=self.has_batch_dim,
        )

    def to_dict(self) -> dict[str, Any]:
        """Return a plain dictionary view copy."""
        return dict(self._cache)

    def apply_to_values(self, fn: Callable[[Any], Any]) -> ActivationCache:
        """Apply a function to every cached value and return a new cache."""
        return ActivationCache(
            {name: fn(value) for name, value in self._cache.items()},
            model=self.model,
            has_batch_dim=self.has_batch_dim,
        )

    def to(self, device: Any, move_model: bool | None = None) -> ActivationCache:
        """Move tensor-like activations to a device when values support `.to()`."""
        self._cache = {name: _move_value(value, device) for name, value in self._cache.items()}
        if move_model:
            model_to = getattr(self.model, "to", None)
            if callable(model_to):
                model_to(device)
        return self

    def cpu(self) -> ActivationCache:
        """Move tensor-like activations to CPU when supported."""
        return self.to("cpu")

    def detach(self) -> ActivationCache:
        """Detach tensor-like activations when values support `.detach()`."""
        return self.apply_to_values(_detach_value)

    def remove_batch_dim(self) -> ActivationCache:
        """Remove singleton batch dimensions in place and return this cache."""
        if not self.has_batch_dim:
            return self
        updated_values: dict[str, Any] = {}
        has_singleton_batch = any(_has_leading_dim(value, 1) for value in self._cache.values())
        for name, value in list(self._cache.items()):
            if _has_leading_dim(value, 1):
                updated_values[name] = _slice_dim(value, 0, dim=0)
                continue
            shape = _shape_of(value)
            if shape and not has_singleton_batch:
                raise ValueError(
                    f"Cannot remove batch dimension from cache with batch size > 1, "
                    f"for key {name} with shape {shape!r}."
                )
            updated_values[name] = value
        self._cache.update(updated_values)
        self.has_batch_dim = False
        return self

    def toggle_autodiff(self, mode: bool = False) -> None:
        """Set PyTorch's global grad-enabled state when PyTorch is available."""
        try:
            import torch
        except ImportError:
            return None
        torch.set_grad_enabled(mode)
        return None

    def apply_slice_to_batch_dim(self, batch_slice: Any) -> ActivationCache:
        """Return a cache sliced along the batch dimension."""
        normalized_slice = _normalize_slice_index(batch_slice)
        if not self.has_batch_dim:
            if normalized_slice == _FULL_SLICE:
                return ActivationCache(
                    dict(self._cache),
                    model=self.model,
                    has_batch_dim=False,
                )
            raise ValueError("Cannot slice batch dimension on a cache without batch dim.")
        has_batch_dim = not isinstance(normalized_slice, int)
        return ActivationCache(
            {
                name: _slice_dim(value, normalized_slice, dim=0)
                for name, value in self._cache.items()
            },
            model=self.model,
            has_batch_dim=has_batch_dim,
        )

    def stack_activation(
        self,
        activation_name: str,
        layer: int | None = None,
        layer_type: str | None = None,
        *,
        sublayer_type: str | None = None,
    ) -> Any:
        """Stack one activation across layers."""
        if sublayer_type is not None:
            layer_type = sublayer_type
        n_layers = self._normalize_layer(layer)
        values = [
            self[(activation_name, current_layer, layer_type)] for current_layer in range(n_layers)
        ]
        return stack_values(values)

    def accumulated_resid(
        self,
        layer: int | None = None,
        incl_mid: bool = False,
        apply_ln: bool = False,
        pos_slice: Any = None,
        mlp_input: bool = False,
        return_labels: bool = False,
        *,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Return residual stream states up to a layer, useful for logit-lens workflows."""
        stack = _normalize_residual_stack_name(stack)
        target_layer = self._normalize_stack_layer(layer, stack=stack)
        values: list[Any] = []
        labels: list[str] = []
        n_layers = self._infer_stack_n_layers(stack)
        max_pre_layer = min(target_layer, n_layers - 1)

        for current_layer in range(max_pre_layer + 1):
            resid_pre_key = _residual_component_key("resid_pre", current_layer, stack=stack)
            resid_mid_key = _residual_component_key("resid_mid", current_layer, stack=stack)
            if resid_pre_key in self:
                values.append(_maybe_slice_pos(self[resid_pre_key], pos_slice))
                labels.append(f"{current_layer}_pre")
            if incl_mid and current_layer < target_layer and resid_mid_key in self:
                values.append(_maybe_slice_pos(self[resid_mid_key], pos_slice))
                labels.append(f"{current_layer}_mid")

        resid_mid_key = _residual_component_key("resid_mid", target_layer, stack=stack)
        if mlp_input and resid_mid_key in self:
            values.append(_maybe_slice_pos(self[resid_mid_key], pos_slice))
            labels.append(f"{target_layer}_mid")
        final_post_key = _residual_component_key("resid_post", n_layers - 1, stack=stack)
        if target_layer >= n_layers and n_layers > 0 and final_post_key in self:
            values.append(_maybe_slice_pos(self[final_post_key], pos_slice))
            labels.append("final_post")
        if not values:
            raise KeyError("No residual stream activations found in cache.")

        residual_stack = stack_values(values)
        if apply_ln:
            residual_stack = self.apply_ln_to_stack(
                residual_stack,
                layer=target_layer,
                mlp_input=mlp_input,
                pos_slice=pos_slice,
                recompute_ln=target_layer == n_layers,
                has_batch_dim=self.has_batch_dim,
                stack=stack,
            )
        if return_labels:
            return residual_stack, labels
        return residual_stack

    def decompose_resid(
        self,
        layer: int | None = None,
        mlp_input: bool = False,
        mode: Literal["all", "mlp", "attn"] = "all",
        apply_ln: bool = False,
        pos_slice: Any = None,
        incl_embeds: bool = True,
        return_labels: bool = False,
        *,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Decompose a residual stream into embedding, attention, and MLP components."""
        stack = _normalize_residual_stack_name(stack)
        target_layer = self._normalize_stack_layer(layer, stack=stack)
        values: list[Any] = []
        labels: list[str] = []
        include_attn = mode != "mlp"
        include_mlp = mode != "attn" and not _model_is_attn_only(self.model)

        for key, label in (("hook_embed", "embed"), ("hook_pos_embed", "pos_embed")):
            if incl_embeds and key in self:
                values.append(_maybe_slice_pos(self[key], pos_slice))
                labels.append(label)

        for current_layer in range(target_layer):
            attn_key = _residual_component_key("attn_out", current_layer, stack=stack)
            cross_attn_key = _residual_component_key("cross_attn_out", current_layer, stack=stack)
            mlp_key = _residual_component_key("mlp_out", current_layer, stack=stack)
            if include_attn and attn_key in self:
                values.append(_maybe_slice_pos(self[attn_key], pos_slice))
                labels.append(f"{current_layer}_attn_out")
            if include_attn and stack == "decoder" and cross_attn_key in self:
                values.append(_maybe_slice_pos(self[cross_attn_key], pos_slice))
                labels.append(f"{current_layer}_cross_attn_out")
            if include_mlp and mlp_key in self:
                values.append(_maybe_slice_pos(self[mlp_key], pos_slice))
                labels.append(f"{current_layer}_mlp_out")

        attn_key = _residual_component_key("attn_out", target_layer, stack=stack)
        cross_attn_key = _residual_component_key("cross_attn_out", target_layer, stack=stack)
        if mlp_input and include_attn and attn_key in self:
            values.append(_maybe_slice_pos(self[attn_key], pos_slice))
            labels.append(f"{target_layer}_attn_out")
        if mlp_input and include_attn and stack == "decoder" and cross_attn_key in self:
            values.append(_maybe_slice_pos(self[cross_attn_key], pos_slice))
            labels.append(f"{target_layer}_cross_attn_out")
        if not values:
            raise KeyError("No residual decomposition activations found in cache.")

        residual_stack = stack_values(values)
        if apply_ln:
            residual_stack = self.apply_ln_to_stack(
                residual_stack,
                layer=target_layer,
                mlp_input=mlp_input,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
                stack=stack,
            )
        if return_labels:
            return residual_stack, labels
        return residual_stack

    def stack_head_results(
        self,
        layer: int | None = None,
        return_labels: bool = False,
        incl_remainder: bool = False,
        pos_slice: Any = None,
        apply_ln: bool = False,
        *,
        component: str = "result",
    ) -> Any:
        """Stack per-head activations from `[batch, pos, head, d_model]` caches."""
        if incl_remainder and component != "result":
            raise ValueError(
                "incl_remainder=True requires residual-space `result` head activations."
            )
        target_layer = self._normalize_layer(layer)
        if component == "result" and any(
            ("z", current_layer) in self and (component, current_layer) not in self
            for current_layer in range(target_layer)
        ):
            try:
                self.compute_head_results(target_layer, store=True)
            except (KeyError, ValueError):
                pass
        values: list[Any] = []
        labels: list[str] = []
        for current_layer in range(target_layer):
            if (component, current_layer) not in self:
                continue
            activation = _maybe_slice_pos(
                self[(component, current_layer)],
                pos_slice,
                dim=_head_vector_pos_dim(component),
            )
            head_dim = _head_axis_after_pos_slice(activation, pos_slice, component=component)
            for head_index in range(_infer_head_count(activation, dim=head_dim)):
                values.append(_slice_dim(activation, head_index, dim=head_dim))
                labels.append(f"L{current_layer}H{head_index}")
        if incl_remainder:
            remainder = _residual_remainder_base(self, target_layer, pos_slice)
            if values:
                remainder = _subtract_values(remainder, _sum_values(values))
            values.append(remainder)
            labels.append("remainder")
        if not values:
            if target_layer == 0:
                head_stack = _empty_component_stack_like_cache(
                    self,
                    pos_slice=pos_slice,
                    has_batch_dim=self.has_batch_dim,
                )
                if apply_ln:
                    head_stack = self.apply_ln_to_stack(
                        head_stack,
                        layer=target_layer,
                        pos_slice=pos_slice,
                        has_batch_dim=self.has_batch_dim,
                    )
                if return_labels:
                    return head_stack, labels
                return head_stack
            raise KeyError(f"No {component!r} head activations found in cache.")
        head_stack = stack_values(values)
        if apply_ln:
            head_stack = self.apply_ln_to_stack(
                head_stack,
                layer=target_layer,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
            )
        if return_labels:
            return head_stack, labels
        return head_stack

    def compute_head_results(
        self,
        layer: int | None = None,
        *,
        store: bool = True,
        pos_slice: Any = None,
        return_labels: bool = False,
    ) -> Any:
        """Compute per-head residual-space results from cached `z` and model `W_O`.

        This fills the common TransformerLens workflow gap where a model exposes
        head outputs `z` and output weights `W_O`, but not cached `result`
        vectors directly.
        """
        target_layer = self._normalize_layer(layer)
        n_layers = self._infer_n_layers()
        max_layer = min(target_layer, n_layers)
        values: list[Any] = []
        labels: list[str] = []

        for current_layer in range(max_layer):
            if ("result", current_layer) in self:
                result = _maybe_slice_pos(
                    self[("result", current_layer)],
                    pos_slice,
                    dim=_head_vector_pos_dim("result"),
                )
                values.append(result)
                labels.append(f"{current_layer}_result")
                continue
            if ("z", current_layer) not in self:
                continue
            w_o = _get_layer_weight(self.model, "W_O", current_layer)
            if w_o is None:
                raise ValueError(
                    "compute_head_results requires a cache with a model exposing W_O "
                    "when cached head result activations are missing."
                )
            from SafeLens.core.analysis import compute_head_results_from_z

            z_activation = _maybe_slice_pos(self[("z", current_layer)], pos_slice, dim=-3)
            result = compute_head_results_from_z(z_activation, w_o)
            if store and pos_slice is None:
                self[f"layer_{current_layer}.result"] = result
                self.cache_dict[f"blocks.{current_layer}.attn.hook_result"] = result
            values.append(result)
            labels.append(f"{current_layer}_result")

        if not values:
            raise KeyError("No cached `z` activations found for head result computation.")
        result_stack = stack_values(values)
        if return_labels:
            return result_stack, labels
        return result_stack

    def get_neuron_results(
        self,
        layer: int,
        neuron_slice: Any = None,
        pos_slice: Any = None,
        project_output_onto: Any = None,
        *,
        component: str = "post",
    ) -> Any:
        """Return one layer's per-neuron residual contributions."""
        if (component, layer) not in self:
            raise KeyError(f"No cached {component!r} neuron activations found for layer {layer}.")
        w_out = _get_layer_weight(self.model, "W_out", layer)
        if w_out is None:
            raise ValueError("get_neuron_results requires a model exposing W_out.")

        neuron_acts = _maybe_slice_pos(self[(component, layer)], pos_slice)
        neuron_count = _infer_last_dim(neuron_acts)
        neuron_indices = _indices_from_slice(neuron_slice, neuron_count)
        neuron_acts = _select_indices_dim(neuron_acts, neuron_indices, dim=-1)
        layer_w_out = _select_indices_dim(w_out, neuron_indices, dim=0)
        if project_output_onto is not None:
            layer_w_out = _project_last_dim(layer_w_out, project_output_onto)
        return _multiply_last_dim_by_matrix(neuron_acts, layer_w_out)

    def _get_cached_ln_scale(
        self,
        layer: int | None,
        mlp_input: bool,
        pos_slice: Any = None,
        batch_slice: Any = None,
    ) -> Any:
        """Return cached layer-norm scale for a residual-stack target."""
        target_layer = self._normalize_layer(layer)
        n_layers = self._infer_n_layers()
        if target_layer == n_layers:
            key: ActivationKey = "ln_final.hook_scale"
        else:
            key = ("scale", target_layer, "ln2" if mlp_input else "ln1")
        try:
            scale = self[key]
        except KeyError as exc:
            resolved_key = key if isinstance(key, str) else get_act_name(*key)
            raise KeyError(
                f"Cached LN scale not found at {resolved_key!r}. apply_ln operations require "
                "this hook to be cached for the requested layer."
            ) from exc
        scale_has_batch_dim = self.has_batch_dim
        if batch_slice is not None and self.has_batch_dim:
            scale = _slice_dim(scale, batch_slice, dim=0)
            scale_has_batch_dim = not isinstance(batch_slice, int)
        if pos_slice is not None:
            pos_dim = _scale_pos_dim(scale, has_batch_dim=scale_has_batch_dim)
            scale = _slice_dim(scale, pos_slice, dim=pos_dim)
        return scale

    def _stack_neuron_results_apply_ln_projected(
        self,
        layer: int,
        pos_slice: Any,
        neuron_slice: Any,
        project_output_onto: Any,
    ) -> Any:
        """Stack LN-applied neuron projections without materializing d_mlp by d_model."""
        scale = self._get_cached_ln_scale(layer, mlp_input=False, pos_slice=pos_slice)
        apply_centering = _uses_centered_layer_norm(self.model)
        projection_sum = _sum_projection_input_dim(project_output_onto) if apply_centering else None
        values: list[Any] = []
        for current_layer in range(layer):
            if ("post", current_layer) not in self:
                continue
            neuron_acts = _maybe_slice_pos(self[("post", current_layer)], pos_slice)
            neuron_indices = _indices_from_slice(neuron_slice, _infer_last_dim(neuron_acts))
            neuron_acts = _select_indices_dim(neuron_acts, neuron_indices, dim=-1)
            w_out = _get_layer_weight(self.model, "W_out", current_layer)
            if w_out is None:
                raise ValueError("stack_neuron_results requires a model exposing W_out.")
            w_out = _select_indices_dim(w_out, neuron_indices, dim=0)
            linear_form = _project_last_dim(w_out, project_output_onto)
            if apply_centering:
                w_means = _mean_last_dim(w_out)
                linear_form = _subtract_values(
                    linear_form,
                    _multiply_by_last_vector(w_means, projection_sum),
                )
            scaled_acts = _divide_values(neuron_acts, scale)
            layer_values = _multiply_last_dim_by_matrix(scaled_acts, linear_form)
            neuron_dim = -1 if len(_shape_of(project_output_onto)) == 1 else -2
            for position in range(len(neuron_indices)):
                values.append(_slice_dim(layer_values, position, dim=neuron_dim))
        if values:
            return stack_values(values)
        return _empty_component_stack_like_cache(
            self,
            pos_slice=pos_slice,
            has_batch_dim=self.has_batch_dim,
            project_output_onto=project_output_onto,
        )

    def _can_fold_ln_neuron_projection(self, layer: int, pos_slice: Any) -> bool:
        try:
            scale = self._get_cached_ln_scale(layer, mlp_input=False, pos_slice=pos_slice)
        except KeyError:
            return False
        scale_shape = _shape_of(scale)
        return not scale_shape or scale_shape[-1] == 1

    def stack_neuron_results(
        self,
        layer: int | None = None,
        pos_slice: Any = None,
        neuron_slice: Any = None,
        return_labels: bool = False,
        incl_remainder: bool = False,
        apply_ln: bool = False,
        project_output_onto: Any = None,
        *,
        component: str = "post",
        require_output_weight: bool = False,
    ) -> Any:
        """Stack per-neuron MLP residual contributions when `W_out` is available."""
        target_layer = self._normalize_layer(layer)
        values: list[Any] = []
        labels: list[str] = []
        can_project_before_stack = (
            project_output_onto is not None and not apply_ln and not incl_remainder
        )
        can_fold_ln_projection = (
            component == "post"
            and project_output_onto is not None
            and apply_ln
            and not incl_remainder
            and self._can_fold_ln_neuron_projection(target_layer, pos_slice)
        )
        results_are_projected = False
        for current_layer in range(target_layer):
            if (component, current_layer) not in self:
                continue
            activation = _maybe_slice_pos(self[(component, current_layer)], pos_slice)
            neuron_indices = _indices_from_slice(neuron_slice, _infer_last_dim(activation))
            if can_fold_ln_projection:
                labels.extend(f"L{current_layer}N{neuron_index}" for neuron_index in neuron_indices)
                continue
            try:
                layer_results = self.get_neuron_results(
                    current_layer,
                    neuron_slice=neuron_indices,
                    pos_slice=pos_slice,
                    project_output_onto=project_output_onto if can_project_before_stack else None,
                    component=component,
                )
                results_are_projected = can_project_before_stack
                neuron_dim = (
                    -1 if results_are_projected and len(_shape_of(project_output_onto)) == 1 else -2
                )
            except ValueError:
                if project_output_onto is not None or require_output_weight or incl_remainder:
                    raise
                layer_results = _select_indices_dim(activation, neuron_indices, dim=-1)
                neuron_dim = -1
                results_are_projected = False
            for position, neuron_index in enumerate(neuron_indices):
                neuron_value = _slice_dim(layer_results, position, dim=neuron_dim)
                values.append(neuron_value)
                labels.append(f"L{current_layer}N{neuron_index}")
        if incl_remainder:
            remainder = _residual_remainder_base(self, target_layer, pos_slice)
            if values:
                remainder = _subtract_values(remainder, _sum_values(values))
            values.append(remainder)
            labels.append("remainder")
        if not values and not (can_fold_ln_projection and labels):
            if target_layer == 0:
                neuron_stack = _empty_component_stack_like_cache(
                    self,
                    pos_slice=pos_slice,
                    has_batch_dim=self.has_batch_dim,
                    project_output_onto=project_output_onto,
                )
                if apply_ln:
                    neuron_stack = self.apply_ln_to_stack(
                        neuron_stack,
                        layer=target_layer,
                        pos_slice=pos_slice,
                        has_batch_dim=self.has_batch_dim,
                    )
                if return_labels:
                    return neuron_stack, labels
                return neuron_stack
            raise KeyError(f"No {component!r} neuron activations found in cache.")
        if can_fold_ln_projection:
            neuron_stack = self._stack_neuron_results_apply_ln_projected(
                target_layer,
                pos_slice,
                neuron_slice,
                project_output_onto,
            )
            results_are_projected = True
        else:
            neuron_stack = stack_values(values)
        if apply_ln and not can_fold_ln_projection:
            neuron_stack = self.apply_ln_to_stack(
                neuron_stack,
                layer=target_layer,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
            )
        if project_output_onto is not None and not results_are_projected:
            neuron_stack = _project_last_dim(neuron_stack, project_output_onto)
        if return_labels:
            return neuron_stack, labels
        return neuron_stack

    def get_full_resid_decomposition(
        self,
        layer: int | None = None,
        mlp_input: bool = False,
        expand_neurons: bool = True,
        apply_ln: bool = False,
        pos_slice: Any = None,
        return_labels: bool = False,
        project_output_onto: Any = None,
        *,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Return a best-effort decomposition into heads, MLP neurons, embeds, and bias."""
        stack = _normalize_residual_stack_name(stack)
        target_layer = self._normalize_stack_layer(layer, stack=stack)
        stacks: list[Any] = []
        labels: list[str] = []
        expanded_neurons = False
        ln_folded = apply_ln and project_output_onto is not None
        bias_reference: Any = None

        def remember_bias_reference(component_stack: Any) -> None:
            nonlocal bias_reference
            if bias_reference is not None:
                return
            component_rows = _unstack_first_dim(component_stack)
            if component_rows:
                bias_reference = component_rows[0]

        def maybe_ln_then_project(component_stack: Any) -> Any:
            if ln_folded:
                component_stack = self.apply_ln_to_stack(
                    component_stack,
                    layer=target_layer,
                    mlp_input=mlp_input,
                    pos_slice=pos_slice,
                    has_batch_dim=self.has_batch_dim,
                    stack=stack,
                )
                return _project_last_dim(component_stack, project_output_onto)
            if project_output_onto is not None:
                return _project_last_dim(component_stack, project_output_onto)
            return component_stack

        def add_stack(component_stack: Any, component_labels: list[str]) -> None:
            stacks.extend(_unstack_first_dim(component_stack))
            labels.extend(component_labels)

        try:
            if stack == "decoder":
                head_stack, head_labels = _decoder_head_result_stack(
                    self,
                    target_layer + (1 if mlp_input else 0),
                    pos_slice=pos_slice,
                )
            else:
                head_stack, head_labels = self.stack_head_results(
                    target_layer + (1 if mlp_input else 0),
                    pos_slice=pos_slice,
                    return_labels=True,
                )
            remember_bias_reference(head_stack)
            head_stack = maybe_ln_then_project(head_stack)
            add_stack(head_stack, head_labels)
        except KeyError:
            try:
                attn_stack, attn_labels = self.decompose_resid(
                    target_layer,
                    mlp_input=mlp_input,
                    mode="attn",
                    incl_embeds=False,
                    pos_slice=pos_slice,
                    return_labels=True,
                    stack=stack,
                )
                remember_bias_reference(attn_stack)
                attn_stack = maybe_ln_then_project(attn_stack)
                add_stack(attn_stack, attn_labels)
            except KeyError:
                pass

        if stack == "decoder" and not _model_is_attn_only(self.model):
            try:
                mlp_stack, mlp_labels = self.decompose_resid(
                    target_layer,
                    mode="mlp",
                    incl_embeds=False,
                    pos_slice=pos_slice,
                    return_labels=True,
                    stack=stack,
                )
                remember_bias_reference(mlp_stack)
                mlp_stack = maybe_ln_then_project(mlp_stack)
                add_stack(mlp_stack, mlp_labels)
            except KeyError:
                pass
        elif expand_neurons and not _model_is_attn_only(self.model):
            try:
                neuron_stack, neuron_labels = self.stack_neuron_results(
                    target_layer,
                    pos_slice=pos_slice,
                    return_labels=True,
                    require_output_weight=True,
                    apply_ln=ln_folded,
                    project_output_onto=project_output_onto,
                )
                if project_output_onto is None and not apply_ln:
                    remember_bias_reference(neuron_stack)
                add_stack(neuron_stack, neuron_labels)
                expanded_neurons = True
            except (KeyError, ValueError):
                try:
                    mlp_stack, mlp_labels = self.decompose_resid(
                        target_layer,
                        mode="mlp",
                        incl_embeds=False,
                        pos_slice=pos_slice,
                        return_labels=True,
                        stack=stack,
                    )
                    remember_bias_reference(mlp_stack)
                    mlp_stack = maybe_ln_then_project(mlp_stack)
                    add_stack(mlp_stack, mlp_labels)
                except KeyError:
                    pass
        elif not _model_is_attn_only(self.model):
            try:
                mlp_stack, mlp_labels = self.decompose_resid(
                    target_layer,
                    mode="mlp",
                    incl_embeds=False,
                    pos_slice=pos_slice,
                    return_labels=True,
                    stack=stack,
                )
                remember_bias_reference(mlp_stack)
                mlp_stack = maybe_ln_then_project(mlp_stack)
                add_stack(mlp_stack, mlp_labels)
            except KeyError:
                pass

        for key, label in (("hook_embed", "embed"), ("hook_pos_embed", "pos_embed")):
            if key in self:
                embed_stack = stack_values([_maybe_slice_pos(self[key], pos_slice)])
                remember_bias_reference(embed_stack)
                embed_stack = maybe_ln_then_project(embed_stack)
                stacks.extend(_unstack_first_dim(embed_stack))
                labels.append(label)

        accumulated_bias = _get_model_attr(self.model, "accumulated_bias")
        if callable(accumulated_bias):
            try:
                bias = accumulated_bias(
                    target_layer,
                    mlp_input,
                    include_mlp_biases=expanded_neurons,
                )
            except TypeError:
                try:
                    bias = accumulated_bias(target_layer, mlp_input)
                except TypeError:
                    bias = accumulated_bias(target_layer)
            if ln_folded:
                bias = _expand_bias_like_for_folded_projection(
                    bias,
                    bias_reference,
                    stacks[0] if stacks else None,
                    project_output_onto,
                )
                bias_stack = stack_values([bias])
                bias_stack = self.apply_ln_to_stack(
                    bias_stack,
                    layer=target_layer,
                    mlp_input=mlp_input,
                    pos_slice=pos_slice,
                    has_batch_dim=self.has_batch_dim,
                    stack=stack,
                )
                bias_stack = _project_last_dim(bias_stack, project_output_onto)
                stacks.extend(_unstack_first_dim(bias_stack))
            else:
                if project_output_onto is not None:
                    bias = _project_last_dim(bias, project_output_onto)
                bias = _expand_bias_like(bias, stacks[0] if stacks else None)
                stacks.append(bias)
            labels.append("bias")

        if not stacks:
            raise KeyError("No activations found for a full residual decomposition.")
        full_stack = stack_values(stacks)
        if apply_ln and not ln_folded:
            full_stack = self.apply_ln_to_stack(
                full_stack,
                layer=target_layer,
                mlp_input=mlp_input,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
                stack=stack,
            )
        if return_labels:
            return full_stack, labels
        return full_stack

    def apply_ln_to_stack(
        self,
        residual_stack: Any,
        layer: int | None = None,
        mlp_input: bool = False,
        pos_slice: Any = None,
        batch_slice: Any = None,
        has_batch_dim: bool | None = None,
        recompute_ln: bool = False,
        *,
        scale_key: ActivationKey | None = None,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Apply cached layer-norm scale to a residual stack when scale is available."""
        stack = _normalize_residual_stack_name(stack)
        resolved_has_batch_dim = self.has_batch_dim if has_batch_dim is None else has_batch_dim
        target_layer = self._normalize_stack_layer(layer, stack=stack)
        n_layers = self._infer_stack_n_layers(stack)
        if _model_explicitly_has_no_layer_norm(self.model):
            return residual_stack
        requires_cached_scale = _model_explicitly_uses_layer_norm(self.model)
        if batch_slice is not None and resolved_has_batch_dim:
            residual_stack = _slice_dim(residual_stack, batch_slice, dim=1)
        if recompute_ln and scale_key is None and target_layer == n_layers:
            ln_final = _get_final_layer_norm(self.model, stack=stack)
            recomputed = _apply_final_layer_norm_to_stack(residual_stack, ln_final)
            if recomputed is not _MISSING:
                return recomputed

        resolved_scale_key = scale_key
        candidates: list[ActivationKey] = []
        if resolved_scale_key is None:
            if target_layer == n_layers:
                candidates = ["ln_final.hook_scale"]
            else:
                requested_layer_norm = _residual_stack_ln_name(stack, mlp_input)
                fallback_layer_norm = _residual_stack_fallback_ln_name(stack, mlp_input)
                candidates = _layer_norm_scale_candidates(target_layer, requested_layer_norm)
                if fallback_layer_norm is not None and not requires_cached_scale:
                    candidates.extend(
                        _layer_norm_scale_candidates(target_layer, fallback_layer_norm)
                    )
            resolved_scale_key = next(
                (candidate for candidate in candidates if candidate in self),
                None,
            )
        if resolved_scale_key is None:
            if requires_cached_scale:
                expected_key = candidates[0] if scale_key is None else scale_key
                expected_name = (
                    get_act_name(*expected_key)
                    if isinstance(expected_key, tuple)
                    else str(expected_key)
                )
                raise KeyError(
                    f"Cached LN scale not found at {expected_name!r}. apply_ln operations "
                    "require this hook to be cached for the requested layer."
                )
            return residual_stack
        scale = self[resolved_scale_key]
        if batch_slice is not None and resolved_has_batch_dim:
            scale = _slice_dim(scale, batch_slice, dim=0)
        if pos_slice is not None:
            pos_dim = _scale_pos_dim_for_residual_stack(
                scale,
                residual_stack,
                has_batch_dim=resolved_has_batch_dim,
            )
            scale = _slice_dim(scale, pos_slice, dim=pos_dim)
        if _uses_centered_layer_norm(self.model):
            residual_stack = _subtract_last_dim_mean(residual_stack)
        return _divide_values(residual_stack, scale)

    def logit_attrs(
        self,
        residual_stack: Any,
        tokens: Any,
        incorrect_tokens: Any = None,
        pos_slice: Any = None,
        batch_slice: Any = None,
        has_batch_dim: bool | None = None,
        *,
        directions: Any = None,
        apply_ln: bool = True,
    ) -> Any:
        """Project residual components onto token residual directions."""
        resolved_has_batch_dim = self.has_batch_dim if has_batch_dim is None else has_batch_dim
        if directions is None:
            if self.model is None:
                directions = tokens
            else:
                directions = self.model.tokens_to_residual_directions(
                    _normalize_logit_tokens(self.model, tokens)
                )
        if incorrect_tokens is not None:
            if self.model is None:
                incorrect_directions = incorrect_tokens
            else:
                incorrect_directions = self.model.tokens_to_residual_directions(
                    _normalize_logit_tokens(self.model, incorrect_tokens)
                )
            if _shape_of(directions) != _shape_of(incorrect_directions):
                raise ValueError(
                    "tokens and incorrect_tokens must resolve to residual directions with the "
                    f"same shape, got {_shape_of(directions)!r} and "
                    f"{_shape_of(incorrect_directions)!r}."
                )
            directions = _subtract_values(directions, incorrect_directions)
        batch_dim = _direction_batch_dim(
            directions,
            residual_stack,
            has_batch_dim=resolved_has_batch_dim,
            prefer_pos_axis=pos_slice is not None,
        )
        residual_batch_dim = _residual_stack_batch_dim(
            residual_stack,
            directions,
            pos_slice=pos_slice,
        )
        if batch_slice is not None and batch_dim is not None:
            directions = _slice_dim(directions, batch_slice, dim=batch_dim)
            direction_has_batch_dim = not isinstance(batch_slice, int)
        else:
            direction_has_batch_dim = resolved_has_batch_dim
        pos_dim = _direction_pos_dim(
            directions,
            residual_stack,
            has_batch_dim=direction_has_batch_dim,
            prefer_pos_axis=pos_slice is not None,
        )
        if pos_slice is not None and pos_dim is not None:
            directions = _slice_dim(directions, pos_slice, dim=pos_dim)
        if apply_ln:
            residual_stack = _slice_residual_stack_for_logit_attrs(
                residual_stack,
                pos_slice,
                batch_slice=batch_slice,
                has_batch_dim=resolved_has_batch_dim,
                residual_batch_dim=residual_batch_dim,
                reference=directions,
            )
            residual_stack = _apply_cached_ln_scale_to_sliced_stack(
                self,
                residual_stack,
                layer=-1,
                pos_slice=pos_slice,
                batch_slice=batch_slice,
            )
        elif batch_slice is not None and resolved_has_batch_dim:
            residual_stack = _slice_dim(
                residual_stack,
                batch_slice,
                dim=residual_batch_dim,
            )
        return _dot_last_dim(residual_stack, directions)

    def residual_stack_to_logits(
        self,
        residual_stack: Any,
        *,
        apply_ln: bool = True,
        pos_slice: Any = None,
        batch_slice: Any = None,
        has_batch_dim: bool | None = None,
        use_unembed_bias: bool = True,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Project residual components through the model unembedding matrix."""
        stack = _normalize_residual_stack_name(stack)
        if self.model is None:
            raise ValueError("residual_stack_to_logits requires a cache with an attached model.")
        unembed = _get_model_attr(self.model, "W_U")
        if unembed is None:
            raise ValueError("residual_stack_to_logits requires a model exposing W_U.")
        unembed_bias = _get_model_attr(self.model, "b_U") if use_unembed_bias else None
        resolved_has_batch_dim = self.has_batch_dim if has_batch_dim is None else has_batch_dim
        residual_batch_dim = _residual_stack_batch_dim(residual_stack, None, pos_slice=pos_slice)
        if apply_ln:
            residual_stack = _maybe_slice_residual_stack_pos(
                residual_stack,
                pos_slice,
                has_batch_dim=resolved_has_batch_dim,
            )
            residual_stack = self.apply_ln_to_stack(
                residual_stack,
                layer=-1,
                pos_slice=pos_slice,
                batch_slice=batch_slice,
                has_batch_dim=resolved_has_batch_dim,
                stack=stack,
            )
        else:
            if batch_slice is not None and resolved_has_batch_dim:
                residual_stack = _slice_dim(residual_stack, batch_slice, dim=residual_batch_dim)
            if pos_slice is not None:
                residual_stack = _slice_dim(
                    residual_stack,
                    pos_slice,
                    dim=_residual_stack_pos_dim_after_optional_batch_slice(
                        residual_stack,
                        has_batch_dim=resolved_has_batch_dim,
                        batch_slice=batch_slice,
                    ),
                )
        from SafeLens.core.analysis import residual_stack_to_logits

        return residual_stack_to_logits(residual_stack, unembed, unembed_bias)

    def accumulated_resid_to_logits(
        self,
        layer: int | None = None,
        incl_mid: bool = False,
        *,
        pos_slice: Any = None,
        return_labels: bool = False,
        use_unembed_bias: bool = True,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Run a logit lens over accumulated residual stream states."""
        stack = _normalize_residual_stack_name(stack)
        residual_stack, labels = self.accumulated_resid(
            layer=layer,
            incl_mid=incl_mid,
            apply_ln=True,
            pos_slice=pos_slice,
            return_labels=True,
            stack=stack,
        )
        logits = self.residual_stack_to_logits(
            residual_stack,
            apply_ln=False,
            use_unembed_bias=use_unembed_bias,
        )
        if return_labels:
            return logits, labels
        return logits

    def decompose_resid_to_logits(
        self,
        layer: int | None = None,
        mlp_input: bool = False,
        mode: Literal["all", "mlp", "attn"] = "all",
        *,
        pos_slice: Any = None,
        incl_embeds: bool = True,
        return_labels: bool = False,
        use_unembed_bias: bool = False,
        stack: Literal["encoder", "decoder"] = "encoder",
    ) -> Any:
        """Project residual-decomposition components through the unembedding."""
        stack = _normalize_residual_stack_name(stack)
        residual_stack, labels = self.decompose_resid(
            layer=layer,
            mlp_input=mlp_input,
            mode=mode,
            apply_ln=True,
            pos_slice=pos_slice,
            incl_embeds=incl_embeds,
            return_labels=True,
            stack=stack,
        )
        logits = self.residual_stack_to_logits(
            residual_stack,
            apply_ln=False,
            use_unembed_bias=use_unembed_bias,
        )
        if return_labels:
            return logits, labels
        return logits

    def _infer_n_layers(self) -> int:
        n_layers = _get_config_int(self.model, ("n_layers", "num_hidden_layers", "num_layers"))
        if n_layers is not None:
            return n_layers
        layers = _layer_indices_from_cache_keys(
            self._cache,
            patterns=(r"(?:blocks|encoder|decoder)\.(\d+)\.", r"layer_(\d+)\."),
        )
        if layers:
            return max(layers) + 1
        return 0

    def _infer_stack_n_layers(self, stack: Literal["encoder", "decoder"]) -> int:
        if stack == "encoder":
            return self._infer_n_layers()
        n_layers = _get_config_int(
            self.model,
            ("n_decoder_layers", "num_decoder_layers", "decoder_layers"),
        )
        if n_layers is not None:
            return n_layers
        layers = _layer_indices_from_cache_keys(
            self._cache,
            patterns=(
                r"decoder\.(\d+)\.",
                r"layer_(\d+)\.(?:decoder_|cross_)",
            ),
        )
        if layers:
            return max(layers) + 1
        return self._infer_n_layers()

    def _normalize_layer(self, layer: int | None) -> int:
        n_layers = self._infer_n_layers()
        if layer is None or layer == -1:
            return n_layers
        if layer < 0:
            return n_layers + layer
        return layer

    def _normalize_stack_layer(
        self,
        layer: int | None,
        *,
        stack: Literal["encoder", "decoder"],
    ) -> int:
        n_layers = self._infer_stack_n_layers(stack)
        if layer is None or layer == -1:
            return n_layers
        if layer < 0:
            return n_layers + layer
        return layer

cache_dict property writable

TransformerLens-compatible view of the underlying activation mapping.

has_embed property

Return whether token embeddings are cached.

has_pos_embed property

Return whether positional embeddings are cached.

accumulated_resid(layer=None, incl_mid=False, apply_ln=False, pos_slice=None, mlp_input=False, return_labels=False, *, stack='encoder')

Return residual stream states up to a layer, useful for logit-lens workflows.

Source code in src/SafeLens/core/hooks.py

def accumulated_resid(
    self,
    layer: int | None = None,
    incl_mid: bool = False,
    apply_ln: bool = False,
    pos_slice: Any = None,
    mlp_input: bool = False,
    return_labels: bool = False,
    *,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Return residual stream states up to a layer, useful for logit-lens workflows."""
    stack = _normalize_residual_stack_name(stack)
    target_layer = self._normalize_stack_layer(layer, stack=stack)
    values: list[Any] = []
    labels: list[str] = []
    n_layers = self._infer_stack_n_layers(stack)
    max_pre_layer = min(target_layer, n_layers - 1)

    for current_layer in range(max_pre_layer + 1):
        resid_pre_key = _residual_component_key("resid_pre", current_layer, stack=stack)
        resid_mid_key = _residual_component_key("resid_mid", current_layer, stack=stack)
        if resid_pre_key in self:
            values.append(_maybe_slice_pos(self[resid_pre_key], pos_slice))
            labels.append(f"{current_layer}_pre")
        if incl_mid and current_layer < target_layer and resid_mid_key in self:
            values.append(_maybe_slice_pos(self[resid_mid_key], pos_slice))
            labels.append(f"{current_layer}_mid")

    resid_mid_key = _residual_component_key("resid_mid", target_layer, stack=stack)
    if mlp_input and resid_mid_key in self:
        values.append(_maybe_slice_pos(self[resid_mid_key], pos_slice))
        labels.append(f"{target_layer}_mid")
    final_post_key = _residual_component_key("resid_post", n_layers - 1, stack=stack)
    if target_layer >= n_layers and n_layers > 0 and final_post_key in self:
        values.append(_maybe_slice_pos(self[final_post_key], pos_slice))
        labels.append("final_post")
    if not values:
        raise KeyError("No residual stream activations found in cache.")

    residual_stack = stack_values(values)
    if apply_ln:
        residual_stack = self.apply_ln_to_stack(
            residual_stack,
            layer=target_layer,
            mlp_input=mlp_input,
            pos_slice=pos_slice,
            recompute_ln=target_layer == n_layers,
            has_batch_dim=self.has_batch_dim,
            stack=stack,
        )
    if return_labels:
        return residual_stack, labels
    return residual_stack

accumulated_resid_to_logits(layer=None, incl_mid=False, *, pos_slice=None, return_labels=False, use_unembed_bias=True, stack='encoder')

Run a logit lens over accumulated residual stream states.

Source code in src/SafeLens/core/hooks.py

def accumulated_resid_to_logits(
    self,
    layer: int | None = None,
    incl_mid: bool = False,
    *,
    pos_slice: Any = None,
    return_labels: bool = False,
    use_unembed_bias: bool = True,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Run a logit lens over accumulated residual stream states."""
    stack = _normalize_residual_stack_name(stack)
    residual_stack, labels = self.accumulated_resid(
        layer=layer,
        incl_mid=incl_mid,
        apply_ln=True,
        pos_slice=pos_slice,
        return_labels=True,
        stack=stack,
    )
    logits = self.residual_stack_to_logits(
        residual_stack,
        apply_ln=False,
        use_unembed_bias=use_unembed_bias,
    )
    if return_labels:
        return logits, labels
    return logits

apply_ln_to_stack(residual_stack, layer=None, mlp_input=False, pos_slice=None, batch_slice=None, has_batch_dim=None, recompute_ln=False, *, scale_key=None, stack='encoder')

Apply cached layer-norm scale to a residual stack when scale is available.

Source code in src/SafeLens/core/hooks.py

def apply_ln_to_stack(
    self,
    residual_stack: Any,
    layer: int | None = None,
    mlp_input: bool = False,
    pos_slice: Any = None,
    batch_slice: Any = None,
    has_batch_dim: bool | None = None,
    recompute_ln: bool = False,
    *,
    scale_key: ActivationKey | None = None,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Apply cached layer-norm scale to a residual stack when scale is available."""
    stack = _normalize_residual_stack_name(stack)
    resolved_has_batch_dim = self.has_batch_dim if has_batch_dim is None else has_batch_dim
    target_layer = self._normalize_stack_layer(layer, stack=stack)
    n_layers = self._infer_stack_n_layers(stack)
    if _model_explicitly_has_no_layer_norm(self.model):
        return residual_stack
    requires_cached_scale = _model_explicitly_uses_layer_norm(self.model)
    if batch_slice is not None and resolved_has_batch_dim:
        residual_stack = _slice_dim(residual_stack, batch_slice, dim=1)
    if recompute_ln and scale_key is None and target_layer == n_layers:
        ln_final = _get_final_layer_norm(self.model, stack=stack)
        recomputed = _apply_final_layer_norm_to_stack(residual_stack, ln_final)
        if recomputed is not _MISSING:
            return recomputed

    resolved_scale_key = scale_key
    candidates: list[ActivationKey] = []
    if resolved_scale_key is None:
        if target_layer == n_layers:
            candidates = ["ln_final.hook_scale"]
        else:
            requested_layer_norm = _residual_stack_ln_name(stack, mlp_input)
            fallback_layer_norm = _residual_stack_fallback_ln_name(stack, mlp_input)
            candidates = _layer_norm_scale_candidates(target_layer, requested_layer_norm)
            if fallback_layer_norm is not None and not requires_cached_scale:
                candidates.extend(
                    _layer_norm_scale_candidates(target_layer, fallback_layer_norm)
                )
        resolved_scale_key = next(
            (candidate for candidate in candidates if candidate in self),
            None,
        )
    if resolved_scale_key is None:
        if requires_cached_scale:
            expected_key = candidates[0] if scale_key is None else scale_key
            expected_name = (
                get_act_name(*expected_key)
                if isinstance(expected_key, tuple)
                else str(expected_key)
            )
            raise KeyError(
                f"Cached LN scale not found at {expected_name!r}. apply_ln operations "
                "require this hook to be cached for the requested layer."
            )
        return residual_stack
    scale = self[resolved_scale_key]
    if batch_slice is not None and resolved_has_batch_dim:
        scale = _slice_dim(scale, batch_slice, dim=0)
    if pos_slice is not None:
        pos_dim = _scale_pos_dim_for_residual_stack(
            scale,
            residual_stack,
            has_batch_dim=resolved_has_batch_dim,
        )
        scale = _slice_dim(scale, pos_slice, dim=pos_dim)
    if _uses_centered_layer_norm(self.model):
        residual_stack = _subtract_last_dim_mean(residual_stack)
    return _divide_values(residual_stack, scale)

apply_slice_to_batch_dim(batch_slice)

Return a cache sliced along the batch dimension.

Source code in src/SafeLens/core/hooks.py

def apply_slice_to_batch_dim(self, batch_slice: Any) -> ActivationCache:
    """Return a cache sliced along the batch dimension."""
    normalized_slice = _normalize_slice_index(batch_slice)
    if not self.has_batch_dim:
        if normalized_slice == _FULL_SLICE:
            return ActivationCache(
                dict(self._cache),
                model=self.model,
                has_batch_dim=False,
            )
        raise ValueError("Cannot slice batch dimension on a cache without batch dim.")
    has_batch_dim = not isinstance(normalized_slice, int)
    return ActivationCache(
        {
            name: _slice_dim(value, normalized_slice, dim=0)
            for name, value in self._cache.items()
        },
        model=self.model,
        has_batch_dim=has_batch_dim,
    )

apply_to_values(fn)

Apply a function to every cached value and return a new cache.

Source code in src/SafeLens/core/hooks.py

def apply_to_values(self, fn: Callable[[Any], Any]) -> ActivationCache:
    """Apply a function to every cached value and return a new cache."""
    return ActivationCache(
        {name: fn(value) for name, value in self._cache.items()},
        model=self.model,
        has_batch_dim=self.has_batch_dim,
    )

clone()

Return a cloned copy when activations support .clone(), otherwise deep-copy values.

Source code in src/SafeLens/core/hooks.py

def clone(self) -> ActivationCache:
    """Return a cloned copy when activations support `.clone()`, otherwise deep-copy values."""
    return ActivationCache(
        {name: clone_activation(value) for name, value in self._cache.items()},
        model=self.model,
        has_batch_dim=self.has_batch_dim,
    )

compute_head_results(layer=None, *, store=True, pos_slice=None, return_labels=False)

Compute per-head residual-space results from cached z and model W_O.

This fills the common TransformerLens workflow gap where a model exposes head outputs z and output weights W_O, but not cached result vectors directly.

Source code in src/SafeLens/core/hooks.py

def compute_head_results(
    self,
    layer: int | None = None,
    *,
    store: bool = True,
    pos_slice: Any = None,
    return_labels: bool = False,
) -> Any:
    """Compute per-head residual-space results from cached `z` and model `W_O`.

    This fills the common TransformerLens workflow gap where a model exposes
    head outputs `z` and output weights `W_O`, but not cached `result`
    vectors directly.
    """
    target_layer = self._normalize_layer(layer)
    n_layers = self._infer_n_layers()
    max_layer = min(target_layer, n_layers)
    values: list[Any] = []
    labels: list[str] = []

    for current_layer in range(max_layer):
        if ("result", current_layer) in self:
            result = _maybe_slice_pos(
                self[("result", current_layer)],
                pos_slice,
                dim=_head_vector_pos_dim("result"),
            )
            values.append(result)
            labels.append(f"{current_layer}_result")
            continue
        if ("z", current_layer) not in self:
            continue
        w_o = _get_layer_weight(self.model, "W_O", current_layer)
        if w_o is None:
            raise ValueError(
                "compute_head_results requires a cache with a model exposing W_O "
                "when cached head result activations are missing."
            )
        from SafeLens.core.analysis import compute_head_results_from_z

        z_activation = _maybe_slice_pos(self[("z", current_layer)], pos_slice, dim=-3)
        result = compute_head_results_from_z(z_activation, w_o)
        if store and pos_slice is None:
            self[f"layer_{current_layer}.result"] = result
            self.cache_dict[f"blocks.{current_layer}.attn.hook_result"] = result
        values.append(result)
        labels.append(f"{current_layer}_result")

    if not values:
        raise KeyError("No cached `z` activations found for head result computation.")
    result_stack = stack_values(values)
    if return_labels:
        return result_stack, labels
    return result_stack

cpu()

Move tensor-like activations to CPU when supported.

Source code in src/SafeLens/core/hooks.py

def cpu(self) -> ActivationCache:
    """Move tensor-like activations to CPU when supported."""
    return self.to("cpu")

decompose_resid(layer=None, mlp_input=False, mode='all', apply_ln=False, pos_slice=None, incl_embeds=True, return_labels=False, *, stack='encoder')

Decompose a residual stream into embedding, attention, and MLP components.

Source code in src/SafeLens/core/hooks.py

def decompose_resid(
    self,
    layer: int | None = None,
    mlp_input: bool = False,
    mode: Literal["all", "mlp", "attn"] = "all",
    apply_ln: bool = False,
    pos_slice: Any = None,
    incl_embeds: bool = True,
    return_labels: bool = False,
    *,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Decompose a residual stream into embedding, attention, and MLP components."""
    stack = _normalize_residual_stack_name(stack)
    target_layer = self._normalize_stack_layer(layer, stack=stack)
    values: list[Any] = []
    labels: list[str] = []
    include_attn = mode != "mlp"
    include_mlp = mode != "attn" and not _model_is_attn_only(self.model)

    for key, label in (("hook_embed", "embed"), ("hook_pos_embed", "pos_embed")):
        if incl_embeds and key in self:
            values.append(_maybe_slice_pos(self[key], pos_slice))
            labels.append(label)

    for current_layer in range(target_layer):
        attn_key = _residual_component_key("attn_out", current_layer, stack=stack)
        cross_attn_key = _residual_component_key("cross_attn_out", current_layer, stack=stack)
        mlp_key = _residual_component_key("mlp_out", current_layer, stack=stack)
        if include_attn and attn_key in self:
            values.append(_maybe_slice_pos(self[attn_key], pos_slice))
            labels.append(f"{current_layer}_attn_out")
        if include_attn and stack == "decoder" and cross_attn_key in self:
            values.append(_maybe_slice_pos(self[cross_attn_key], pos_slice))
            labels.append(f"{current_layer}_cross_attn_out")
        if include_mlp and mlp_key in self:
            values.append(_maybe_slice_pos(self[mlp_key], pos_slice))
            labels.append(f"{current_layer}_mlp_out")

    attn_key = _residual_component_key("attn_out", target_layer, stack=stack)
    cross_attn_key = _residual_component_key("cross_attn_out", target_layer, stack=stack)
    if mlp_input and include_attn and attn_key in self:
        values.append(_maybe_slice_pos(self[attn_key], pos_slice))
        labels.append(f"{target_layer}_attn_out")
    if mlp_input and include_attn and stack == "decoder" and cross_attn_key in self:
        values.append(_maybe_slice_pos(self[cross_attn_key], pos_slice))
        labels.append(f"{target_layer}_cross_attn_out")
    if not values:
        raise KeyError("No residual decomposition activations found in cache.")

    residual_stack = stack_values(values)
    if apply_ln:
        residual_stack = self.apply_ln_to_stack(
            residual_stack,
            layer=target_layer,
            mlp_input=mlp_input,
            pos_slice=pos_slice,
            has_batch_dim=self.has_batch_dim,
            stack=stack,
        )
    if return_labels:
        return residual_stack, labels
    return residual_stack

decompose_resid_to_logits(layer=None, mlp_input=False, mode='all', *, pos_slice=None, incl_embeds=True, return_labels=False, use_unembed_bias=False, stack='encoder')

Project residual-decomposition components through the unembedding.

Source code in src/SafeLens/core/hooks.py

def decompose_resid_to_logits(
    self,
    layer: int | None = None,
    mlp_input: bool = False,
    mode: Literal["all", "mlp", "attn"] = "all",
    *,
    pos_slice: Any = None,
    incl_embeds: bool = True,
    return_labels: bool = False,
    use_unembed_bias: bool = False,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Project residual-decomposition components through the unembedding."""
    stack = _normalize_residual_stack_name(stack)
    residual_stack, labels = self.decompose_resid(
        layer=layer,
        mlp_input=mlp_input,
        mode=mode,
        apply_ln=True,
        pos_slice=pos_slice,
        incl_embeds=incl_embeds,
        return_labels=True,
        stack=stack,
    )
    logits = self.residual_stack_to_logits(
        residual_stack,
        apply_ln=False,
        use_unembed_bias=use_unembed_bias,
    )
    if return_labels:
        return logits, labels
    return logits

detach()

Detach tensor-like activations when values support .detach().

Source code in src/SafeLens/core/hooks.py

def detach(self) -> ActivationCache:
    """Detach tensor-like activations when values support `.detach()`."""
    return self.apply_to_values(_detach_value)

get_activation(name)

Return one cached activation.

Source code in src/SafeLens/core/hooks.py

def get_activation(self, name: str) -> Any:
    """Return one cached activation."""
    return self[name]

get_full_resid_decomposition(layer=None, mlp_input=False, expand_neurons=True, apply_ln=False, pos_slice=None, return_labels=False, project_output_onto=None, *, stack='encoder')

Return a best-effort decomposition into heads, MLP neurons, embeds, and bias.

Source code in src/SafeLens/core/hooks.py

def get_full_resid_decomposition(
    self,
    layer: int | None = None,
    mlp_input: bool = False,
    expand_neurons: bool = True,
    apply_ln: bool = False,
    pos_slice: Any = None,
    return_labels: bool = False,
    project_output_onto: Any = None,
    *,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Return a best-effort decomposition into heads, MLP neurons, embeds, and bias."""
    stack = _normalize_residual_stack_name(stack)
    target_layer = self._normalize_stack_layer(layer, stack=stack)
    stacks: list[Any] = []
    labels: list[str] = []
    expanded_neurons = False
    ln_folded = apply_ln and project_output_onto is not None
    bias_reference: Any = None

    def remember_bias_reference(component_stack: Any) -> None:
        nonlocal bias_reference
        if bias_reference is not None:
            return
        component_rows = _unstack_first_dim(component_stack)
        if component_rows:
            bias_reference = component_rows[0]

    def maybe_ln_then_project(component_stack: Any) -> Any:
        if ln_folded:
            component_stack = self.apply_ln_to_stack(
                component_stack,
                layer=target_layer,
                mlp_input=mlp_input,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
                stack=stack,
            )
            return _project_last_dim(component_stack, project_output_onto)
        if project_output_onto is not None:
            return _project_last_dim(component_stack, project_output_onto)
        return component_stack

    def add_stack(component_stack: Any, component_labels: list[str]) -> None:
        stacks.extend(_unstack_first_dim(component_stack))
        labels.extend(component_labels)

    try:
        if stack == "decoder":
            head_stack, head_labels = _decoder_head_result_stack(
                self,
                target_layer + (1 if mlp_input else 0),
                pos_slice=pos_slice,
            )
        else:
            head_stack, head_labels = self.stack_head_results(
                target_layer + (1 if mlp_input else 0),
                pos_slice=pos_slice,
                return_labels=True,
            )
        remember_bias_reference(head_stack)
        head_stack = maybe_ln_then_project(head_stack)
        add_stack(head_stack, head_labels)
    except KeyError:
        try:
            attn_stack, attn_labels = self.decompose_resid(
                target_layer,
                mlp_input=mlp_input,
                mode="attn",
                incl_embeds=False,
                pos_slice=pos_slice,
                return_labels=True,
                stack=stack,
            )
            remember_bias_reference(attn_stack)
            attn_stack = maybe_ln_then_project(attn_stack)
            add_stack(attn_stack, attn_labels)
        except KeyError:
            pass

    if stack == "decoder" and not _model_is_attn_only(self.model):
        try:
            mlp_stack, mlp_labels = self.decompose_resid(
                target_layer,
                mode="mlp",
                incl_embeds=False,
                pos_slice=pos_slice,
                return_labels=True,
                stack=stack,
            )
            remember_bias_reference(mlp_stack)
            mlp_stack = maybe_ln_then_project(mlp_stack)
            add_stack(mlp_stack, mlp_labels)
        except KeyError:
            pass
    elif expand_neurons and not _model_is_attn_only(self.model):
        try:
            neuron_stack, neuron_labels = self.stack_neuron_results(
                target_layer,
                pos_slice=pos_slice,
                return_labels=True,
                require_output_weight=True,
                apply_ln=ln_folded,
                project_output_onto=project_output_onto,
            )
            if project_output_onto is None and not apply_ln:
                remember_bias_reference(neuron_stack)
            add_stack(neuron_stack, neuron_labels)
            expanded_neurons = True
        except (KeyError, ValueError):
            try:
                mlp_stack, mlp_labels = self.decompose_resid(
                    target_layer,
                    mode="mlp",
                    incl_embeds=False,
                    pos_slice=pos_slice,
                    return_labels=True,
                    stack=stack,
                )
                remember_bias_reference(mlp_stack)
                mlp_stack = maybe_ln_then_project(mlp_stack)
                add_stack(mlp_stack, mlp_labels)
            except KeyError:
                pass
    elif not _model_is_attn_only(self.model):
        try:
            mlp_stack, mlp_labels = self.decompose_resid(
                target_layer,
                mode="mlp",
                incl_embeds=False,
                pos_slice=pos_slice,
                return_labels=True,
                stack=stack,
            )
            remember_bias_reference(mlp_stack)
            mlp_stack = maybe_ln_then_project(mlp_stack)
            add_stack(mlp_stack, mlp_labels)
        except KeyError:
            pass

    for key, label in (("hook_embed", "embed"), ("hook_pos_embed", "pos_embed")):
        if key in self:
            embed_stack = stack_values([_maybe_slice_pos(self[key], pos_slice)])
            remember_bias_reference(embed_stack)
            embed_stack = maybe_ln_then_project(embed_stack)
            stacks.extend(_unstack_first_dim(embed_stack))
            labels.append(label)

    accumulated_bias = _get_model_attr(self.model, "accumulated_bias")
    if callable(accumulated_bias):
        try:
            bias = accumulated_bias(
                target_layer,
                mlp_input,
                include_mlp_biases=expanded_neurons,
            )
        except TypeError:
            try:
                bias = accumulated_bias(target_layer, mlp_input)
            except TypeError:
                bias = accumulated_bias(target_layer)
        if ln_folded:
            bias = _expand_bias_like_for_folded_projection(
                bias,
                bias_reference,
                stacks[0] if stacks else None,
                project_output_onto,
            )
            bias_stack = stack_values([bias])
            bias_stack = self.apply_ln_to_stack(
                bias_stack,
                layer=target_layer,
                mlp_input=mlp_input,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
                stack=stack,
            )
            bias_stack = _project_last_dim(bias_stack, project_output_onto)
            stacks.extend(_unstack_first_dim(bias_stack))
        else:
            if project_output_onto is not None:
                bias = _project_last_dim(bias, project_output_onto)
            bias = _expand_bias_like(bias, stacks[0] if stacks else None)
            stacks.append(bias)
        labels.append("bias")

    if not stacks:
        raise KeyError("No activations found for a full residual decomposition.")
    full_stack = stack_values(stacks)
    if apply_ln and not ln_folded:
        full_stack = self.apply_ln_to_stack(
            full_stack,
            layer=target_layer,
            mlp_input=mlp_input,
            pos_slice=pos_slice,
            has_batch_dim=self.has_batch_dim,
            stack=stack,
        )
    if return_labels:
        return full_stack, labels
    return full_stack

get_neuron_results(layer, neuron_slice=None, pos_slice=None, project_output_onto=None, *, component='post')

Return one layer's per-neuron residual contributions.

Source code in src/SafeLens/core/hooks.py

def get_neuron_results(
    self,
    layer: int,
    neuron_slice: Any = None,
    pos_slice: Any = None,
    project_output_onto: Any = None,
    *,
    component: str = "post",
) -> Any:
    """Return one layer's per-neuron residual contributions."""
    if (component, layer) not in self:
        raise KeyError(f"No cached {component!r} neuron activations found for layer {layer}.")
    w_out = _get_layer_weight(self.model, "W_out", layer)
    if w_out is None:
        raise ValueError("get_neuron_results requires a model exposing W_out.")

    neuron_acts = _maybe_slice_pos(self[(component, layer)], pos_slice)
    neuron_count = _infer_last_dim(neuron_acts)
    neuron_indices = _indices_from_slice(neuron_slice, neuron_count)
    neuron_acts = _select_indices_dim(neuron_acts, neuron_indices, dim=-1)
    layer_w_out = _select_indices_dim(w_out, neuron_indices, dim=0)
    if project_output_onto is not None:
        layer_w_out = _project_last_dim(layer_w_out, project_output_onto)
    return _multiply_last_dim_by_matrix(neuron_acts, layer_w_out)

items()

Return cached activation items, matching TransformerLens' mapping API.

Source code in src/SafeLens/core/hooks.py

def items(self) -> Any:
    """Return cached activation items, matching TransformerLens' mapping API."""
    return self._cache.items()

keys()

Return cached activation names, matching TransformerLens' mapping API.

Source code in src/SafeLens/core/hooks.py

def keys(self) -> Any:
    """Return cached activation names, matching TransformerLens' mapping API."""
    return self._cache.keys()

keys_matching(names_filter)

Return activation names matching a TransformerLens-style names filter.

Source code in src/SafeLens/core/hooks.py

def keys_matching(self, names_filter: NamesFilter) -> list[str]:
    """Return activation names matching a TransformerLens-style names filter."""
    return [name for name in self._cache if _cache_key_matches_filter(name, names_filter)]

logit_attrs(residual_stack, tokens, incorrect_tokens=None, pos_slice=None, batch_slice=None, has_batch_dim=None, *, directions=None, apply_ln=True)

Project residual components onto token residual directions.

Source code in src/SafeLens/core/hooks.py

def logit_attrs(
    self,
    residual_stack: Any,
    tokens: Any,
    incorrect_tokens: Any = None,
    pos_slice: Any = None,
    batch_slice: Any = None,
    has_batch_dim: bool | None = None,
    *,
    directions: Any = None,
    apply_ln: bool = True,
) -> Any:
    """Project residual components onto token residual directions."""
    resolved_has_batch_dim = self.has_batch_dim if has_batch_dim is None else has_batch_dim
    if directions is None:
        if self.model is None:
            directions = tokens
        else:
            directions = self.model.tokens_to_residual_directions(
                _normalize_logit_tokens(self.model, tokens)
            )
    if incorrect_tokens is not None:
        if self.model is None:
            incorrect_directions = incorrect_tokens
        else:
            incorrect_directions = self.model.tokens_to_residual_directions(
                _normalize_logit_tokens(self.model, incorrect_tokens)
            )
        if _shape_of(directions) != _shape_of(incorrect_directions):
            raise ValueError(
                "tokens and incorrect_tokens must resolve to residual directions with the "
                f"same shape, got {_shape_of(directions)!r} and "
                f"{_shape_of(incorrect_directions)!r}."
            )
        directions = _subtract_values(directions, incorrect_directions)
    batch_dim = _direction_batch_dim(
        directions,
        residual_stack,
        has_batch_dim=resolved_has_batch_dim,
        prefer_pos_axis=pos_slice is not None,
    )
    residual_batch_dim = _residual_stack_batch_dim(
        residual_stack,
        directions,
        pos_slice=pos_slice,
    )
    if batch_slice is not None and batch_dim is not None:
        directions = _slice_dim(directions, batch_slice, dim=batch_dim)
        direction_has_batch_dim = not isinstance(batch_slice, int)
    else:
        direction_has_batch_dim = resolved_has_batch_dim
    pos_dim = _direction_pos_dim(
        directions,
        residual_stack,
        has_batch_dim=direction_has_batch_dim,
        prefer_pos_axis=pos_slice is not None,
    )
    if pos_slice is not None and pos_dim is not None:
        directions = _slice_dim(directions, pos_slice, dim=pos_dim)
    if apply_ln:
        residual_stack = _slice_residual_stack_for_logit_attrs(
            residual_stack,
            pos_slice,
            batch_slice=batch_slice,
            has_batch_dim=resolved_has_batch_dim,
            residual_batch_dim=residual_batch_dim,
            reference=directions,
        )
        residual_stack = _apply_cached_ln_scale_to_sliced_stack(
            self,
            residual_stack,
            layer=-1,
            pos_slice=pos_slice,
            batch_slice=batch_slice,
        )
    elif batch_slice is not None and resolved_has_batch_dim:
        residual_stack = _slice_dim(
            residual_stack,
            batch_slice,
            dim=residual_batch_dim,
        )
    return _dot_last_dim(residual_stack, directions)

remove_batch_dim()

Remove singleton batch dimensions in place and return this cache.

Source code in src/SafeLens/core/hooks.py

def remove_batch_dim(self) -> ActivationCache:
    """Remove singleton batch dimensions in place and return this cache."""
    if not self.has_batch_dim:
        return self
    updated_values: dict[str, Any] = {}
    has_singleton_batch = any(_has_leading_dim(value, 1) for value in self._cache.values())
    for name, value in list(self._cache.items()):
        if _has_leading_dim(value, 1):
            updated_values[name] = _slice_dim(value, 0, dim=0)
            continue
        shape = _shape_of(value)
        if shape and not has_singleton_batch:
            raise ValueError(
                f"Cannot remove batch dimension from cache with batch size > 1, "
                f"for key {name} with shape {shape!r}."
            )
        updated_values[name] = value
    self._cache.update(updated_values)
    self.has_batch_dim = False
    return self

residual_stack_to_logits(residual_stack, *, apply_ln=True, pos_slice=None, batch_slice=None, has_batch_dim=None, use_unembed_bias=True, stack='encoder')

Project residual components through the model unembedding matrix.

Source code in src/SafeLens/core/hooks.py

def residual_stack_to_logits(
    self,
    residual_stack: Any,
    *,
    apply_ln: bool = True,
    pos_slice: Any = None,
    batch_slice: Any = None,
    has_batch_dim: bool | None = None,
    use_unembed_bias: bool = True,
    stack: Literal["encoder", "decoder"] = "encoder",
) -> Any:
    """Project residual components through the model unembedding matrix."""
    stack = _normalize_residual_stack_name(stack)
    if self.model is None:
        raise ValueError("residual_stack_to_logits requires a cache with an attached model.")
    unembed = _get_model_attr(self.model, "W_U")
    if unembed is None:
        raise ValueError("residual_stack_to_logits requires a model exposing W_U.")
    unembed_bias = _get_model_attr(self.model, "b_U") if use_unembed_bias else None
    resolved_has_batch_dim = self.has_batch_dim if has_batch_dim is None else has_batch_dim
    residual_batch_dim = _residual_stack_batch_dim(residual_stack, None, pos_slice=pos_slice)
    if apply_ln:
        residual_stack = _maybe_slice_residual_stack_pos(
            residual_stack,
            pos_slice,
            has_batch_dim=resolved_has_batch_dim,
        )
        residual_stack = self.apply_ln_to_stack(
            residual_stack,
            layer=-1,
            pos_slice=pos_slice,
            batch_slice=batch_slice,
            has_batch_dim=resolved_has_batch_dim,
            stack=stack,
        )
    else:
        if batch_slice is not None and resolved_has_batch_dim:
            residual_stack = _slice_dim(residual_stack, batch_slice, dim=residual_batch_dim)
        if pos_slice is not None:
            residual_stack = _slice_dim(
                residual_stack,
                pos_slice,
                dim=_residual_stack_pos_dim_after_optional_batch_slice(
                    residual_stack,
                    has_batch_dim=resolved_has_batch_dim,
                    batch_slice=batch_slice,
                ),
            )
    from SafeLens.core.analysis import residual_stack_to_logits

    return residual_stack_to_logits(residual_stack, unembed, unembed_bias)

resolve_key(key)

Resolve exact, SafeLens-style, or TransformerLens-style activation keys.

Source code in src/SafeLens/core/hooks.py

def resolve_key(self, key: ActivationKey) -> str:
    """Resolve exact, SafeLens-style, or TransformerLens-style activation keys."""
    candidates = activation_name_candidates(
        key,
        n_layers=self._infer_n_layers(),
        decoder_n_layers=self._infer_stack_n_layers("decoder"),
    )
    for candidate in candidates:
        if candidate in self._cache:
            return candidate
    raise KeyError(f"Unknown activation key {key!r}. Tried {candidates!r}.")

select(names_filter)

Return a new cache containing only matching activation names.

Source code in src/SafeLens/core/hooks.py

def select(self, names_filter: NamesFilter) -> ActivationCache:
    """Return a new cache containing only matching activation names."""
    return ActivationCache(
        {
            name: value
            for name, value in self._cache.items()
            if _cache_key_matches_filter(name, names_filter)
        },
        model=self.model,
        has_batch_dim=self.has_batch_dim,
    )

stack_activation(activation_name, layer=None, layer_type=None, *, sublayer_type=None)

Stack one activation across layers.

Source code in src/SafeLens/core/hooks.py

def stack_activation(
    self,
    activation_name: str,
    layer: int | None = None,
    layer_type: str | None = None,
    *,
    sublayer_type: str | None = None,
) -> Any:
    """Stack one activation across layers."""
    if sublayer_type is not None:
        layer_type = sublayer_type
    n_layers = self._normalize_layer(layer)
    values = [
        self[(activation_name, current_layer, layer_type)] for current_layer in range(n_layers)
    ]
    return stack_values(values)

stack_head_results(layer=None, return_labels=False, incl_remainder=False, pos_slice=None, apply_ln=False, *, component='result')

Stack per-head activations from [batch, pos, head, d_model] caches.

Source code in src/SafeLens/core/hooks.py

def stack_head_results(
    self,
    layer: int | None = None,
    return_labels: bool = False,
    incl_remainder: bool = False,
    pos_slice: Any = None,
    apply_ln: bool = False,
    *,
    component: str = "result",
) -> Any:
    """Stack per-head activations from `[batch, pos, head, d_model]` caches."""
    if incl_remainder and component != "result":
        raise ValueError(
            "incl_remainder=True requires residual-space `result` head activations."
        )
    target_layer = self._normalize_layer(layer)
    if component == "result" and any(
        ("z", current_layer) in self and (component, current_layer) not in self
        for current_layer in range(target_layer)
    ):
        try:
            self.compute_head_results(target_layer, store=True)
        except (KeyError, ValueError):
            pass
    values: list[Any] = []
    labels: list[str] = []
    for current_layer in range(target_layer):
        if (component, current_layer) not in self:
            continue
        activation = _maybe_slice_pos(
            self[(component, current_layer)],
            pos_slice,
            dim=_head_vector_pos_dim(component),
        )
        head_dim = _head_axis_after_pos_slice(activation, pos_slice, component=component)
        for head_index in range(_infer_head_count(activation, dim=head_dim)):
            values.append(_slice_dim(activation, head_index, dim=head_dim))
            labels.append(f"L{current_layer}H{head_index}")
    if incl_remainder:
        remainder = _residual_remainder_base(self, target_layer, pos_slice)
        if values:
            remainder = _subtract_values(remainder, _sum_values(values))
        values.append(remainder)
        labels.append("remainder")
    if not values:
        if target_layer == 0:
            head_stack = _empty_component_stack_like_cache(
                self,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
            )
            if apply_ln:
                head_stack = self.apply_ln_to_stack(
                    head_stack,
                    layer=target_layer,
                    pos_slice=pos_slice,
                    has_batch_dim=self.has_batch_dim,
                )
            if return_labels:
                return head_stack, labels
            return head_stack
        raise KeyError(f"No {component!r} head activations found in cache.")
    head_stack = stack_values(values)
    if apply_ln:
        head_stack = self.apply_ln_to_stack(
            head_stack,
            layer=target_layer,
            pos_slice=pos_slice,
            has_batch_dim=self.has_batch_dim,
        )
    if return_labels:
        return head_stack, labels
    return head_stack

stack_neuron_results(layer=None, pos_slice=None, neuron_slice=None, return_labels=False, incl_remainder=False, apply_ln=False, project_output_onto=None, *, component='post', require_output_weight=False)

Stack per-neuron MLP residual contributions when W_out is available.

Source code in src/SafeLens/core/hooks.py

def stack_neuron_results(
    self,
    layer: int | None = None,
    pos_slice: Any = None,
    neuron_slice: Any = None,
    return_labels: bool = False,
    incl_remainder: bool = False,
    apply_ln: bool = False,
    project_output_onto: Any = None,
    *,
    component: str = "post",
    require_output_weight: bool = False,
) -> Any:
    """Stack per-neuron MLP residual contributions when `W_out` is available."""
    target_layer = self._normalize_layer(layer)
    values: list[Any] = []
    labels: list[str] = []
    can_project_before_stack = (
        project_output_onto is not None and not apply_ln and not incl_remainder
    )
    can_fold_ln_projection = (
        component == "post"
        and project_output_onto is not None
        and apply_ln
        and not incl_remainder
        and self._can_fold_ln_neuron_projection(target_layer, pos_slice)
    )
    results_are_projected = False
    for current_layer in range(target_layer):
        if (component, current_layer) not in self:
            continue
        activation = _maybe_slice_pos(self[(component, current_layer)], pos_slice)
        neuron_indices = _indices_from_slice(neuron_slice, _infer_last_dim(activation))
        if can_fold_ln_projection:
            labels.extend(f"L{current_layer}N{neuron_index}" for neuron_index in neuron_indices)
            continue
        try:
            layer_results = self.get_neuron_results(
                current_layer,
                neuron_slice=neuron_indices,
                pos_slice=pos_slice,
                project_output_onto=project_output_onto if can_project_before_stack else None,
                component=component,
            )
            results_are_projected = can_project_before_stack
            neuron_dim = (
                -1 if results_are_projected and len(_shape_of(project_output_onto)) == 1 else -2
            )
        except ValueError:
            if project_output_onto is not None or require_output_weight or incl_remainder:
                raise
            layer_results = _select_indices_dim(activation, neuron_indices, dim=-1)
            neuron_dim = -1
            results_are_projected = False
        for position, neuron_index in enumerate(neuron_indices):
            neuron_value = _slice_dim(layer_results, position, dim=neuron_dim)
            values.append(neuron_value)
            labels.append(f"L{current_layer}N{neuron_index}")
    if incl_remainder:
        remainder = _residual_remainder_base(self, target_layer, pos_slice)
        if values:
            remainder = _subtract_values(remainder, _sum_values(values))
        values.append(remainder)
        labels.append("remainder")
    if not values and not (can_fold_ln_projection and labels):
        if target_layer == 0:
            neuron_stack = _empty_component_stack_like_cache(
                self,
                pos_slice=pos_slice,
                has_batch_dim=self.has_batch_dim,
                project_output_onto=project_output_onto,
            )
            if apply_ln:
                neuron_stack = self.apply_ln_to_stack(
                    neuron_stack,
                    layer=target_layer,
                    pos_slice=pos_slice,
                    has_batch_dim=self.has_batch_dim,
                )
            if return_labels:
                return neuron_stack, labels
            return neuron_stack
        raise KeyError(f"No {component!r} neuron activations found in cache.")
    if can_fold_ln_projection:
        neuron_stack = self._stack_neuron_results_apply_ln_projected(
            target_layer,
            pos_slice,
            neuron_slice,
            project_output_onto,
        )
        results_are_projected = True
    else:
        neuron_stack = stack_values(values)
    if apply_ln and not can_fold_ln_projection:
        neuron_stack = self.apply_ln_to_stack(
            neuron_stack,
            layer=target_layer,
            pos_slice=pos_slice,
            has_batch_dim=self.has_batch_dim,
        )
    if project_output_onto is not None and not results_are_projected:
        neuron_stack = _project_last_dim(neuron_stack, project_output_onto)
    if return_labels:
        return neuron_stack, labels
    return neuron_stack

storage_key(key)

Return the existing or canonical storage name for an activation key.

Source code in src/SafeLens/core/hooks.py

def storage_key(self, key: ActivationKey) -> str:
    """Return the existing or canonical storage name for an activation key."""
    try:
        return self.resolve_key(key)
    except KeyError:
        if isinstance(key, tuple) and key:
            tuple_key = list(key)
            if len(tuple_key) == 1:
                top_level_name = _TOP_LEVEL_ACT_NAMES.get(str(tuple_key[0]))
                if top_level_name is not None:
                    return top_level_name
            layer = tuple_key[1] if len(tuple_key) >= 2 else None
            raw_name = _strip_hook_prefix(str(tuple_key[0]))
            name = _ACT_NAME_ALIASES.get(raw_name, raw_name)
            if layer == -1:
                stack = _activation_key_stack_name(name)
                n_layers = self._infer_stack_n_layers(stack)
                if n_layers > 0:
                    layer = n_layers - 1
            layer_type = (
                str(tuple_key[2]) if len(tuple_key) >= 3 and tuple_key[2] is not None else None
            )
            if layer_type is not None:
                layer_type = _LAYER_TYPE_ALIASES.get(layer_type, layer_type)
            if layer_type and layer is not None:
                return f"{activation_name_for_layer(layer)}.{layer_type}.{name}"
            return safelens_act_name(name, layer)
        if isinstance(key, str):
            canonical_key = _canonical_storage_key_for_string(key)
            if canonical_key != key:
                return canonical_key
        candidates = activation_name_candidates(
            key,
            n_layers=self._infer_n_layers(),
            decoder_n_layers=self._infer_stack_n_layers("decoder"),
        )
        if candidates:
            return candidates[0]
        return str(key)

store(name, activation, *, detach=True, clone=False, device=None)

Store an activation, optionally detaching, cloning, or moving it.

Source code in src/SafeLens/core/hooks.py

def store(
    self,
    name: str,
    activation: Any,
    *,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
) -> None:
    """Store an activation, optionally detaching, cloning, or moving it."""
    self._cache[name] = prepare_activation_for_cache(
        activation,
        detach=detach,
        clone=clone,
        device=device,
    )

to(device, move_model=None)

Move tensor-like activations to a device when values support .to().

Source code in src/SafeLens/core/hooks.py

def to(self, device: Any, move_model: bool | None = None) -> ActivationCache:
    """Move tensor-like activations to a device when values support `.to()`."""
    self._cache = {name: _move_value(value, device) for name, value in self._cache.items()}
    if move_model:
        model_to = getattr(self.model, "to", None)
        if callable(model_to):
            model_to(device)
    return self

to_dict()

Return a plain dictionary view copy.

Source code in src/SafeLens/core/hooks.py

def to_dict(self) -> dict[str, Any]:
    """Return a plain dictionary view copy."""
    return dict(self._cache)

toggle_autodiff(mode=False)

Set PyTorch's global grad-enabled state when PyTorch is available.

Source code in src/SafeLens/core/hooks.py

def toggle_autodiff(self, mode: bool = False) -> None:
    """Set PyTorch's global grad-enabled state when PyTorch is available."""
    try:
        import torch
    except ImportError:
        return None
    torch.set_grad_enabled(mode)
    return None

values()

Return cached activation values, matching TransformerLens' mapping API.

Source code in src/SafeLens/core/hooks.py

def values(self) -> Any:
    """Return cached activation values, matching TransformerLens' mapping API."""
    return self._cache.values()

HookPoint

Dependency-free identity hook point for instrumenting custom SafeLens models.

Source code in src/SafeLens/core/hooks.py

class HookPoint:
    """Dependency-free identity hook point for instrumenting custom SafeLens models."""

    def __init__(self, name: str | None = None) -> None:
        self.name = name
        self.ctx: dict[str, Any] = {}
        self.fwd_hooks: list[tuple[HookFn, LensHandle]] = []
        self.bwd_hooks: list[tuple[HookFn, LensHandle]] = []
        self.hook_conversion: Any = None
        self.backward_scale: float = 1.0

    def __repr__(self) -> str:
        """Return a TransformerLens-style hook summary."""
        bits = [f"name={self.name!r}"] if self.name is not None else []
        if self.fwd_hooks:
            bits.append(f"{len(self.fwd_hooks)} fwd")
        if self.bwd_hooks:
            bits.append(f"{len(self.bwd_hooks)} bwd")
        return f"HookPoint({', '.join(bits)})" if bits else "HookPoint()"

    def __call__(self, activation: Any) -> Any:
        """Run forward hooks over an activation and return the final value."""
        output = activation
        for hook_fn, _handle in list(self.fwd_hooks):
            hook_input = self.hook_conversion.convert(output) if self.hook_conversion else output
            patched = _call_hookpoint_fn(hook_fn, hook_input, self)
            if patched is not None:
                output = self.hook_conversion.revert(patched) if self.hook_conversion else patched
        if self.bwd_hooks:
            output = self._register_backward_hooks(output)
        return output

    def forward(self, activation: Any) -> Any:
        """TransformerLens-style module entrypoint."""
        return self(activation)

    def add_perma_hook(
        self,
        hook_fn: HookFn | None = None,
        dir: Literal["fwd", "bwd"] = "fwd",
        *,
        hook: HookFn | None = None,
    ) -> LensHandle:
        """Register a permanent hook."""
        resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
        return self.add_hook(resolved_hook, dir=dir, is_permanent=True)

    def add_hook(
        self,
        hook_fn: HookFn | None = None,
        dir: Literal["fwd", "bwd"] = "fwd",
        *,
        hook: HookFn | None = None,
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
        alias_names: Sequence[str] | None = None,
    ) -> LensHandle:
        """Register a hook and return a removable handle."""
        hook_fn = _resolve_hook_argument(hook_fn, hook=hook)
        user_hook = hook_fn
        if alias_names is not None:
            hook_fn = _alias_hook_fn(hook_fn, self, alias_names)
        hook_list = self._hook_list(dir)
        record: tuple[HookFn, LensHandle]

        def remove_record() -> None:
            if record in hook_list:
                hook_list.remove(record)

        handle = LensHandle(
            remove_record,
            is_permanent=is_permanent,
            level=level,
            user_hook=user_hook,
            is_cache=bool(getattr(user_hook, "_safelens_is_cache_hook", False)),
        )
        record = (hook_fn, handle)
        if prepend:
            hook_list.insert(0, record)
        else:
            hook_list.append(record)
        return handle

    def has_hooks(
        self,
        dir: HookDirection = "both",
        *,
        including_permanent: bool = True,
        level: int | None = None,
    ) -> bool:
        """Return whether matching hooks are registered."""
        return any(
            _handle_matches(handle, including_permanent=including_permanent, level=level)
            for _hook_fn, handle in self._matching_hook_records(dir)
        )

    def remove_hooks(
        self,
        dir: HookDirection = "fwd",
        *,
        including_permanent: bool = False,
        level: int | None = None,
    ) -> None:
        """Remove hooks matching permanence and context-level filters."""
        for _hook_fn, handle in list(self._matching_hook_records(dir)):
            if _handle_matches(handle, including_permanent=including_permanent, level=level):
                handle.remove()

    def clear_context(self) -> None:
        """Clear this hook point's mutable context dictionary."""
        self.ctx.clear()

    def enable_reshape(self, hook_conversion: Any = None) -> None:
        """Set an optional conversion applied around user hooks."""
        self.hook_conversion = hook_conversion

    def layer(self) -> int:
        """Extract the layer index from names like `blocks.3.attn.hook_q`."""
        if self.name is None:
            raise ValueError("Cannot infer layer from an unnamed HookPoint.")
        match = re.search(r"(?:blocks\.|layer_)(\d+)", self.name)
        if match is None:
            raise ValueError(f"Cannot infer layer from hook name {self.name!r}.")
        return int(match.group(1))

    def _hook_list(self, dir: Literal["fwd", "bwd"]) -> list[tuple[HookFn, LensHandle]]:
        if dir == "fwd":
            return self.fwd_hooks
        if dir == "bwd":
            return self.bwd_hooks
        raise ValueError(f"Invalid hook direction {dir!r}.")

    def _matching_hook_records(self, dir: HookDirection) -> list[tuple[HookFn, LensHandle]]:
        if dir == "fwd":
            return list(self.fwd_hooks)
        if dir == "bwd":
            return list(self.bwd_hooks)
        if dir == "both":
            return [*self.fwd_hooks, *self.bwd_hooks]
        raise ValueError(f"Invalid hook direction {dir!r}.")

    def _register_backward_hooks(self, activation: Any) -> Any:
        register_hook = getattr(activation, "register_hook", None)
        if not callable(register_hook):
            return activation

        def backward_hook(grad: Any) -> Any:
            output_grad = grad
            for hook_fn, _handle in list(self.bwd_hooks):
                hook_grad = _scale_backward_gradient_for_hooks(output_grad, self.backward_scale)
                if self.hook_conversion is not None:
                    hook_grad = self.hook_conversion.convert(hook_grad)
                patched = _call_hookpoint_fn(hook_fn, hook_grad, self)
                if patched is not None:
                    if self.hook_conversion is not None:
                        patched = self.hook_conversion.revert(patched)
                    output_grad = _unwrap_scaled_gradient_value(patched)
            return output_grad

        try:
            register_hook(backward_hook)
        except RuntimeError:
            return activation
        return activation

__call__(activation)

Run forward hooks over an activation and return the final value.

Source code in src/SafeLens/core/hooks.py

def __call__(self, activation: Any) -> Any:
    """Run forward hooks over an activation and return the final value."""
    output = activation
    for hook_fn, _handle in list(self.fwd_hooks):
        hook_input = self.hook_conversion.convert(output) if self.hook_conversion else output
        patched = _call_hookpoint_fn(hook_fn, hook_input, self)
        if patched is not None:
            output = self.hook_conversion.revert(patched) if self.hook_conversion else patched
    if self.bwd_hooks:
        output = self._register_backward_hooks(output)
    return output

__repr__()

Return a TransformerLens-style hook summary.

Source code in src/SafeLens/core/hooks.py

def __repr__(self) -> str:
    """Return a TransformerLens-style hook summary."""
    bits = [f"name={self.name!r}"] if self.name is not None else []
    if self.fwd_hooks:
        bits.append(f"{len(self.fwd_hooks)} fwd")
    if self.bwd_hooks:
        bits.append(f"{len(self.bwd_hooks)} bwd")
    return f"HookPoint({', '.join(bits)})" if bits else "HookPoint()"

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

Register a hook and return a removable handle.

Source code in src/SafeLens/core/hooks.py

def add_hook(
    self,
    hook_fn: HookFn | None = None,
    dir: Literal["fwd", "bwd"] = "fwd",
    *,
    hook: HookFn | None = None,
    is_permanent: bool = False,
    level: int | None = None,
    prepend: bool = False,
    alias_names: Sequence[str] | None = None,
) -> LensHandle:
    """Register a hook and return a removable handle."""
    hook_fn = _resolve_hook_argument(hook_fn, hook=hook)
    user_hook = hook_fn
    if alias_names is not None:
        hook_fn = _alias_hook_fn(hook_fn, self, alias_names)
    hook_list = self._hook_list(dir)
    record: tuple[HookFn, LensHandle]

    def remove_record() -> None:
        if record in hook_list:
            hook_list.remove(record)

    handle = LensHandle(
        remove_record,
        is_permanent=is_permanent,
        level=level,
        user_hook=user_hook,
        is_cache=bool(getattr(user_hook, "_safelens_is_cache_hook", False)),
    )
    record = (hook_fn, handle)
    if prepend:
        hook_list.insert(0, record)
    else:
        hook_list.append(record)
    return handle

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

Register a permanent hook.

Source code in src/SafeLens/core/hooks.py

def add_perma_hook(
    self,
    hook_fn: HookFn | None = None,
    dir: Literal["fwd", "bwd"] = "fwd",
    *,
    hook: HookFn | None = None,
) -> LensHandle:
    """Register a permanent hook."""
    resolved_hook = _resolve_hook_argument(hook_fn, hook=hook)
    return self.add_hook(resolved_hook, dir=dir, is_permanent=True)

clear_context()

Clear this hook point's mutable context dictionary.

Source code in src/SafeLens/core/hooks.py

def clear_context(self) -> None:
    """Clear this hook point's mutable context dictionary."""
    self.ctx.clear()

enable_reshape(hook_conversion=None)

Set an optional conversion applied around user hooks.

Source code in src/SafeLens/core/hooks.py

def enable_reshape(self, hook_conversion: Any = None) -> None:
    """Set an optional conversion applied around user hooks."""
    self.hook_conversion = hook_conversion

forward(activation)

TransformerLens-style module entrypoint.

Source code in src/SafeLens/core/hooks.py

def forward(self, activation: Any) -> Any:
    """TransformerLens-style module entrypoint."""
    return self(activation)

has_hooks(dir='both', *, including_permanent=True, level=None)

Return whether matching hooks are registered.

Source code in src/SafeLens/core/hooks.py

def has_hooks(
    self,
    dir: HookDirection = "both",
    *,
    including_permanent: bool = True,
    level: int | None = None,
) -> bool:
    """Return whether matching hooks are registered."""
    return any(
        _handle_matches(handle, including_permanent=including_permanent, level=level)
        for _hook_fn, handle in self._matching_hook_records(dir)
    )

layer()

Extract the layer index from names like blocks.3.attn.hook_q.

Source code in src/SafeLens/core/hooks.py

def layer(self) -> int:
    """Extract the layer index from names like `blocks.3.attn.hook_q`."""
    if self.name is None:
        raise ValueError("Cannot infer layer from an unnamed HookPoint.")
    match = re.search(r"(?:blocks\.|layer_)(\d+)", self.name)
    if match is None:
        raise ValueError(f"Cannot infer layer from hook name {self.name!r}.")
    return int(match.group(1))

remove_hooks(dir='fwd', *, including_permanent=False, level=None)

Remove hooks matching permanence and context-level filters.

Source code in src/SafeLens/core/hooks.py

def remove_hooks(
    self,
    dir: HookDirection = "fwd",
    *,
    including_permanent: bool = False,
    level: int | None = None,
) -> None:
    """Remove hooks matching permanence and context-level filters."""
    for _hook_fn, handle in list(self._matching_hook_records(dir)):
        if _handle_matches(handle, including_permanent=including_permanent, level=level):
            handle.remove()

LensHandle dataclass

Small removable hook handle with permanence and context-level metadata.

Source code in src/SafeLens/core/hooks.py

@dataclass
class LensHandle:
    """Small removable hook handle with permanence and context-level metadata."""

    remove_fn: Callable[[], None]
    is_permanent: bool = False
    level: int | None = None
    user_hook: Callable[..., Any] | None = None
    is_cache: bool = False
    removed: bool = False

    @property
    def context_level(self) -> int | None:
        """TransformerLens-compatible alias for the hook context level."""
        return self.level

    @context_level.setter
    def context_level(self, value: int | None) -> None:
        self.level = value

    @property
    def hook(self) -> LensHandle:
        """TransformerLens-compatible removable-handle reference."""
        return self

    def remove(self) -> None:
        """Remove this hook once."""
        if not self.removed:
            self.remove_fn()
            self.removed = True

context_level property writable

TransformerLens-compatible alias for the hook context level.

hook property

TransformerLens-compatible removable-handle reference.

remove()

Remove this hook once.

Source code in src/SafeLens/core/hooks.py

def remove(self) -> None:
    """Remove this hook once."""
    if not self.removed:
        self.remove_fn()
        self.removed = True

RemovableHandle

Bases: Protocol

Protocol for PyTorch-style removable hook handles.

Source code in src/SafeLens/core/hooks.py

class RemovableHandle(Protocol):
    """Protocol for PyTorch-style removable hook handles."""

    def remove(self) -> None:
        """Remove a registered hook."""

remove()

Remove a registered hook.

Source code in src/SafeLens/core/hooks.py

def remove(self) -> None:
    """Remove a registered hook."""

activation_name_candidates(key, *, n_layers=0, decoder_n_layers=None)

Return possible cache names for an activation key.

Source code in src/SafeLens/core/hooks.py

def activation_name_candidates(
    key: ActivationKey,
    *,
    n_layers: int = 0,
    decoder_n_layers: int | None = None,
) -> list[str]:
    """Return possible cache names for an activation key."""
    if isinstance(key, str):
        candidates = [key]
        top_level_name = _TOP_LEVEL_ACT_NAMES.get(key)
        if top_level_name is not None:
            candidates.append(top_level_name)
        candidates.extend(_safelens_names_from_transformer_lens_name(key))
        candidates.extend(_transformer_lens_names_from_safelens_name(key))
        tl_name = get_act_name(key)
        candidates.append(tl_name)
        candidates.extend(_safelens_names_from_transformer_lens_name(tl_name))
        if "." not in key:
            candidates.append(safelens_act_name(key))
        return _unique(candidates)

    if not key:
        return []

    raw_name = str(key[0])
    name = _strip_hook_prefix(raw_name)
    if len(key) == 1:
        top_level_name = _TOP_LEVEL_ACT_NAMES.get(raw_name) or _TOP_LEVEL_ACT_NAMES.get(name)
        if top_level_name is not None:
            return _unique([raw_name, name, top_level_name])

    original_layer = key[1] if len(key) >= 2 else None
    layer = original_layer
    layer_type = str(key[2]) if len(key) >= 3 and key[2] is not None else None
    aliased_name = _ACT_NAME_ALIASES.get(name, name)
    aliased_layer_type = (
        _LAYER_TYPE_ALIASES.get(layer_type, layer_type) if layer_type is not None else None
    )
    if layer == -1:
        stack_n_layers = (
            decoder_n_layers if _activation_key_stack_name(aliased_name) == "decoder" else n_layers
        )
        if stack_n_layers:
            layer = stack_n_layers - 1

    normalized_key = (name, layer, *key[2:])
    candidates = [".".join(str(item) for item in normalized_key if item is not None)]
    if raw_name != name or original_layer != layer:
        candidates.append(".".join(str(item) for item in key if item is not None))
    candidates.append(get_act_name(aliased_name, layer, layer_type))
    if layer is not None:
        candidates.append(safelens_act_name(aliased_name, layer))
        prefixed_attention = _prefixed_attention_parts(aliased_name)
        if prefixed_attention is not None:
            prefix, base_component = prefixed_attention
            tl_layer_type = "attn" if prefix == "decoder" else "cross_attn"
            compat_layer_type = "decoder_attn" if prefix == "decoder" else "cross_attn"
            candidates.extend(
                [
                    f"decoder.{layer}.{tl_layer_type}.hook_{base_component}",
                    f"blocks.{layer}.{compat_layer_type}.hook_{base_component}",
                    f"{activation_name_for_layer(layer)}.{compat_layer_type}.{base_component}",
                ]
            )
        decoder_top_level = _decoder_top_level_base(aliased_name)
        if decoder_top_level is not None:
            candidates.extend(
                [
                    f"decoder.{layer}.hook_{decoder_top_level}",
                    f"blocks.{layer}.hook_{aliased_name}",
                ]
            )
        elif aliased_name in _CROSS_TOP_LEVEL_ACTS:
            candidates.extend(
                [
                    f"decoder.{layer}.hook_{aliased_name}",
                    f"blocks.{layer}.hook_{aliased_name}",
                ]
            )
        decoder_mlp_base = _decoder_mlp_base(aliased_name)
        if decoder_mlp_base is not None:
            candidates.extend(
                [
                    f"decoder.{layer}.mlp.hook_{decoder_mlp_base}",
                    f"blocks.{layer}.decoder_mlp.hook_{decoder_mlp_base}",
                    f"{activation_name_for_layer(layer)}.decoder_mlp.{decoder_mlp_base}",
                ]
            )
        decoder_ln = _decoder_ln_parts(aliased_name)
        if decoder_ln is not None:
            ln_layer_type, ln_component = decoder_ln
            candidates.extend(
                [
                    f"decoder.{layer}.{ln_layer_type}.hook_{ln_component}",
                    f"blocks.{layer}.{ln_layer_type}.hook_{ln_component}",
                    f"{activation_name_for_layer(layer)}.{ln_layer_type}.{ln_component}",
                ]
            )
        if layer_type is None:
            if aliased_name in _ENCODER_TOP_LEVEL_ACTS:
                candidates.extend(
                    [
                        f"encoder.{layer}.hook_{aliased_name}",
                        f"blocks.{layer}.hook_{aliased_name}",
                    ]
                )
            elif aliased_name in _ATTN_ACTS:
                candidates.extend(
                    [
                        f"encoder.{layer}.attn.hook_{aliased_name}",
                        f"{activation_name_for_layer(layer)}.attn.{aliased_name}",
                        f"blocks.{layer}.attn.hook_{aliased_name}",
                    ]
                )
            elif aliased_name in _MLP_ACTS:
                candidates.extend(
                    [
                        f"encoder.{layer}.mlp.hook_{aliased_name}",
                        f"{activation_name_for_layer(layer)}.mlp.{aliased_name}",
                        f"blocks.{layer}.mlp.hook_{aliased_name}",
                    ]
                )
        if aliased_layer_type:
            stack_candidates = []
            if aliased_layer_type in {"attn", "mlp", "ln1", "ln2"}:
                stack_candidates.append(f"encoder.{layer}.{aliased_layer_type}.hook_{aliased_name}")
            candidates.extend(
                [
                    *stack_candidates,
                    f"{activation_name_for_layer(layer)}.{aliased_layer_type}.{aliased_name}",
                    f"blocks.{layer}.{aliased_layer_type}.hook_{aliased_name}",
                ]
            )
    return _unique(candidates)

activation_name_for_layer(layer)

Return the canonical cache name for a layer reference.

Source code in src/SafeLens/core/hooks.py

def activation_name_for_layer(layer: LayerRef) -> str:
    """Return the canonical cache name for a layer reference."""
    if isinstance(layer, int):
        return f"layer_{layer}"
    return str(layer)

cache_activations(model, batch, layers, *, names_filter=None, detach=True, clone=False, device=None)

Run a model while caching selected layer activations through temporary hooks.

Source code in src/SafeLens/core/hooks.py

def cache_activations(
    model: ModelWrapper,
    batch: Batch,
    layers: Sequence[LayerRef],
    *,
    names_filter: NamesFilter = None,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
) -> tuple[Any, ActivationCache]:
    """Run a model while caching selected layer activations through temporary hooks."""
    cache = ActivationCache()
    hooks = []
    for layer in layers:
        name = activation_name_for_layer(layer)
        if matches_names_filter(name, names_filter):
            hooks.append(
                (
                    layer,
                    make_cache_hook(cache, name, detach=detach, clone=clone, device=device),
                )
            )
    output, _ = run_with_hooks(model, batch, hooks)
    return output, cache

clone_activation(activation)

Clone tensor-like values and deep-copy everything else.

Source code in src/SafeLens/core/hooks.py

def clone_activation(activation: Any) -> Any:
    """Clone tensor-like values and deep-copy everything else."""
    clone = getattr(activation, "clone", None)
    if callable(clone):
        return clone()
    return deepcopy(activation)

extract_hook_output(args, kwargs)

Extract an activation from either PyTorch-style or keyword-style hook calls.

Source code in src/SafeLens/core/hooks.py

def extract_hook_output(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
    """Extract an activation from either PyTorch-style or keyword-style hook calls."""
    if len(args) >= 3:
        return args[2]
    if len(args) == 2 and _looks_like_hook_context(args[1]):
        return args[0]
    if "output" in kwargs:
        return kwargs["output"]
    if "activation" in kwargs:
        return kwargs["activation"]
    return _MISSING

get_act_name(name, layer=None, layer_type=None)

Convert common TransformerLens activation shorthands into hook names.

Source code in src/SafeLens/core/hooks.py

def get_act_name(
    name: str,
    layer: int | str | None = None,
    layer_type: str | None = None,
) -> str:
    """Convert common TransformerLens activation shorthands into hook names."""
    if ("." in name or name.startswith("hook_")) and layer is None and layer_type is None:
        return name

    name = _strip_hook_prefix(name)
    match = re.match(r"([a-z_]+)(\d+)([a-z]?.*)", name)
    if match is not None:
        name, parsed_layer, parsed_layer_type = match.groups()
        layer = parsed_layer
        layer_type = parsed_layer_type or layer_type

    name = _ACT_NAME_ALIASES.get(name, name)
    if name in _ATTN_ACTS:
        layer_type = "attn"
    elif name in _MLP_ACTS:
        layer_type = "mlp"
    elif layer_type in _LAYER_TYPE_ALIASES:
        layer_type = _LAYER_TYPE_ALIASES[layer_type]

    full_name = ""
    if layer is not None:
        full_name += f"blocks.{layer}."
    if layer_type:
        full_name += f"{layer_type}."
    full_name += f"hook_{name}"
    if name in _LAYER_NORM_ACTS and layer is None:
        full_name = f"ln_final.{full_name}"
    return full_name

has_hook_output(args, kwargs)

Return whether a hook call contains an activation output.

Source code in src/SafeLens/core/hooks.py

def has_hook_output(args: tuple[Any, ...], kwargs: dict[str, Any]) -> bool:
    """Return whether a hook call contains an activation output."""
    return extract_hook_output(args, kwargs) is not _MISSING

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

Create a hook that stores its activation in an ActivationCache.

Source code in src/SafeLens/core/hooks.py

def make_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 hook that stores its activation in an `ActivationCache`."""
    normalized_pos_slice = _normalize_cache_pos_slice(pos_slice)

    def cache_hook(*args: Any, **kwargs: Any) -> None:
        activation = extract_hook_output(args, kwargs)
        if activation is _MISSING:
            return None
        if remove_batch_dim:
            activation = _remove_singleton_batch(activation)
            cache.has_batch_dim = False
        if normalized_pos_slice is not None:
            activation = _slice_dim(
                activation,
                normalized_pos_slice,
                dim=_cache_pos_dim(name, activation),
            )
        cache.store(name, activation, detach=detach, clone=clone, device=device)
        return None

    cache_hook._safelens_is_cache_hook = True  # type: ignore[attr-defined]
    return cache_hook

matches_names_filter(name, names_filter=None)

Return whether an activation name matches a TransformerLens-style names filter.

Source code in src/SafeLens/core/hooks.py

def matches_names_filter(name: str, names_filter: NamesFilter = None) -> bool:
    """Return whether an activation name matches a TransformerLens-style names filter."""
    if names_filter is None:
        return True
    if isinstance(names_filter, str):
        return _activation_names_equivalent(name, names_filter)
    if callable(names_filter):
        return bool(names_filter(name))
    return any(_activation_names_equivalent(name, candidate) for candidate in names_filter)

prepare_activation_for_cache(activation, *, detach=True, clone=False, device=None)

Prepare an activation for caching without requiring a torch dependency.

Source code in src/SafeLens/core/hooks.py

def prepare_activation_for_cache(
    activation: Any,
    *,
    detach: bool = True,
    clone: bool = False,
    device: Any = None,
) -> Any:
    """Prepare an activation for caching without requiring a torch dependency."""
    value = activation
    detach_fn = getattr(value, "detach", None)
    if detach and callable(detach_fn):
        value = detach_fn()
    if clone:
        value = clone_activation(value)
    to_fn = getattr(value, "to", None)
    if device is not None and callable(to_fn):
        value = to_fn(device)
    return value

run_with_hooks(model, batch, hooks, *, layers=None)

Run a model with temporary hooks and remove them when the run finishes.

Source code in src/SafeLens/core/hooks.py

def run_with_hooks(
    model: ModelWrapper,
    batch: Batch,
    hooks: Iterable[tuple[LayerRef, HookFn]],
    *,
    layers: Sequence[LayerRef] | None = None,
) -> tuple[Any, dict[str, Any]]:
    """Run a model with temporary hooks and remove them when the run finishes."""
    with temporary_hooks(model, hooks):
        return model.run_with_cache(batch, layers=layers)

safelens_act_name(name, layer=None)

Return SafeLens-style activation names such as layer_0.resid_pre.

Source code in src/SafeLens/core/hooks.py

def safelens_act_name(name: str, layer: int | str | None = None) -> str:
    """Return SafeLens-style activation names such as `layer_0.resid_pre`."""
    name = _strip_hook_prefix(name)
    name = _ACT_NAME_ALIASES.get(name, name)
    if layer is None:
        return f"hook_{name}"
    return f"{activation_name_for_layer(layer)}.{name}"

stack_values(values)

Stack tensor-like values or fall back to a list copy.

Source code in src/SafeLens/core/hooks.py

def stack_values(values: Sequence[Any]) -> Any:
    """Stack tensor-like values or fall back to a list copy."""
    if not values:
        return []
    first = values[0]
    module = type(first).__module__.split(".")[0]
    if module == "torch":
        try:
            import torch

            return torch.stack(list(values))
        except Exception:
            pass
    if module == "numpy":
        try:
            import numpy as np

            return np.stack(list(values))
        except Exception:
            pass
    return [clone_activation(value) for value in values]

temporary_hooks(model, hooks)

Register hooks for one context and always remove them afterward.

Source code in src/SafeLens/core/hooks.py

@contextmanager
def temporary_hooks(
    model: ModelWrapper,
    hooks: Iterable[tuple[LayerRef, HookFn]],
) -> Iterator[list[Any]]:
    """Register hooks for one context and always remove them afterward."""
    handles: list[Any] = []
    try:
        for layer, hook_fn in hooks:
            handles.append(model.add_hook(layer, hook_fn))
        yield handles
    finally:
        for handle in reversed(handles):
            remove = getattr(handle, "remove", None)
            if callable(remove):
                remove()