transformer_lens.benchmarks.encoder_common module

Shared benchmark implementations for non-text encoder bridges.

Audio and vision encoders take one raw tensor (waveform, spectrogram, or pixels) where text models take token ids; given that tensor, the forward, run_with_cache, and perturbation-stability checks are modality-independent. The modality modules (audio.py, vision.py) wrap these with their result names, HF reference-forward kwarg, and critical component lists.

transformer_lens.benchmarks.encoder_common.benchmark_encoder_cache(bridge: TransformerBridge, test_input: Tensor, name: str, critical_components: Sequence[str], min_found: int = 3) → BenchmarkResult

run_with_cache() benchmark on a raw modality tensor.

Verifies that critical hooks fire and produce valid tensors: the critical_components this architecture actually declares in its component mapping, plus the first and last block.

Parameters:

  • bridge – TransformerBridge model to test

  • test_input – Raw modality tensor (waveform, spectrogram, or pixels)

  • name – Result name (e.g. “audio_cache”, “vision_cache”)

  • critical_components – Component-mapping names whose hook_out must be cached when the architecture declares them

  • min_found – Minimum critical hooks present to still pass with a warning

transformer_lens.benchmarks.encoder_common.benchmark_encoder_forward(bridge: TransformerBridge, test_input: Tensor, name: str, ref_input_key: str, reference_model: Module | None = None) → BenchmarkResult

Forward-pass benchmark on a raw modality tensor.

Compares bridge output against the HF native model on the same input when a reference is given. Bare encoders compare last_hidden_state; head models (CTC, classification) compare logits.

Parameters:

  • bridge – TransformerBridge model to test

  • test_input – Raw modality tensor (waveform, spectrogram, or pixels)

  • name – Result name (e.g. “audio_forward”, “vision_forward”)

  • ref_input_key – Kwarg the HF reference takes the tensor under (“input_values” for audio, “pixel_values” for vision)

  • reference_model – Optional HF reference model for comparison

transformer_lens.benchmarks.encoder_common.benchmark_encoder_representation_stability(bridge: TransformerBridge, test_input: Tensor, name: str) → BenchmarkResult

Representation stability under small input perturbations.

Similar inputs should produce similar hidden states. Skipped for tiny-random models (random weights won’t produce stable representations).

Parameters:

  • bridge – TransformerBridge model to test

  • test_input – Raw modality tensor (waveform, spectrogram, or pixels)

  • name – Result name (e.g. “audio_representation_stability”)

transformer_lens.benchmarks.encoder_common.extract_encoder_states(out) → Tensor | None

Pull the hidden-state/logit tensor out of a tensor, BaseModelOutput, or head output.