Analysis Utilities

This module collects small, framework-light helpers used by logit lens, attribution, and ablation experiments.

Core utilities:

• softmax, log_softmax, and logits_to_log_probs.
• per_token_cross_entropy_loss, cross_entropy_loss, lm_log_probs, lm_cross_entropy_loss, and lm_accuracy.
• topk_tokens and logit_diff.
• residual_stack_to_logits, compute_head_results_from_z, and direct_logit_attribution.
• attention_pattern_score, previous_token_attention_score, and induction_attention_score for causal attention-pattern diagonal workflows, including repeated-token induction stripes via repeat_length.
• zero_ablation_hook, mean_ablation_hook, and replace_activation_hook.

Example:

from SafeLens.core.analysis import cross_entropy_loss, logit_diff

loss = cross_entropy_loss([[0.0, 0.0]], [1])
score = logit_diff([[[1.0, 4.0], [7.0, 2.0]]], 0, 1)

Ablation hooks can be attached to HookPoint or HookedRoot objects:

from SafeLens.core.analysis import zero_ablation_hook
from SafeLens.core.hooked_root import HookedRoot

root = HookedRoot()
hook = root.add_hook_point("blocks.0.hook_resid_pre")

with root.hooks(fwd_hooks=[("blocks.0.hook_resid_pre", zero_ablation_hook)]):
    assert hook([1, 2, 3]) == [0, 0, 0]

Small analysis helpers for logits, losses, ablations, and head detection.

and_values(left, right)

Elementwise boolean and for nested-list masks.

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

def and_values(left: Any, right: Any) -> Any:
    """Elementwise boolean and for nested-list masks."""
    try:
        if hasattr(left, "shape") or hasattr(right, "shape"):
            return left.bool() & right.bool()
    except Exception:
        pass
    if _is_sequence(left) and _is_sequence(right):
        return [
            and_values(left_item, right_item)
            for left_item, right_item in zip(left, right, strict=False)
        ]
    return bool(left) and bool(right)

argmax_last_dim(values)

Return argmax indices over the final dimension.

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

def argmax_last_dim(values: Any) -> Any:
    """Return argmax indices over the final dimension."""
    try:
        if hasattr(values, "shape"):
            return values.argmax(dim=-1)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(values, "shape"):
            return np.asarray(values).argmax(axis=-1)
    except Exception:
        pass
    if _is_sequence(values):
        if not values:
            return []
        if values and _is_sequence(values[0]):
            return [argmax_last_dim(item) for item in values]
        return max(range(len(values)), key=lambda index: float(values[index]))
    return values

attention_pattern_score(pattern, offset=-1, *, min_dest_pos=None)

Average attention paid to a fixed source-position offset.

pattern is expected to end in [dest_pos, src_pos], with any number of leading batch/layer/head dimensions preserved. offset=-1 scores previous token attention. Causal induction heads also attend backwards in the attention matrix, so induction-style matching is the same negative diagonal shifted by the repeat length in repeated-token prompts. min_dest_pos excludes diagonal entries before a destination position, which is useful for skipping the first copy of a repeated prompt.

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

def attention_pattern_score(
    pattern: Any,
    offset: int = -1,
    *,
    min_dest_pos: int | None = None,
) -> Any:
    """Average attention paid to a fixed source-position offset.

    `pattern` is expected to end in `[dest_pos, src_pos]`, with any number of
    leading batch/layer/head dimensions preserved. `offset=-1` scores previous
    token attention. Causal induction heads also attend backwards in the
    attention matrix, so induction-style matching is the same negative diagonal
    shifted by the repeat length in repeated-token prompts. `min_dest_pos`
    excludes diagonal entries before a destination position, which is useful for
    skipping the first copy of a repeated prompt.
    """
    try:
        import torch

        if hasattr(pattern, "shape") and isinstance(pattern, torch.Tensor):
            diag = torch.diagonal(pattern, offset=offset, dim1=-2, dim2=-1)
            diag = _slice_diagonal_by_dest_pos(diag, offset, min_dest_pos)
            if diag.shape[-1] == 0:
                return torch.zeros(pattern.shape[:-2], dtype=pattern.dtype, device=pattern.device)
            return diag.mean(dim=-1)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(pattern, "shape"):
            array = np.asarray(pattern)
            numpy_diag = np.diagonal(array, offset=offset, axis1=-2, axis2=-1)
            numpy_diag = _slice_diagonal_by_dest_pos(numpy_diag, offset, min_dest_pos)
            if numpy_diag.shape[-1] == 0:
                return np.zeros(array.shape[:-2], dtype=array.dtype)
            return numpy_diag.mean(axis=-1)
    except Exception:
        pass
    return _attention_pattern_score_nested(pattern, offset, min_dest_pos)

causal_lm_loss_mask(attention_mask)

Return mask for valid causal LM targets from an input attention mask.

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

def causal_lm_loss_mask(attention_mask: Any) -> Any:
    """Return mask for valid causal LM targets from an input attention mask."""
    try:
        import torch

        if hasattr(attention_mask, "shape"):
            if isinstance(attention_mask, torch.Tensor):
                return attention_mask[..., :-1].bool() & attention_mask[..., 1:].bool()
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(attention_mask, "shape"):
            mask = np.asarray(attention_mask).astype(bool)
            return mask[..., :-1] & mask[..., 1:]
    except Exception:
        pass
    previous_mask = slice_last_dim(attention_mask, stop=-1)
    next_mask = slice_last_dim(attention_mask, start=1)
    return and_values(previous_mask, next_mask)

