transformer_lens.model_bridge.driver_protocol module

Driver protocol: the contract every model-execution backend satisfies.

class transformer_lens.model_bridge.driver_protocol.Driver(*args, **kwargs)

Bases: Protocol

The forward-pass contract. Hook installation is the driver’s problem.

forward has two dialects:

  • Module-replacement drivers (TransformersDriver): hooks fire via the bridge’s HookPoint system during the real torch forward, so capture/ intervene/max_new_tokens are not served here — conforming drivers raise NotImplementedError on them rather than silently ignore.

  • Spec drivers (vLLM, Inspect): no local module, so capture names hook points to record and intervene carries declarative edit specs; results come back in ForwardResult.captured.

architecture: str
bridge_config: TransformerBridgeConfig
close() → None
forward(input_ids: TensorLike | None = None, *, capture: tuple[str, ...] = (), intervene: Mapping[str, Callable[[TensorLike], TensorLike] | Mapping[str, Any]] | None = None, max_new_tokens: int = 1, return_logits: bool = True, **kwargs: Any) → ForwardResult
non_fireable_hook_points: frozenset[str]
provides_sequence_logits: bool
supported_hook_points: frozenset[str]
supports(feature: str) → bool

Capability flag over KNOWN_FEATURES. The bridge consults “parameters”; the others are caller-facing declarations.

tokenizer: Any
class transformer_lens.model_bridge.driver_protocol.ForwardResult(logits: ~transformer_lens.model_bridge.driver_protocol.TensorLike | None = None, captured: ~typing.Mapping[str, ~transformer_lens.model_bridge.driver_protocol.TensorLike] = <factory>, new_tokens: ~transformer_lens.model_bridge.driver_protocol.TensorLike | None = None, raw_output: ~typing.Any = None)

Bases: object

One forward call’s outputs. Tensors are native to the driver’s framework.

captured: Mapping[str, TensorLike]
logits: TensorLike | None = None
new_tokens: TensorLike | None = None
raw_output: Any = None
class transformer_lens.model_bridge.driver_protocol.TensorLike(*args, **kwargs)

Bases: Protocol

Quacks like a tensor: __array__ + shape + dtype.

property dtype: Any
property shape: Any
transformer_lens.model_bridge.driver_protocol.to_torch(t: TensorLike, *, dtype: dtype | None = None) → Tensor

Convert any TensorLike to torch.Tensor at the bridge boundary.

Order: torch passthrough → DLPack (jax/mlx/tf/cupy/numpy≥1.22, preserves device) → __array__ + from_numpy (CPU only).

transformer_lens.model_bridge.driver_protocol.validate_driver(driver: Any, *, after_bridge_construction: bool = False) → None

Stronger than isinstance(driver, Driver): checks types, signatures, and (optionally) post-construction state.

Parameters:

after_bridge_construction – when True, also requires at least one of supported_hook_points / non_fireable_hook_points non-empty (the bridge backfills the former, so empty-on-both means the driver silently degrades to “supports nothing”).

Raises:

TypeError – with a message naming the contract violation.