Skip to content

docs: migrate executable doctests to TransformerBridge - #1576

Merged
jlarson4 merged 4 commits into
TransformerLensOrg:dev-4.xfrom
Austin1serb:docs/1565-bridge-doctests
Aug 6, 2026
Merged

docs: migrate executable doctests to TransformerBridge#1576
jlarson4 merged 4 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

Comment thread transformer_lens/ActivationCache.py Outdated
Comment thread transformer_lens/ActivationCache.py Outdated

@Canonik Canonik left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Ran the doctests from this branch locally to double-check. ActivationCache.py and
exploratory_utils.py pass as-is (nice, the compat bridge reproduces the old float table
in test_prompt exactly). The IOIDataset one breaks on a CUDA machine though:

RuntimeError: Expected all tensors to be on the same device, but got index is on cpu,
different from other tensors on cuda:0

boot_transformers("gpt2") with no device arg picks cuda when available, but ioi_eval
feeds CPU tensors straight from its DataLoader, and the bridge forward passes input_ids
through to the HF model unchanged. The old HookedTransformer version wasn't affected
because input_to_embed moves tokens to the model device (HookedTransformer.py:412).
CI is green here only because the runners are CPU-only. Easiest fix:

>>> model = TransformerBridge.boot_transformers("gpt2", device="cpu")

same pattern as the gpt2_bridge_compat fixture in tests/conftest.py. Verified locally
that with device="cpu" the doctest passes end to end on this branch.

Two small things while you're in these files:

  • the comment at the bottom of exploratory_utils.py still says the docstring "won't be
    tested with PyTest". Not true: --collect-only shows the doctest item is collected and
    run; the pytest.mark.skip only hits the accidentally-collected function item. Worth
    fixing the comment since this PR makes that doctest live on the bridge.
  • ActivationCache.py:63 still cross-references
    :meth:`transformer_lens.HookedTransformer.HookedTransformer.run_with_cache` , which
    will dangle once HookedTransformer is removed. Could point at the bridge's
    run_with_cache instead.

For anyone looking at the red check: Full Code Coverage fails on
test_optimizer_compatibility.py::test_bridge_hooked_parity_multi_step_optimization at
0.001032 vs a 0.001 threshold, a numerics flake unrelated to this doctest-only diff.

Rest looks right to me: dropping the "Loaded pretrained model ..." lines matches boot
printing nothing, the compat-mode choices line up with the acceptance-test precedent
(test_activation_cache.py, test_evals.py), and the HF repo ids are the documented
preference over the deprecated aliases.

@Austin1serb

Copy link
Copy Markdown
Author

I have addressed the requested changes below:

  1. Pinned the IOI doctest to device="cpu" to avoid the device mismatch.
  2. Updated the exploratory_utils.py comment to match the current doctest behavior.
  3. Fixed the ActivationCache cross-reference to point to TransformerBridge.run_with_cache.

@jlarson4

jlarson4 commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

Great work on this @Austin1serb, looks great, merging it now!

@jlarson4
jlarson4 merged commit 69f5bd2 into TransformerLensOrg:dev-4.x Aug 6, 2026
49 of 50 checks passed
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.

[Proposal] Rewrite executable doctests that construct HookedTransformer onto TransformerBridge

3 participants