compute_head_attention_similarity_score(attention_pattern, detection_pattern, *, exclude_bos, exclude_current_token, error_measure)

Compute similarity between a single head pattern and a detector pattern.

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

def compute_head_attention_similarity_score(
    attention_pattern: Any,
    detection_pattern: Any,
    *,
    exclude_bos: bool,
    exclude_current_token: bool,
    error_measure: str,
) -> float:
    """Compute similarity between a single head pattern and a detector pattern."""
    if error_measure not in ERROR_MEASURES:
        raise ValueError(
            f"Invalid error_measure={error_measure!r}; valid values are {ERROR_MEASURES}"
        )
    _validate_square_attention_pattern(attention_pattern)
    try:
        import torch

        if isinstance(attention_pattern, torch.Tensor) or isinstance(
            detection_pattern, torch.Tensor
        ):
            attention = _as_torch_float_tensor(attention_pattern, like=detection_pattern).clone()
            detection = _as_torch_float_tensor(detection_pattern, like=attention)
            if error_measure == "mul":
                if exclude_bos:
                    attention[:, 0] = 0
                if exclude_current_token:
                    attention.fill_diagonal_(0)
                return float(((attention * detection).sum() / attention.sum()).item())
            abs_diff = (attention - detection).abs()
            if not torch.allclose(abs_diff, torch.tril(abs_diff)):
                raise AssertionError(
                    "Attention pattern and detection pattern differ above the diagonal."
                )
            if exclude_bos:
                abs_diff[:, 0] = 0
            if exclude_current_token:
                abs_diff.fill_diagonal_(0)
            return 1 - round(float(abs_diff.mean().item() * len(abs_diff)), 3)
    except ImportError:
        pass
    try:
        import numpy as np

        if hasattr(attention_pattern, "shape") or hasattr(detection_pattern, "shape"):
            attention = np.asarray(attention_pattern, dtype=float).copy()
            detection = np.asarray(detection_pattern, dtype=float)
            if error_measure == "mul":
                if exclude_bos:
                    attention[:, 0] = 0
                if exclude_current_token:
                    np.fill_diagonal(attention, 0)
                denominator = attention.sum()
                return float(
                    np.nan if denominator == 0 else (attention * detection).sum() / denominator
                )
            numpy_abs_diff = np.abs(attention - detection)
            if not np.allclose(numpy_abs_diff, np.tril(numpy_abs_diff)):
                raise AssertionError(
                    "Attention pattern and detection pattern differ above the diagonal."
                )
            if exclude_bos:
                numpy_abs_diff[:, 0] = 0
            if exclude_current_token:
                np.fill_diagonal(numpy_abs_diff, 0)
            return 1 - round(float(numpy_abs_diff.mean() * len(numpy_abs_diff)), 3)
    except ImportError:
        pass

    attention = _nested_float_matrix(attention_pattern)
    detection = _nested_float_matrix(detection_pattern)
    if error_measure == "mul":
        numerator = 0.0
        denominator = 0.0
        for dest, row in enumerate(attention):
            for src, value in enumerate(row):
                if exclude_bos and src == 0:
                    value = 0.0
                if exclude_current_token and src == dest:
                    value = 0.0
                numerator += value * detection[dest][src]
                denominator += value
        return numerator / denominator if denominator else float("nan")

    nested_abs_diff: list[list[float]] = []
    for dest, row in enumerate(attention):
        diff_row = []
        for src, value in enumerate(row):
            diff = abs(value - detection[dest][src])
            if src > dest and diff:
                raise AssertionError(
                    "Attention pattern and detection pattern differ above the diagonal."
                )
            if exclude_bos and src == 0:
                diff = 0.0
            if exclude_current_token and src == dest:
                diff = 0.0
            diff_row.append(diff)
        nested_abs_diff.append(diff_row)
    total = sum(sum(row) for row in nested_abs_diff)
    size = len(nested_abs_diff)
    return 1 - round(total / max(1, size * size) * size, 3)

compute_head_results_from_z(z, W_O, *, has_layer_axis=None)

Project per-head z activations through W_O into residual-space results.

z is expected to end in [head, d_head], and W_O should be shaped [head, d_head, d_model] or [layer, head, d_head, d_model]. Leading batch/position dimensions are preserved. When z is stacked by layer but omits an explicit batch dimension, pass has_layer_axis=True to disambiguate the leading axis from a batch/position axis.

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

