aquin.compute.attention_routing Namespace Reference

Functions
int	_count_induction_pairs (list[int] token_ids)
tuple[float, float]	_score_head (torch.Tensor pattern, list[int] token_ids)
list[dict]	_heads_from_pattern_stack (list[tuple[int, torch.Tensor]] patterns, list[int] token_ids)
list[tuple[int, torch.Tensor]]	_patterns_from_tl_cache (Any model, torch.Tensor tokens, int n_layers)
tuple[list[tuple[int, torch.Tensor]], str|None]	_patterns_from_hf_output_attentions (Any model, torch.Tensor tokens)
dict	run_attention_routing (Any model, str prompt, *, int top_k=5, str model_id="")

Function Documentation

_count_induction_pairs()

int _count_induction_pairs ( list[int] token_ids )

protected

Positions (q, j) where token[j] repeats at q and j+1 is valid.

Definition at line 19 of file attention_routing.py.

Referenced by run_attention_routing().

_heads_from_pattern_stack()

list[dict] _heads_from_pattern_stack ( list[tuple[int, torch.Tensor]] patterns, list[int] token_ids )

protected

patterns: list of (layer_idx, tensor [n_heads, seq, seq]).

Definition at line 53 of file attention_routing.py.

References _score_head().

Referenced by run_attention_routing().

_patterns_from_hf_output_attentions()

tuple[list[tuple[int, torch.Tensor]], str | None] _patterns_from_hf_output_attentions ( Any model, torch.Tensor tokens )

protected

HF-native / HfLlmShim path.
TransformerLens patterns are unavailable on hf_only models (LFM, …).

Definition at line 98 of file attention_routing.py.

Referenced by run_attention_routing().

_patterns_from_tl_cache()

list[tuple[int, torch.Tensor]] _patterns_from_tl_cache ( Any model, torch.Tensor tokens, int n_layers )

protected

Definition at line 81 of file attention_routing.py.

Referenced by run_attention_routing().

_score_head()

tuple[float, float] _score_head ( torch.Tensor pattern, list[int] token_ids )

protected

pattern: (seq, seq) causal attn weights for one head.

Definition at line 29 of file attention_routing.py.

Referenced by _heads_from_pattern_stack().

run_attention_routing()

dict run_attention_routing	(	Any	model,
		str	prompt,
		*	,
		int	top_k = 5,
		str	model_id = "" )

Definition at line 142 of file attention_routing.py.

References _count_induction_pairs(), _heads_from_pattern_stack(), _patterns_from_hf_output_attentions(), and _patterns_from_tl_cache().