transformer_lens.benchmarks.hook_structure module

Hook structure validation benchmarks.

This module provides structure-only validation of hooks. It checks hook existence, registration, firing, and shape compatibility without comparing activation values.

transformer_lens.benchmarks.hook_structure.benchmark_activation_cache_structure(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None, prepend_bos: bool | None = None) → BenchmarkResult

Benchmark activation cache for structural correctness (keys, shapes).

This checks: - Cache returns expected keys - Cache tensor shapes are compatible - run_with_cache works correctly

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing

  • reference_model – Optional HookedTransformer for comparison

  • prepend_bos – Whether to prepend BOS token. If None, uses model default.

Returns:

BenchmarkResult with structural validation details

transformer_lens.benchmarks.hook_structure.benchmark_backward_hooks_structure(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None, prepend_bos: bool | None = None) → BenchmarkResult

Benchmark backward hooks for structural correctness (existence, firing, shapes).

This checks: - All reference backward hooks exist in bridge - Hooks can be registered - Hooks fire during backward pass - Gradient tensor shapes are compatible

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing

  • reference_model – Optional HookedTransformer for comparison

  • prepend_bos – Whether to prepend BOS token. If None, uses model default.

Returns:

BenchmarkResult with structural validation details

transformer_lens.benchmarks.hook_structure.benchmark_forward_hooks_structure(bridge: TransformerBridge, test_text: str, reference_model: HookedTransformer | None = None, prepend_bos: bool | None = None) → BenchmarkResult

Benchmark forward hooks for structural correctness (existence, firing, shapes).

This checks: - All reference hooks exist in bridge - Hooks can be registered - Hooks fire during forward pass - Hook tensor shapes are compatible

Parameters:

  • bridge – TransformerBridge model to test

  • test_text – Input text for testing

  • reference_model – Optional HookedTransformer for comparison

  • prepend_bos – Whether to prepend BOS token. If None, uses model default.

Returns:

BenchmarkResult with structural validation details