def compute_head_results_from_z(
    z: Any,
    W_O: Any,
    *,
    has_layer_axis: bool | None = None,
) -> Any:
    """Project per-head `z` activations through `W_O` into residual-space results.

    `z` is expected to end in `[head, d_head]`, and `W_O` should be shaped
    `[head, d_head, d_model]` or `[layer, head, d_head, d_model]`.
    Leading batch/position dimensions are preserved. When `z` is stacked by
    layer but omits an explicit batch dimension, pass `has_layer_axis=True` to
    disambiguate the leading axis from a batch/position axis.
    """
    try:
        import torch

        if hasattr(z, "shape") or hasattr(W_O, "shape"):
            if not hasattr(z, "shape"):
                z = torch.as_tensor(
                    z,
                    dtype=getattr(W_O, "dtype", None),
                    device=getattr(W_O, "device", None),
                )
            if not hasattr(W_O, "shape"):
                W_O = torch.as_tensor(
                    W_O,
                    dtype=getattr(z, "dtype", None),
                    device=getattr(z, "device", None),
                )
            if _has_aligned_layer_axis(z, W_O, has_layer_axis=has_layer_axis):
                return torch.einsum("l...hd,lhdm->l...hm", z, W_O)
            W_O = _normalize_w_o_for_z(z, W_O, has_layer_axis=has_layer_axis)
            return torch.einsum("...hd,hdm->...hm", z, W_O)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(z, "shape") and hasattr(W_O, "shape"):
            if _has_aligned_layer_axis(z, W_O, has_layer_axis=has_layer_axis):
                return np.einsum("l...hd,lhdm->l...hm", z, W_O)
            W_O = _normalize_w_o_for_z(z, W_O, has_layer_axis=has_layer_axis)
            return np.einsum("...hd,hdm->...hm", z, W_O)
    except Exception:
        pass

    W_O = _normalize_w_o_for_z(z, W_O, has_layer_axis=has_layer_axis)
    return _head_results_from_nested(z, W_O, has_layer_axis=has_layer_axis)

cross_entropy_loss(logits, tokens)

Return mean cross-entropy loss.

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

def cross_entropy_loss(logits: Any, tokens: Any) -> float:
    """Return mean cross-entropy loss."""
    losses = flatten(per_token_cross_entropy_loss(logits, tokens))
    return sum(float(loss) for loss in losses) / max(1, len(losses))

detect_head(model, seq, detection_pattern, heads=None, cache=None, *, exclude_bos=False, exclude_current_token=False, error_measure='mul')

Search cached attention patterns for TransformerLens-style head patterns.

This mirrors transformer_lens.head_detector.detect_head without depending on TransformerLens. detection_pattern can be one of "previous_token_head", "duplicate_token_head", "induction_head", or an explicit square lower-triangular pattern. Returned scores are shaped [n_layers, n_heads] and unselected heads are set to -1.

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

def detect_head(
    model: Any,
    seq: str | Sequence[str] | Any,
    detection_pattern: Any,
    heads: Sequence[tuple[int, int]] | Mapping[int, Sequence[int]] | None = None,
    cache: ActivationCache | Mapping[Any, Any] | None = None,
    *,
    exclude_bos: bool = False,
    exclude_current_token: bool = False,
    error_measure: str = "mul",
) -> Any:
    """Search cached attention patterns for TransformerLens-style head patterns.

    This mirrors ``transformer_lens.head_detector.detect_head`` without depending
    on TransformerLens. ``detection_pattern`` can be one of
    ``"previous_token_head"``, ``"duplicate_token_head"``, ``"induction_head"``,
    or an explicit square lower-triangular pattern. Returned scores are shaped
    ``[n_layers, n_heads]`` and unselected heads are set to ``-1``.
    """
    if error_measure not in ERROR_MEASURES:
        raise ValueError(
            f"Invalid error_measure={error_measure!r}; valid values are {ERROR_MEASURES}"
        )
    if isinstance(detection_pattern, str) and _is_string_sequence(seq) and cache is None:
        scores = [
            detect_head(
                model,
                item,
                detection_pattern,
                heads=heads,
                cache=None,
                exclude_bos=exclude_bos,
                exclude_current_token=exclude_current_token,
                error_measure=error_measure,
            )
            for item in seq
        ]
        return _mean_score_matrices(scores)

    tokens = _tokenize_head_detector_sequence(model, seq)
    seq_len = _sequence_length(tokens)
    cfg = _model_cfg(model)
    n_ctx = _cfg_int(cfg, "n_ctx")
    if seq_len <= 1 or (n_ctx is not None and seq_len >= n_ctx):
        raise ValueError(SEQ_LEN_ERR)

    if isinstance(detection_pattern, str):
        if detection_pattern not in HEAD_NAMES:
            raise ValueError(INVALID_HEAD_NAME_ERR % detection_pattern)
        detection_pattern = _named_head_detection_pattern(detection_pattern, tokens)

    detection_pattern = _move_pattern_to_token_device(detection_pattern, tokens, cfg)
    _validate_detection_pattern(detection_pattern, seq_len)
    if error_measure == "mul" and not _pattern_values_are_binary(detection_pattern):
        logging.warning(
            "Using detection pattern with values other than 0 or 1 with error_measure 'mul'"
        )

    if cache is None:
        run_with_cache = getattr(model, "run_with_cache", None)
        if not callable(run_with_cache):
            raise TypeError("detect_head requires `cache` or a model with `run_with_cache`.")
        cache_result = run_with_cache(tokens, remove_batch_dim=True)
        if not isinstance(cache_result, tuple | list) or len(cache_result) < 2:
            raise TypeError("model.run_with_cache must return an (output, cache) pair.")
        cache = cache_result[1]
        if cache is None:
            raise TypeError("model.run_with_cache returned None for cache.")
    resolved_cache = _ensure_activation_cache(cache, model)
    n_layers, n_heads = _infer_head_score_shape(model, resolved_cache, heads)
    layer_to_heads = _normalize_head_selection(heads, n_layers=n_layers, n_heads=n_heads)
    matches = _make_head_score_matrix(n_layers, n_heads, cfg, tokens, resolved_cache)

    for layer, layer_heads in layer_to_heads.items():
        layer_patterns = _normalize_cached_layer_attention_pattern(
            _get_cached_attention_pattern(resolved_cache, layer),
            n_heads=n_heads,
            seq_len=seq_len,
        )
        for head in layer_heads:
            head_pattern = _index_head_attention_pattern(layer_patterns, head)
            score = compute_head_attention_similarity_score(
                head_pattern,
                detection_pattern=detection_pattern,
                exclude_bos=exclude_bos,
                exclude_current_token=exclude_current_token,
                error_measure=error_measure,
            )
            _set_head_score(matches, layer, head, score)
    return matches

