Skip to content

docs: migrate executable doctests to TransformerBridge - #1576

Open
Austin1serb wants to merge 2 commits into
TransformerLensOrg:dev-4.xfrom
Austin1serb:docs/1565-bridge-doctests
Open

docs: migrate executable doctests to TransformerBridge#1576
Austin1serb wants to merge 2 commits into
TransformerLensOrg:dev-4.xfrom
Austin1serb:docs/1565-bridge-doctests

Conversation

@Austin1serb

Copy link
Copy Markdown

Description

Migrates executable doctests in ActivationCache, evals, and exploratory_utils from the deprecated HookedTransformer loader to TransformerBridge.

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

  • This change requires a documentation update

Checklist

  • No additional explanatory source comments were required
  • The executable documentation was updated
  • My changes generate no new warnings
  • The focused and full docstring suites pass locally
  • New and existing unit tests pass locally
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

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-format
  • git diff --check

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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']

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants