docs: migrate executable doctests to TransformerBridge - #1576
docs: migrate executable doctests to TransformerBridge#1576Austin1serb wants to merge 3 commits into
Conversation
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for putting this together @Austin1serb! Looks great! Just one couple small notes below
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?
There was a problem hiding this comment.
Okay yeah. I revised the ActivationCache documentation to show it supports both HookedTransformer and TransformerBridge, and clarified that the remaining limitation is that the advanced helpers expect the TransformerLens weight-processing interface and a complete activation cache.
| >>> _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'] |
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?
There was a problem hiding this comment.
Got it.
I updated both examples so they show a better section of the cache instead of highlighting the alias
2 participants
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