direct_logit_attribution(residual_stack, token_directions)

Project residual components onto token directions.

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

def direct_logit_attribution(residual_stack: Any, token_directions: Any) -> Any:
    """Project residual components onto token directions."""
    return dot_last_dim(residual_stack, token_directions)

dot_last_dim(left, right)

Dot product over final dimension.

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

def dot_last_dim(left: Any, right: Any) -> Any:
    """Dot product over final dimension."""
    try:
        import torch

        if isinstance(left, torch.Tensor) or isinstance(right, torch.Tensor):
            if not isinstance(left, torch.Tensor):
                left = torch.as_tensor(
                    left,
                    dtype=getattr(right, "dtype", None),
                    device=getattr(right, "device", None),
                )
            if not isinstance(right, torch.Tensor):
                right = torch.as_tensor(
                    right,
                    dtype=getattr(left, "dtype", None),
                    device=getattr(left, "device", None),
                )
            return (left * right).sum(dim=-1)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(left, "shape") or hasattr(right, "shape"):
            return (np.asarray(left) * np.asarray(right)).sum(axis=-1)
    except Exception:
        pass
    try:
        return (left * right).sum(dim=-1)
    except Exception:
        pass
    return _dot_nested(left, right)

equal_values(left, right)

Elementwise equality returning numeric 0/1 values for list backends.

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

def equal_values(left: Any, right: Any) -> Any:
    """Elementwise equality returning numeric 0/1 values for list backends."""
    try:
        import torch

        if hasattr(left, "shape") and isinstance(left, torch.Tensor):
            if not hasattr(right, "shape"):
                right = torch.as_tensor(right, device=left.device)
            else:
                right = right.to(device=left.device)
            return left == right
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(left, "shape") or hasattr(right, "shape"):
            return np.asarray(left) == np.asarray(right)
    except Exception:
        pass
    if _is_sequence(left) and _is_sequence(right):
        return [
            equal_values(left_item, right_item)
            for left_item, right_item in zip(left, right, strict=False)
        ]
    return 1.0 if left == right else 0.0

flatten(value)

Flatten nested lists or tensor-like values.

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

def flatten(value: Any) -> list[Any]:
    """Flatten nested lists or tensor-like values."""
    if hasattr(value, "shape"):
        try:
            return list(value.reshape(-1))
        except Exception:
            pass
    tolist = getattr(value, "tolist", None)
    if callable(tolist):
        value = tolist()
    if _is_sequence(value):
        result: list[Any] = []
        for item in value:
            result.extend(flatten(item))
        return result
    return [value]

gather_last_dim(values, indices)

Gather values at final-dimension indices.

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

def gather_last_dim(values: Any, indices: Any) -> Any:
    """Gather values at final-dimension indices."""
    try:
        import torch

        if hasattr(values, "shape"):
            if not hasattr(indices, "shape"):
                indices = torch.as_tensor(indices, dtype=torch.long, device=values.device)
            else:
                indices = indices.to(device=values.device, dtype=torch.long)
            if indices.ndim == 0:
                return values[..., int(indices.item())]
            return values.gather(-1, indices.unsqueeze(-1)).squeeze(-1)
    except Exception:
        pass
    try:
        if hasattr(values, "shape") and hasattr(indices, "shape"):
            if len(getattr(indices, "shape", ())) == 0:
                item = getattr(indices, "item", None)
                index = int(cast(Any, item() if callable(item) else indices))
                return values[..., index]
            return values.gather(-1, indices.unsqueeze(-1)).squeeze(-1)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(values, "shape") or hasattr(indices, "shape"):
            values_array = np.asarray(values)
            index_array = np.asarray(indices)
            if index_array.ndim == 0:
                return values_array[..., int(index_array.item())]
            return np.take_along_axis(
                values_array, np.expand_dims(index_array, -1), axis=-1
            ).squeeze(-1)
    except Exception:
        pass
    if _is_sequence(indices):
        if indices and _is_sequence(indices[0]):
            return [
                gather_last_dim(value_row, index_row)
                for value_row, index_row in zip(values, indices, strict=True)
            ]
        return [row[index] for row, index in zip(values, indices, strict=True)]
    if _is_sequence(values) and values and _is_sequence(values[0]):
        return [gather_last_dim(value_row, indices) for value_row in values]
    return values[indices]

get_duplicate_token_head_detection_pattern(tokens)

Return a pattern whose entries mark earlier equal tokens.

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

def get_duplicate_token_head_detection_pattern(tokens: Any) -> Any:
    """Return a pattern whose entries mark earlier equal tokens."""
    values = _token_sequence_values(tokens)
    seq_len = len(values)
    try:
        import torch

        if isinstance(tokens, torch.Tensor):
            token_tensor = torch.as_tensor(values, device=tokens.device)
            pattern = token_tensor[:, None].eq(token_tensor[None, :]).to(torch.float32)
            pattern.fill_diagonal_(0)
            return torch.tril(pattern)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(tokens, "shape"):
            token_array = np.asarray(values)
            pattern = (token_array[:, None] == token_array[None, :]).astype(float)
            np.fill_diagonal(pattern, 0)
            return np.tril(pattern)
    except Exception:
        pass
    return [
        [1.0 if dest > src and values[dest] == values[src] else 0.0 for src in range(seq_len)]
        for dest in range(seq_len)
    ]

