transformer_lens.model_bridge.bridge_core module

Framework-agnostic bridge surface shared by TransformerBridge and RemoteBridge.

class transformer_lens.model_bridge.bridge_core.BridgeCore(adapter: ArchitectureAdapter, tokenizer: Any, driver: Any)

Bases: object

Framework-agnostic bridge surface: hooks, cache, run_with_*, driver wiring.

Holds state shared by every bridge (adapter, cfg, tokenizer, driver, hook registries). Subclasses add framework-specific state — TransformerBridge walks the wrapped nn.Module; RemoteBridge builds components from adapter metadata.

__init__(adapter: ArchitectureAdapter, tokenizer: Any, driver: Any) → None

Subclasses call this AFTER nn.Module.__init__ (if applicable), then do framework-specific setup.

add_caching_hooks(names_filter: str | List[str] | Callable[[str], bool] | None = None, incl_bwd: bool = False, device: Any = None, remove_batch_dim: bool = False, cache: dict | None = None) → dict

Attach caching hooks to the model (does not run it). Returns the cache dict.

Mirrors HookedRootModule.add_caching_hooks. The hooks persist until reset_hooks().

add_hook(name: str | Callable[[str], bool], hook_fn: Any, dir: Literal['fwd', 'bwd'] = 'fwd', is_permanent: bool = False) → None

Add a hook to a specific component or to all components matching a filter.

Parameters:

  • name – Either a string hook point name (e.g. “blocks.0.attn.hook_q”) or a callable filter (str) -> bool that is applied to every hook point name; the hook is added to each point where the filter returns True.

  • hook_fn – The hook function (activation, hook) -> activation | None.

  • dir – Hook direction, "fwd" or "bwd".

  • is_permanent – If True the hook survives reset_hooks() calls.

add_perma_hook(name: str | Callable[[str], bool], hook_fn: Callable, dir: Literal['fwd', 'bwd'] = 'fwd') → None

Add a permanent hook that survives reset_hooks() calls.

Convenience wrapper for add_hook(..., is_permanent=True). To remove, call reset_hooks(including_permanent=True) or remove from the underlying HookPoint directly.

cache_all(cache: dict | None, incl_bwd: bool = False, device: Any = None, remove_batch_dim: bool = False) → None

Deprecated: cache every activation. Use run_with_cache / add_caching_hooks.

cache_some(cache: dict | None, names: Callable[[str], bool], incl_bwd: bool = False, device: Any = None, remove_batch_dim: bool = False) → None

Deprecated: cache activations matching names. Use run_with_cache.

check_hooks_to_add(hook_point: HookPoint, hook_point_name: str, hook: Callable, dir: Literal['fwd', 'bwd'] = 'fwd', is_permanent: bool = False, prepend: bool = False) → None

Validate a hook before it is added; override to add checks.

Raises for a gated-off hook point — a targeted attach there would silently never fire. Every explicit-name attach path routes through here; filter sweeps pre-skip gated matches (with a warning) before reaching it, since a filter was not necessarily targeting them.

Gating keys on the POINT’s own canonical name, not the requested spelling: on adapters with hook_alias_overrides a gated HT name (e.g. blocks.0.hook_mlp_in on BERT) resolves to an always-firing point, and refusing it would reject a hook that works. Driver-fireability is enforced separately in _check_hook_fireable.

clear_contexts() → None

Clear the stored ctx on every hook point.

clear_hook_registry() → None

Clear the hook registry and force re-initialization.

close() → None

Release driver-managed resources. Idempotent — safe to call multiple times.

forward(*args: Any, **kwargs: Any) → Any

Subclasses implement how the driver gets called.

get_caching_hooks(names_filter: str | List[str] | Callable[[str], bool] | None = None, incl_bwd: bool = False, device: Any = None, remove_batch_dim: bool = False, cache: dict | None = None, pos_slice: Slice | int | Tuple[int] | Tuple[int, int] | Tuple[int, int, int] | List[int] | Tensor | ndarray | None = None) → Tuple[dict, list, list]

Build caching hooks without adding them. Mirrors HookedRootModule.get_caching_hooks.

Returns (cache, fwd_hooks, bwd_hooks) where each hook is a (name, hook_fn) pair suitable for hooks() / run_with_hooks. Activations are keyed by the HookPoint’s canonical name; backward hooks append "_grad". bwd_hooks is empty unless incl_bwd.

get_hook_point(hook_name: str) → HookPoint | None

Get a hook point by name from the bridge’s hook system.

