From d07fdd34dcc2f21f2c417a8ace79851b7649674e Mon Sep 17 00:00:00 2001 From: Austin1serb Date: Thu, 30 Jul 2026 23:00:19 -0600 Subject: [PATCH 1/3] docs: migrate executable doctests to TransformerBridge --- transformer_lens/ActivationCache.py | 28 +++++++++---------- transformer_lens/evals.py | 11 ++++---- .../utilities/exploratory_utils.py | 7 +++-- 3 files changed, 24 insertions(+), 22 deletions(-) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index d40fcc109..c2bd9f96a 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -71,9 +71,9 @@ class ActivationCache: the model predicting "road". This kind of analysis commonly falls under the category of "logit attribution" or "direct logit attribution" (DLA). - >>> from transformer_lens import HookedTransformer - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Why did the chicken cross the") >>> residual_stream, labels = cache.decompose_resid(return_labels=True, mode="attn") @@ -270,12 +270,12 @@ def keys(self): Examples: - >>> from transformer_lens import HookedTransformer - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Some prompt") >>> list(cache.keys())[0:3] - ['hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre'] + ['embed.hook_in', 'hook_embed', 'embed.hook_out'] Returns: List of all keys. @@ -306,16 +306,16 @@ def __iter__(self) -> Iterator[str]: Examples: - >>> from transformer_lens import HookedTransformer - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Some prompt") >>> cache_interesting_names = [] >>> for key in cache: ... if not key.startswith("blocks.") or key.startswith("blocks.0"): ... cache_interesting_names.append(key) >>> print(cache_interesting_names[0:3]) - ['hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre'] + ['embed.hook_in', 'hook_embed', 'embed.hook_out'] Returns: Iterator over the cache. @@ -392,12 +392,12 @@ def accumulated_resid( Logit Lens analysis can be done as follows: - >>> from transformer_lens import HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge >>> import torch >>> import pandas as pd - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M", device="cpu", fold_ln=True) - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M", device="cpu") + >>> model.enable_compatibility_mode() >>> prompt = "Why did the chicken cross the" >>> answer = " road" diff --git a/transformer_lens/evals.py b/transformer_lens/evals.py index 8423f301e..74512cf8a 100644 --- a/transformer_lens/evals.py +++ b/transformer_lens/evals.py @@ -323,10 +323,10 @@ class IOIDataset(Dataset): .. code-block:: python >>> from transformer_lens.evals import ioi_eval, IOIDataset - >>> from transformer_lens.HookedTransformer import HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge - >>> model = HookedTransformer.from_pretrained('gpt2-small') - Loaded pretrained model gpt2-small into HookedTransformer + >>> model = TransformerBridge.boot_transformers("gpt2") + >>> model.enable_compatibility_mode() >>> # Evaluate on a deterministic dataset (seed makes results reproducible) >>> ds = IOIDataset(tokenizer=model.tokenizer, num_samples=100, seed=42) @@ -551,10 +551,11 @@ def mmlu_eval( .. code-block:: python - >>> from transformer_lens import HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge >>> from transformer_lens.evals import mmlu_eval - >>> model = HookedTransformer.from_pretrained("gpt2-small") # doctest: +SKIP + >>> model = TransformerBridge.boot_transformers("gpt2") # doctest: +SKIP + >>> model.enable_compatibility_mode() # doctest: +SKIP >>> results = mmlu_eval(model, subjects="abstract_algebra", num_samples=10) # doctest: +SKIP >>> print(f"Accuracy: {results['accuracy']:.2%}") # doctest: +SKIP """ diff --git a/transformer_lens/utilities/exploratory_utils.py b/transformer_lens/utilities/exploratory_utils.py index 426aee6a0..cfa1721f6 100644 --- a/transformer_lens/utilities/exploratory_utils.py +++ b/transformer_lens/utilities/exploratory_utils.py @@ -31,9 +31,10 @@ def test_prompt( Examples: - >>> from transformer_lens import HookedTransformer, utilities - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens import utilities + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> prompt = "Why did the elephant cross the" >>> answer = "road" From d5413641597e94dbe884cc3b501c6cbf35691743 Mon Sep 17 00:00:00 2001 From: Austin1serb Date: Wed, 5 Aug 2026 05:32:36 -0600 Subject: [PATCH 2/3] docs: address ActivationCache review feedback --- transformer_lens/ActivationCache.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index c2bd9f96a..5580ea87e 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -101,9 +101,10 @@ class ActivationCache: Warning: :class:`ActivationCache` is designed to be used with - :class:`transformer_lens.HookedTransformer`, and will not work with other models. It's also - designed to be used with all activations of :class:`transformer_lens.HookedTransformer` being - cached, and some internal methods will break without that. + :class:`transformer_lens.HookedTransformer` or + :class:`transformer_lens.model_bridge.TransformerBridge`. Advanced helpers expect the model to + expose the TransformerLens weight-processing interface and generally expect a complete cache; + some internal methods may break with other models or partial caches. The biggest footgun and source of bugs in this code will be keeping track of indexes, dimensions, and the numbers of each. There are several kinds of activations: @@ -274,8 +275,10 @@ def keys(self): >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Some prompt") - >>> list(cache.keys())[0:3] - ['embed.hook_in', 'hook_embed', 'embed.hook_out'] + >>> list(cache.keys())[0:8] + ['embed.hook_in', 'hook_embed', 'embed.hook_out', 'pos_embed.hook_in', + 'hook_pos_embed', 'pos_embed.hook_out', 'blocks.0.hook_in', + 'blocks.0.hook_resid_pre'] Returns: List of all keys. @@ -314,8 +317,10 @@ def __iter__(self) -> Iterator[str]: >>> for key in cache: ... if not key.startswith("blocks.") or key.startswith("blocks.0"): ... cache_interesting_names.append(key) - >>> print(cache_interesting_names[0:3]) - ['embed.hook_in', 'hook_embed', 'embed.hook_out'] + >>> print(cache_interesting_names[0:8]) + ['embed.hook_in', 'hook_embed', 'embed.hook_out', 'pos_embed.hook_in', + 'hook_pos_embed', 'pos_embed.hook_out', 'blocks.0.hook_in', + 'blocks.0.hook_resid_pre'] Returns: Iterator over the cache. From c4333927d9b5c103696466828c1c5a14796ef2cf Mon Sep 17 00:00:00 2001 From: Austin1serb Date: Wed, 5 Aug 2026 09:18:39 -0600 Subject: [PATCH 3/3] docs: address CUDA doctest review --- transformer_lens/ActivationCache.py | 2 +- transformer_lens/evals.py | 2 +- transformer_lens/utilities/exploratory_utils.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 5580ea87e..ed16cf322 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -60,7 +60,7 @@ class ActivationCache: The :class:`ActivationCache` is at the core of Transformer Lens. It is a wrapper that stores all important activations from a forward pass of the model, and provides a variety of helper functions to investigate them. The common way to access it is to run the model with - :meth:`transformer_lens.HookedTransformer.HookedTransformer.run_with_cache`. + :meth:`transformer_lens.model_bridge.TransformerBridge.run_with_cache`. Examples: diff --git a/transformer_lens/evals.py b/transformer_lens/evals.py index 74512cf8a..ee093a121 100644 --- a/transformer_lens/evals.py +++ b/transformer_lens/evals.py @@ -325,7 +325,7 @@ class IOIDataset(Dataset): >>> from transformer_lens.evals import ioi_eval, IOIDataset >>> from transformer_lens.model_bridge import TransformerBridge - >>> model = TransformerBridge.boot_transformers("gpt2") + >>> model = TransformerBridge.boot_transformers("gpt2", device="cpu") >>> model.enable_compatibility_mode() >>> # Evaluate on a deterministic dataset (seed makes results reproducible) diff --git a/transformer_lens/utilities/exploratory_utils.py b/transformer_lens/utilities/exploratory_utils.py index cfa1721f6..3e69d53b9 100644 --- a/transformer_lens/utilities/exploratory_utils.py +++ b/transformer_lens/utilities/exploratory_utils.py @@ -157,8 +157,8 @@ def test_prompt( try: import pytest - # Note: Docstring won't be tested with PyTest (it's ignored), as it thinks this is a regular unit - # test (because its name is prefixed `test_`). + # Note: PyTest collects and runs this docstring as a doctest. The skip marker only applies to the + # accidentally collected function item (because its name is prefixed `test_`). pytest.mark.skip(test_prompt) except ModuleNotFoundError: pass # disregard if pytest not in env