get_induction_head_detection_pattern(tokens)

Return a duplicate-token pattern shifted right for induction heads.

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

def get_induction_head_detection_pattern(tokens: Any) -> Any:
    """Return a duplicate-token pattern shifted right for induction heads."""
    duplicate_pattern = get_duplicate_token_head_detection_pattern(tokens)
    try:
        import torch

        if isinstance(duplicate_pattern, torch.Tensor):
            shifted = torch.roll(duplicate_pattern, shifts=1, dims=1)
            shifted[:, 0] = 0
            return torch.tril(shifted)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(duplicate_pattern, "shape"):
            numpy_shifted = np.roll(np.asarray(duplicate_pattern), shift=1, axis=1)
            numpy_shifted[:, 0] = 0
            return np.tril(numpy_shifted)
    except Exception:
        pass
    seq_len = len(duplicate_pattern)
    return [
        [
            float(src > 0 and dest >= src and duplicate_pattern[dest][src - 1])
            for src in range(seq_len)
        ]
        for dest in range(seq_len)
    ]

get_previous_token_head_detection_pattern(tokens)

Return a lower-triangular pattern for attention to the previous token.

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

def get_previous_token_head_detection_pattern(tokens: Any) -> Any:
    """Return a lower-triangular pattern for attention to the previous token."""
    seq_len = _sequence_length(tokens)
    try:
        import torch

        if isinstance(tokens, torch.Tensor):
            pattern = torch.zeros((seq_len, seq_len), dtype=torch.float32, device=tokens.device)
            if seq_len > 1:
                pattern[1:, :-1] = torch.eye(
                    seq_len - 1, dtype=pattern.dtype, device=pattern.device
                )
            return torch.tril(pattern)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(tokens, "shape"):
            numpy_pattern = np.zeros((seq_len, seq_len), dtype=float)
            if seq_len > 1:
                numpy_pattern[1:, :-1] = np.eye(seq_len - 1)
            return np.tril(numpy_pattern)
    except Exception:
        pass
    return [[1.0 if dest == src + 1 else 0.0 for src in range(seq_len)] for dest in range(seq_len)]

get_supported_heads()

Print and return supported TransformerLens-style head detector names.

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

def get_supported_heads() -> list[str]:
    """Print and return supported TransformerLens-style head detector names."""
    heads = [str(name) for name in HEAD_NAMES]
    print(f"Supported heads: {heads}")
    return heads

induction_attention_score(pattern, *, offset=-1, repeat_length=None)

Score induction-head style attention on a causal backward diagonal.

For the minimal [A][B][A] -> [B] setup this is the previous-token diagonal (offset=-1). For a repeated sequence of length N, induction attention from the second copy to the next token after the first copy is on offset=1-N.

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

def induction_attention_score(
    pattern: Any,
    *,
    offset: int = -1,
    repeat_length: int | None = None,
) -> Any:
    """Score induction-head style attention on a causal backward diagonal.

    For the minimal `[A][B][A] -> [B]` setup this is the previous-token
    diagonal (`offset=-1`). For a repeated sequence of length `N`, induction
    attention from the second copy to the next token after the first copy is on
    `offset=1-N`.
    """
    if repeat_length is not None:
        if repeat_length < 2:
            raise ValueError("repeat_length must be at least 2 for induction attention.")
        offset = 1 - repeat_length
        return attention_pattern_score(pattern, offset=offset, min_dest_pos=repeat_length)
    return attention_pattern_score(pattern, offset=offset)

is_valid_number(value)

Return whether a flattened scalar should contribute to an aggregate.

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

def is_valid_number(value: Any) -> bool:
    """Return whether a flattened scalar should contribute to an aggregate."""
    if value is None:
        return False
    try:
        return not math.isnan(float(value))
    except (TypeError, ValueError):
        return True

lm_accuracy(logits, tokens, attention_mask=None, *, per_token=False)

Return next-token prediction accuracy for causal language modeling.

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

def lm_accuracy(
    logits: Any,
    tokens: Any,
    attention_mask: Any | None = None,
    *,
    per_token: bool = False,
) -> Any:
    """Return next-token prediction accuracy for causal language modeling."""
    predictions = argmax_last_dim(slice_second_last_dim(logits, stop=-1))
    targets = slice_last_dim(tokens, start=1)
    correct = equal_values(predictions, targets)
    if attention_mask is None:
        if per_token:
            return correct
        values = flatten(correct)
        if not values:
            return float("nan")
        return sum(float(value) for value in values) / len(values)
    masked = mask_values(correct, causal_lm_loss_mask(attention_mask))
    if per_token:
        return masked
    values = [float(value) for value in flatten(masked) if is_valid_number(value)]
    if not values:
        return float("nan")
    return sum(values) / len(values)

lm_cross_entropy_loss(logits, tokens, attention_mask=None, *, per_token=False)

Return causal LM cross-entropy using logits before each target token.

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