hook_aliases: Dict[str, str | List[str]] = {'hook_embed': ['embed_ln.hook_out', 'embed.hook_out'], 'hook_pos_embed': ['pos_embed.hook_out', 'rotary_emb.hook_out'], 'hook_unembed': 'unembed.hook_out'}
property hook_dict: dict[str, HookPoint]

All HookPoint objects, including aliases — TransformerLens-compatible.

hook_points() → Iterable[HookPoint]

All HookPoint instances (registry is canonical and complete).

hooks(fwd_hooks: List = [], bwd_hooks: List = [], reset_hooks_end: bool = True, clear_contexts: bool = False) → Any

Context manager for temporarily adding hooks.

reset_hooks_end removes the hooks this context added when it exits — hooks the caller attached beforehand are left alone either way; clear_contexts also wipes the touched hook points’ ctx.

Example

with model.hooks(fwd_hooks=[(“hook_embed”, my_hook)]):

output = model(“Hello world”)

loss_fn(logits: Tensor, tokens: Tensor, attention_mask: Tensor | None = None, per_token: bool = False) → Tensor

Cross-entropy loss matching HookedTransformer’s formula (log_softmax + gather).

property mod_dict: Dict[str, Any]

Module/hook name -> object, HookedRootModule-compatible.

Union of the named-module tree and the aliased hook view, so both canonical (blocks.0.mlp.hook_out) and HT-style (blocks.0.hook_mlp_out) names resolve to the same HookPoint.

remove_all_hook_fns(direction: Literal['fwd', 'bwd', 'both'] = 'both', including_permanent: bool = False, level: int | None = None) → None

Remove hook functions from every hook point.

reset_hooks(clear_contexts: bool = True, direction: Literal['fwd', 'bwd', 'both'] = 'both', including_permanent: bool = False, level: int | None = None) → None

Remove hooks from every hook point; mirrors HookedRootModule.reset_hooks.

The hook registry is canonical and complete (every component’s HookPoint is registered), so a single pass covers the whole model.

Parameters:

  • clear_contexts – Also clear each hook point’s stored ctx.

  • direction – Which direction(s) to remove — "fwd", "bwd", or "both".

  • including_permanent – If True, also remove hooks added via add_perma_hook.

  • level – If set, only remove hooks registered at this context level.

run_with_cache(input: str | List[str] | Tensor, return_cache_object: Literal[True] = True, remove_batch_dim: bool = False, **kwargs) → Tuple[Any, ActivationCache]
run_with_cache(input: str | List[str] | Tensor, return_cache_object: Literal[False], remove_batch_dim: bool = False, **kwargs) → Tuple[Any, Dict[str, Tensor]]

Run the model and cache activations. Returns (output, cache).

stop_at_layer raises StopAtLayerException to stop early. start_at_layer treats input as the residual entering block k (see forward()); blocks below k are excluded from the cache to match HookedTransformer. pos_slice slices each cached activation along its position dimension (dim 1 for resid/per-head/token-id activations; the query position -2 for attention patterns/scores). incl_bwd also caches gradients under "<name>_grad" by running output.backward(); the caller must request a scalar output (return_type="loss") and the model must be on the gradients-capable transformers driver. reset_hooks_end removes the hooks this call added when it finishes — hooks the caller attached beforehand are left alone either way; clear_contexts also wipes the touched hook points’ ctx. device offloads cached activations (matches ActivationCache.to); the model and inputs stay where the caller put them.

run_with_hooks(input: Any, fwd_hooks: List[Tuple[str | Callable, Callable]] = [], bwd_hooks: List[Tuple[str | Callable, Callable]] = [], reset_hooks_end: bool = True, clear_contexts: bool = False, return_type: str | None = 'logits', stop_at_layer: int | None = None, start_at_layer: int | None = None, remove_batch_dim: bool = False, **kwargs: Any) → Any

Run the model with specified forward and backward hooks.

stop_at_layer raises StopAtLayerException to stop early (KV cache cleaned up on stop). start_at_layer treats input as the residual entering block k (see forward()); hooks on blocks below k are skipped to match HookedTransformer. remove_batch_dim squeezes/unsqueezes the batch dim around hook callbacks (batch_size==1 only). reset_hooks_end removes the hooks this call added when it finishes — hooks the caller attached beforehand are left alone either way; clear_contexts also wipes the touched hook points’ ctx.

to_tokens(*args: Any, **kwargs: Any) → Any

Subclasses implement against their tokenizer surface.

property tokenizer: Any

The tokenizer used for encoding/decoding text.

transformer_lens.model_bridge.bridge_core.build_alias_to_canonical_map(hook_dict: Any, prefix: str = '') → dict

Map alias hook names to their canonical names (where .name differs from the key).