docs: migrate executable doctests to TransformerBridge - #1576
Open
Austin1serb wants to merge 2 commits into
Open
docs: migrate executable doctests to TransformerBridge#1576Austin1serb wants to merge 2 commits into
Austin1serb wants to merge 2 commits into
Conversation
jlarson4
reviewed
Jul 31, 2026
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for putting this together @Austin1serb! Looks great! Just one couple small notes below
| Warning: | ||
|
|
||
| :class:`ActivationCache` is designed to be used with | ||
| :class:`transformer_lens.HookedTransformer`, and will not work with other models. It's also |
Collaborator
There was a problem hiding this comment.
This block still says ActivationCache is designed for HookedTransformer and "will not work with other models" Can we update this comment and any surrounding examples?
| >>> _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'] |
Collaborator
There was a problem hiding this comment.
Two of these three keys are the same tensor, hook_embed is the compat alias for embed.hook_out. The slice shows duplication rather than cache structure, and __iter__ now prints an identical list below (line 318). Can the slice span enough keys to still teach the layout?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Migrates executable doctests in
ActivationCache,evals, andexploratory_utilsfrom the deprecatedHookedTransformerloader toTransformerBridge.The examples now use official Hugging Face model IDs and enable compatibility mode where their existing numerics and legacy hook aliases require it. Obsolete loader-output assertions were removed, and cache-key examples were updated to reflect the bridge's canonical hooks alongside compatibility aliases.
The skipped MMLU example was also migrated so no doctest in the three surviving modules constructs
HookedTransformer.Fixes #1565
Type of change
Checklist
Testing
uv run pytest --doctest-modules transformer_lens/ActivationCache.py transformer_lens/evals.py transformer_lens/utilities/exploratory_utils.py -q(9 passed,1 skipped)make docstring-test(17 passed,24 skipped)make unit-test(5139 passed,61 skipped,44 deselected,10 xfailed)uv run mypy .make check-formatgit diff --check