def lm_cross_entropy_loss(
    logits: Any,
    tokens: Any,
    attention_mask: Any | None = None,
    *,
    per_token: bool = False,
) -> Any:
    """Return causal LM cross-entropy using logits before each target token."""
    losses = negate_values(lm_log_probs(logits, tokens))
    if attention_mask is not None:
        mask = causal_lm_loss_mask(attention_mask)
        if per_token:
            return zero_mask_values(losses, mask)
        losses = mask_values(losses, mask)
    if per_token:
        return losses
    values = [float(value) for value in flatten(losses) if is_valid_number(value)]
    if not values:
        return float("nan")
    return sum(values) / len(values)

lm_log_probs(logits, tokens, attention_mask=None)

Return next-token log-probabilities for causal language modeling.

Logits at position i are gathered at token i + 1, matching TransformerLens' language-model loss convention.

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

def lm_log_probs(logits: Any, tokens: Any, attention_mask: Any | None = None) -> Any:
    """Return next-token log-probabilities for causal language modeling.

    Logits at position `i` are gathered at token `i + 1`, matching
    TransformerLens' language-model loss convention.
    """
    shifted_logits = slice_second_last_dim(logits, stop=-1)
    shifted_tokens = slice_last_dim(tokens, start=1)
    log_probs = logits_to_log_probs(shifted_logits, shifted_tokens)
    if attention_mask is None:
        return log_probs
    return mask_values(log_probs, causal_lm_loss_mask(attention_mask))

log_softmax(values)

Apply log-softmax over the final dimension.

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

def log_softmax(values: Any) -> Any:
    """Apply log-softmax over the final dimension."""
    try:
        import torch

        if hasattr(values, "shape"):
            return torch.log_softmax(values, dim=-1)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(values, "shape"):
            probs = softmax(values)
            return np.log(probs)
    except Exception:
        pass
    probs = softmax(values)
    if _is_sequence(probs) and probs and _is_sequence(probs[0]):
        return [log_softmax(row) for row in values]
    return [math.log(float(prob)) for prob in probs]

logit_diff(logits, correct_token, incorrect_token, *, pos=-1)

Return logit difference at one position.

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

def logit_diff(logits: Any, correct_token: int, incorrect_token: int, *, pos: int = -1) -> float:
    """Return logit difference at one position."""
    shape = getattr(logits, "shape", None)
    if shape is not None:
        if len(shape) >= 3:
            value = logits[0, pos, correct_token] - logits[0, pos, incorrect_token]
        elif len(shape) == 2:
            value = logits[pos, correct_token] - logits[pos, incorrect_token]
        else:
            value = logits[correct_token] - logits[incorrect_token]
        item = getattr(value, "item", None)
        return float(cast(Any, item() if callable(item) else value))

    if _is_sequence(logits) and logits and _is_sequence(logits[0]):
        if logits[0] and _is_sequence(logits[0][0]):
            row = logits[0][pos]
        else:
            row = logits[pos]
    else:
        row = logits
    return float(row[correct_token]) - float(row[incorrect_token])

logits_to_df(logits, tokenizer=None, top_k=None)

Convert a 1-D logit vector into a probability-sorted pandas DataFrame.

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

def logits_to_df(logits: Any, tokenizer: Any | None = None, top_k: int | None = None) -> Any:
    """Convert a 1-D logit vector into a probability-sorted pandas DataFrame."""

    import pandas as pd

    values = [float(value) for value in _as_flat_list(logits)]
    log_probs = _log_softmax_vector(values)
    probabilities = [math.exp(value) for value in log_probs]
    order = sorted(range(len(values)), key=lambda index: probabilities[index], reverse=True)
    if top_k is not None:
        order = order[:top_k]

    data: dict[str, Any] = {"token_index": order}
    if tokenizer is not None:
        data["token_string"] = [_decode_token_for_dataframe(tokenizer, index) for index in order]
    data["logit"] = [values[index] for index in order]
    data["log_prob"] = [log_probs[index] for index in order]
    data["probability"] = [probabilities[index] for index in order]
    return pd.DataFrame(data)

logits_to_log_probs(logits, tokens=None)

Convert logits to log probabilities, optionally gathering token log-probs.

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

def logits_to_log_probs(logits: Any, tokens: Any | None = None) -> Any:
    """Convert logits to log probabilities, optionally gathering token log-probs."""
    log_probs = log_softmax(logits)
    if tokens is None:
        return log_probs
    return gather_last_dim(log_probs, tokens)

map_values(value, fn)

Map over nested list leaves.

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

def map_values(value: Any, fn: Callable[[Any], Any]) -> Any:
    """Map over nested list leaves."""
    try:
        import torch

        if hasattr(value, "shape"):
            try:
                mapped = fn(value)
            except Exception:
                return torch.zeros_like(value) if fn(1) == 0 else value
            if mapped is None:
                return None
            if hasattr(mapped, "shape"):
                return mapped
            if mapped == 0:
                return torch.zeros_like(value)
            return mapped
    except Exception:
        pass
    if _is_sequence(value):
        return [map_values(item, fn) for item in value]
    return fn(value)

mask_values(values, mask)

Set values to None wherever mask is false.

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

def mask_values(values: Any, mask: Any) -> Any:
    """Set values to `None` wherever mask is false."""
    try:
        import torch

        if hasattr(values, "shape"):
            if not hasattr(mask, "shape"):
                mask = torch.as_tensor(mask, dtype=torch.bool, device=values.device)
            else:
                mask = mask.to(device=values.device, dtype=torch.bool)
            if values.dtype.is_floating_point:
                fill_value = torch.full_like(values, float("nan"))
                return torch.where(mask, values, fill_value)
            return torch.where(mask, values.float(), torch.full_like(values.float(), float("nan")))
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(values, "shape") or hasattr(mask, "shape"):
            values_array = np.asarray(values)
            mask_array = np.asarray(mask).astype(bool)
            float_values = values_array.astype(float, copy=False)
            return np.where(mask_array, float_values, np.full_like(float_values, np.nan))
    except Exception:
        pass
    if _is_sequence(values) and _is_sequence(mask):
        return [
            mask_values(value_item, mask_item)
            for value_item, mask_item in zip(values, mask, strict=False)
        ]
    return values if bool(mask) else None

matmul_last_dim(left, right)

Multiply left[..., d] @ right[d, out] for nested-list values.

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

def matmul_last_dim(left: Any, right: Any) -> Any:
    """Multiply `left[..., d] @ right[d, out]` for nested-list values."""
    if _shape_of(left) == ():
        return left
    if _is_vector(left):
        return _matvec(left, right)
    return [matmul_last_dim(item, right) for item in left]

mean_ablation_hook(activation, hook=None)

Hook that replaces values with the activation mean.

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

def mean_ablation_hook(activation: Any, hook: HookPoint | None = None) -> Any:
    """Hook that replaces values with the activation mean."""
    _ = hook
    try:
        import torch

        if hasattr(activation, "shape"):
            return torch.zeros_like(activation) + activation.float().mean()
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(activation, "shape"):
            return np.zeros_like(activation) + np.asarray(activation, dtype=float).mean()
    except Exception:
        pass
    values = [float(value) for value in flatten(activation)]
    mean_value = sum(values) / max(1, len(values))
    return map_values(activation, lambda _value: mean_value)

negate_values(value)

Negate nested values.

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

def negate_values(value: Any) -> Any:
    """Negate nested values."""
    if value is None:
        return None
    try:
        return -value
    except Exception:
        pass
    return map_values(value, lambda item: None if item is None else -float(item))

per_token_cross_entropy_loss(logits, tokens)

Return negative log-probability for each target token.

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

def per_token_cross_entropy_loss(logits: Any, tokens: Any) -> Any:
    """Return negative log-probability for each target token."""
    gathered = logits_to_log_probs(logits, tokens)
    return negate_values(gathered)

previous_token_attention_score(pattern)

Score attention to the immediately previous token.

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

def previous_token_attention_score(pattern: Any) -> Any:
    """Score attention to the immediately previous token."""
    return attention_pattern_score(pattern, offset=-1)

replace_activation_hook(replacement)

Return a hook that replaces the full activation.

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

def replace_activation_hook(replacement: Any) -> Callable[[Any, HookPoint | None], Any]:
    """Return a hook that replaces the full activation."""

    def hook(_activation: Any, _hook: HookPoint | None = None) -> Any:
        return clone_activation(replacement)

    return hook

residual_stack_to_logits(residual_stack, unembed, unembed_bias=None)

Project residual components through an unembedding matrix and optional bias.

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

def residual_stack_to_logits(
    residual_stack: Any,
    unembed: Any,
    unembed_bias: Any | None = None,
) -> Any:
    """Project residual components through an unembedding matrix and optional bias."""
    try:
        import torch

        if hasattr(residual_stack, "shape") or hasattr(unembed, "shape"):
            if not hasattr(residual_stack, "shape"):
                residual_stack = torch.as_tensor(
                    residual_stack,
                    dtype=getattr(unembed, "dtype", None),
                    device=getattr(unembed, "device", None),
                )
            if not hasattr(unembed, "shape"):
                unembed = torch.as_tensor(
                    unembed,
                    dtype=getattr(residual_stack, "dtype", None),
                    device=getattr(residual_stack, "device", None),
                )
            logits = residual_stack @ unembed
            return _add_unembed_bias(logits, unembed_bias)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(residual_stack, "shape") or hasattr(unembed, "shape"):
            logits = np.matmul(residual_stack, unembed)
            return _add_unembed_bias(logits, unembed_bias)
    except Exception:
        pass
    logits = matmul_last_dim(residual_stack, unembed)
    return _add_unembed_bias(logits, unembed_bias)

sample_logits(final_logits, top_k=None, top_p=None, temperature=1.0, freq_penalty=0.0, repetition_penalty=1.0, tokens=None)

Sample token IDs from final logits with TransformerLens-style controls.

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

def sample_logits(
    final_logits: Any,
    top_k: int | None = None,
    top_p: float | None = None,
    temperature: float = 1.0,
    freq_penalty: float = 0.0,
    repetition_penalty: float = 1.0,
    tokens: Any | None = None,
) -> Any:
    """Sample token IDs from final logits with TransformerLens-style controls."""

    if top_k is not None:
        assert top_k > 0, "top_k has to be greater than 0"
    if top_p is not None:
        assert 1.0 >= top_p > 0.0, "top_p has to be in (0, 1]"
    assert temperature >= 0.0, "temperature has to be non-negative"
    assert freq_penalty >= 0.0, "freq_penalty has to be non-negative"
    assert repetition_penalty > 0.0, "repetition_penalty has to be greater than 0"

    torch_result = _sample_logits_torch(
        final_logits,
        top_k=top_k,
        top_p=top_p,
        temperature=temperature,
        freq_penalty=freq_penalty,
        repetition_penalty=repetition_penalty,
        tokens=tokens,
    )
    if torch_result is not None:
        return torch_result

    return _sample_logits_python(
        final_logits,
        top_k=top_k,
        top_p=top_p,
        temperature=temperature,
        freq_penalty=freq_penalty,
        repetition_penalty=repetition_penalty,
        tokens=tokens,
    )

slice_last_dim(value, *, start=None, stop=None)

Slice the last dimension of tensor-like or nested-list values.

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

def slice_last_dim(value: Any, *, start: int | None = None, stop: int | None = None) -> Any:
    """Slice the last dimension of tensor-like or nested-list values."""
    try:
        if hasattr(value, "shape"):
            return value[..., slice(start, stop)]
    except Exception:
        pass
    return _slice_nested_dim(value, slice(start, stop), dim=-1)

slice_second_last_dim(value, *, start=None, stop=None)

Slice the second-last dimension of tensor-like or nested-list values.

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

def slice_second_last_dim(value: Any, *, start: int | None = None, stop: int | None = None) -> Any:
    """Slice the second-last dimension of tensor-like or nested-list values."""
    try:
        if hasattr(value, "shape"):
            return value[..., slice(start, stop), :]
    except Exception:
        pass
    return _slice_nested_dim(value, slice(start, stop), dim=-2)

softmax(values)

Apply softmax over the final dimension.

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

def softmax(values: Any) -> Any:
    """Apply softmax over the final dimension."""
    try:
        import torch

        if hasattr(values, "shape"):
            return torch.softmax(values, dim=-1)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(values, "shape"):
            array = np.asarray(values)
            shifted = array - np.max(array, axis=-1, keepdims=True)
            exps = np.exp(shifted)
            return exps / np.sum(exps, axis=-1, keepdims=True)
    except Exception:
        pass
    if _is_sequence(values) and not values:
        return []
    if _is_sequence(values) and values and _is_sequence(values[0]):
        return [softmax(row) for row in values]
    max_value = max(float(value) for value in values)
    exps = [math.exp(float(value) - max_value) for value in values]
    total = sum(exps)
    return [value / total for value in exps]

test_prompt(*args, **kwargs)

Run a TransformerLens-style prompt sanity check.

Supports both SafeLens' structured call shape test_prompt(model, prompt, correct_token, incorrect_token=None, ...) and TransformerLens' exploratory call shape test_prompt(prompt, answer, model, ...). Both return a structured result; TL-style calls additionally include answer-token ranks.

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

def test_prompt(
    *args: Any,
    **kwargs: Any,
) -> dict[str, Any]:
    """Run a TransformerLens-style prompt sanity check.

    Supports both SafeLens' structured call shape
    ``test_prompt(model, prompt, correct_token, incorrect_token=None, ...)`` and
    TransformerLens' exploratory call shape
    ``test_prompt(prompt, answer, model, ...)``. Both return a structured result;
    TL-style calls additionally include answer-token ranks.
    """
    if len(args) >= 3 and isinstance(args[0], str) and not isinstance(args[2], str | int):
        return _test_prompt_transformerlens(*args, **kwargs)
    return _test_prompt_structured(*args, **kwargs)

topk_tokens(logits, k=5)

Return top-k token indices and values for the final dimension.

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

def topk_tokens(logits: Any, k: int = 5) -> Any:
    """Return top-k token indices and values for the final dimension."""
    k = _clamp_top_k(logits, k)
    try:
        import torch

        if hasattr(logits, "shape"):
            values, indices = torch.topk(logits, k, dim=-1)
            return indices, values
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(logits, "shape"):
            array = np.asarray(logits)
            sorted_indices = np.argsort(array, axis=-1)[..., ::-1][..., :k]
            sorted_values = np.take_along_axis(array, sorted_indices, axis=-1)
            return sorted_indices, sorted_values
    except Exception:
        pass
    if _is_sequence(logits) and logits and _is_sequence(logits[0]):
        return [topk_tokens(row, k=k) for row in logits]
    pairs = sorted(enumerate(logits), key=lambda item: float(item[1]), reverse=True)[:k]
    return [index for index, _value in pairs], [value for _index, value in pairs]

zero_ablation_hook(activation, hook=None)

Hook that replaces an activation with zeros.

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

def zero_ablation_hook(activation: Any, hook: HookPoint | None = None) -> Any:
    """Hook that replaces an activation with zeros."""
    _ = hook
    try:
        import torch

        if hasattr(activation, "shape"):
            return torch.zeros_like(activation)
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(activation, "shape"):
            return np.zeros_like(activation)
    except Exception:
        pass
    return map_values(activation, lambda _value: 0)

zero_mask_values(values, mask)

Set values to zero wherever mask is false.

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

def zero_mask_values(values: Any, mask: Any) -> Any:
    """Set values to zero wherever mask is false."""
    try:
        import torch

        if hasattr(values, "shape"):
            if not hasattr(mask, "shape"):
                mask = torch.as_tensor(mask, dtype=torch.bool, device=values.device)
            else:
                mask = mask.to(device=values.device, dtype=torch.bool)
            return torch.where(mask, values, torch.zeros_like(values))
    except Exception:
        pass
    try:
        import numpy as np

        if hasattr(values, "shape") or hasattr(mask, "shape"):
            return np.where(np.asarray(mask).astype(bool), np.asarray(values), 0)
    except Exception:
        pass
    if _is_sequence(values) and _is_sequence(mask):
        return [
            zero_mask_values(value_item, mask_item)
            for value_item, mask_item in zip(values, mask, strict=False)
        ]
    return values if bool(mask) else 0.0