interlens.interp¶

def capture_activations(model: "PreTrainedModel", input_ids: torch.Tensor, spec: CaptureSpec) -> list[CapturedSite]:
	"""Run one clean forward pass over ``input_ids`` and return ``[(layer, site, tensor[seq, d_model])]``.

	Design choice: capture is a *separate forward pass* over the full (prompt + generated) sequence rather than
	accumulating hooks across the multi-step decode loop. This is simpler and provably complete — every position
	is present in one pass — at the cost of one extra forward. Residual-stream activations come from
	``output_hidden_states`` (no hooks needed); ``attn``/``mlp`` sublayer outputs come from forward hooks on the
	corresponding submodules.

	Note: ``attn`` captures the attention *sublayer output* (post-o_proj), which is available under any attention
	backend. Attention *weights/patterns* are NOT captured here — those require ``attn_implementation='eager'`` +
	``output_attentions`` and are out of scope for the default kernel.
	"""
	layers = decoder_layers(model)
	want = spec.layers if spec.layers is not None else tuple(range(len(layers)))
	sites = set(spec.sites)

	results: list[CapturedSite] = []
	handles = []
	hook_store: dict[tuple[int, Site], torch.Tensor] = {}

	def make_hook(layer_idx: int, site: Site):
		def hook(module, inputs, output):
			hs = output[0] if isinstance(output, tuple) else output
			hook_store[(layer_idx, site)] = hs.detach()[0]  # [seq, d_model]
		return hook

	for li in want:
		if "attn" in sites:
			handles.append(layers[li].self_attn.register_forward_hook(make_hook(li, "attn")))
		if "mlp" in sites:
			handles.append(layers[li].mlp.register_forward_hook(make_hook(li, "mlp")))

	need_hidden = "residual" in sites
	try:
		with torch.inference_mode():
			out = model(input_ids, output_hidden_states=need_hidden, use_cache=False)
	finally:
		for h in handles:
			h.remove()

	if need_hidden:
		# hidden_states is (embeddings, layer_0_out, ..., layer_{L-1}_out); layer li output is index li+1.
		hs = out.hidden_states
		for li in want:
			results.append(CapturedSite(li, "residual", hs[li + 1][0]))
	for (li, site), tensor in hook_store.items():
		results.append(CapturedSite(li, site, tensor))

	return results

def capture_router_logits(model: "PreTrainedModel", input_ids: torch.Tensor, layers: tuple[int, ...] | None = None,
                          top_k_only: bool = False, offload: str = "cpu") -> list[RoutingCapture]:
	"""One clean forward pass over ``input_ids`` (``[1, seq]``); return per-MoE-layer ``RoutingCapture``.

	Parameters:
		model: an HF MoE causal LM whose ``forward`` accepts ``output_router_logits=True`` (OLMoE, Qwen-MoE,
			Mixtral, ...).
		input_ids: full token sequence to route, shape ``[1, seq]`` (batch of 1 — replay one view at a time).
		layers: decoder-layer indices to keep (default: all sparse layers, per ``moe_layer_indices``).
		top_k_only: drop the full ``[seq, n_experts]`` logits and keep only top-k ids/probs (int16/fp16,
			~8x smaller — use for long-sequence sweeps over large MoEs like Qwen3-30B-A3B).
		offload: device for the returned tensors (default ``"cpu"`` so GPU memory is freed immediately).
	"""
	sparse = moe_layer_indices(model)
	want = set(layers if layers is not None else sparse)
	k = moe_topk(model)

	with torch.inference_mode():
		out = model(input_ids.to(model.device), output_router_logits=True, use_cache=False)
	if getattr(out, "router_logits", None) is None:
		raise ValueError(f"{type(model).__name__} did not return router_logits — not an MoE forward?")

	results: list[RoutingCapture] = []
	for layer_idx, logits in zip(sparse, out.router_logits):
		if layer_idx not in want:
			continue
		logits = logits.detach().float()  # [seq, n_experts] (HF flattens batch=1 into seq)
		if logits.dim() == 3:
			logits = logits[0]
		probs = torch.softmax(logits, dim=-1)
		topk_probs, topk_ids = probs.topk(k, dim=-1)
		results.append(RoutingCapture(
			layer=layer_idx,
			router_logits=None if top_k_only else logits.to(offload),
			topk_experts=topk_ids.to(torch.int16).to(offload),
			topk_probs=topk_probs.to(torch.float16).to(offload),
		))
	return results

def continuation_logprob(
	model: "PreTrainedModel",
	*,
	target_ids: torch.Tensor,
	prefix_ids: torch.Tensor | None = None,
	prefix_embeds: torch.Tensor | None = None,
	attention_mask: torch.Tensor | None = None,
	reduction: str = "mean",
	checkpoint: bool = False,
) -> torch.Tensor:
	"""Differentiable teacher-forced logprob of ``target_ids`` continuing a prefix, under ``model``.

	The grad-enabled analog of ``train_to_steer/rl_elicit.py::cont_logprob`` (which is ``@torch.no_grad`` and returns
	a float). Supply the prefix as ``prefix_ids`` (``[batch, P]`` long) OR ``prefix_embeds`` (``[batch, P, d_model]``
	float — e.g. a learnable soft prompt, or soft embeddings bridged from model A). ``target_ids`` is ``[batch, T]``
	long (or ``[T]`` broadcast to batch 1). The prefix and target are embedded via ``model.get_input_embeddings()``
	and run in one forward; the returned scalar (or per-example vector, see ``reduction``) is the logprob the model
	assigns to the exact target tokens, still attached to the graph.

	``reduction``: ``"mean"`` (mean over target tokens, mean over batch -> scalar; the reward used by rl_elicit),
	``"sum"`` (sum over target tokens, mean over batch), or ``"none"`` (``[batch, T]`` per-token logprobs). Use
	``"mean"`` as a length-normalized elicitation reward, ``"sum"`` when total sequence likelihood matters, ``"none"``
	for custom weighting. ``checkpoint`` is forwarded to save memory when backpropagating through a large frozen B.
	"""
	if (prefix_ids is None) == (prefix_embeds is None):
		raise ValueError("pass exactly one of prefix_ids or prefix_embeds")
	embed = model.get_input_embeddings()
	device = embed.weight.device
	if target_ids.dim() == 1:
		target_ids = target_ids.unsqueeze(0)
	target_ids = target_ids.to(device)

	if prefix_embeds is None:
		if prefix_ids.dim() == 1:
			prefix_ids = prefix_ids.unsqueeze(0)
		prefix_embeds = embed(prefix_ids.to(device))
	prefix_embeds = prefix_embeds.to(device)
	if prefix_embeds.shape[0] == 1 and target_ids.shape[0] > 1:
		prefix_embeds = prefix_embeds.expand(target_ids.shape[0], -1, -1)

	target_embeds = embed(target_ids)
	full = torch.cat([prefix_embeds, target_embeds], dim=1)
	fout = forward_with_grad(model, inputs_embeds=full, attention_mask=attention_mask, checkpoint=checkpoint)

	P = prefix_embeds.shape[1]
	T = target_ids.shape[1]
	# Logits at position i predict token i+1; the T target tokens sit at positions P..P+T-1, so the logits that
	# predict them are at positions P-1..P+T-2.
	pred_logits = fout.logits[:, P - 1: P + T - 1, :].float()
	logp = torch.log_softmax(pred_logits, dim=-1)
	tok_logp = logp.gather(2, target_ids.unsqueeze(2)).squeeze(2)  # [batch, T]

	if reduction == "none":
		return tok_logp
	if reduction == "sum":
		return tok_logp.sum(dim=1).mean()
	if reduction == "mean":
		return tok_logp.mean()
	raise ValueError(f"unknown reduction {reduction!r}")

def decoder_layers(model: "PreTrainedModel"):
	"""Return the list of transformer decoder layer modules, across common HF architectures.

	Steering/patching hooks and per-layer capture all need the ordered layer stack. Qwen and Gemma both expose
	it at ``model.model.layers``; a few fallbacks cover other families so the interp layer isn't Qwen/Gemma-only.
	Multimodal image-text-to-text wrappers (Qwen 3.5, Gemma 4, …) nest the text decoder one level deeper — its
	layers live at ``model.language_model.layers`` (or ``.model.language_model.model.layers``) — so those paths are
	included too. PEFT/adapter wrappers (``PeftModelForCausalLM``) are unwrapped first so hooks land on the real
	decoder layers.
	"""
	if hasattr(model, "get_base_model"):  # peft wrapper -> underlying transformers model
		try:
			model = model.get_base_model()
		except Exception:
			pass
	for path in ("model.layers", "transformer.h", "gpt_neox.layers", "model.decoder.layers",
	             "model.language_model.layers", "language_model.layers", "model.language_model.model.layers"):
		obj = model
		try:
			for attr in path.split("."):
				obj = getattr(obj, attr)
			return obj
		except AttributeError:
			continue
	raise AttributeError(f"could not locate decoder layers on {type(model).__name__}")

def forward_with_grad(
	model: "PreTrainedModel",
	*,
	input_ids: torch.Tensor | None = None,
	inputs_embeds: torch.Tensor | None = None,
	attention_mask: torch.Tensor | None = None,
	capture: GradCaptureSpec | None = None,
	checkpoint: bool = False,
) -> GradForwardOutput:
	"""One grad-connected forward pass over ``input_ids`` XOR ``inputs_embeds``.

	Exactly one of ``input_ids`` (``[batch, seq]`` long) or ``inputs_embeds`` (``[batch, seq, d_model]`` float,
	typically the output of ``bridge.soft_embed`` or a learnable soft prompt) must be given. ``attention_mask`` is
	optional (``[batch, seq]``). Returns ``logits`` and any requested hidden activations, all still attached to the
	graph so ``.backward()`` flows back to ``inputs_embeds`` / soft-prompt params / an upstream model A.

	``checkpoint=True`` enables HF gradient checkpointing for this call (``use_reentrant=False``) and restores the
	prior setting afterwards; it recomputes layer activations on the backward pass to trade compute for memory, and
	produces identical gradients. Residual sites come from ``output_hidden_states`` (grad-connected); ``attn``/``mlp``
	sites come from non-detaching forward hooks on the submodules.
	"""
	if (input_ids is None) == (inputs_embeds is None):
		raise ValueError("pass exactly one of input_ids or inputs_embeds")

	sites: set[Site] = set(capture.sites) if capture is not None else set()
	need_hidden = "residual" in sites
	need_hooks = bool(sites - {"residual"})  # attn/mlp need submodule hooks; residual comes from hidden_states

	handles = []
	hook_store: dict[tuple[Site, int], torch.Tensor] = {}

	def make_hook(layer_idx: int, site: Site):
		def hook(module, inputs, output):
			hs = output[0] if isinstance(output, tuple) else output
			hook_store[(site, layer_idx)] = hs  # keep grad — do NOT detach
		return hook

	if need_hooks:  # only touch the layer stack when a submodule hook is actually required
		layers = decoder_layers(model)
		hook_layers = capture.layers if capture.layers is not None else tuple(range(len(layers)))
		for li in hook_layers:
			if "attn" in sites:
				handles.append(layers[li].self_attn.register_forward_hook(make_hook(li, "attn")))
			if "mlp" in sites:
				handles.append(layers[li].mlp.register_forward_hook(make_hook(li, "mlp")))

	was_checkpointing = getattr(model, "is_gradient_checkpointing", False)
	try:
		if checkpoint and not was_checkpointing:
			model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
		out = model(
			input_ids=input_ids,
			inputs_embeds=inputs_embeds,
			attention_mask=attention_mask,
			output_hidden_states=need_hidden,
			use_cache=False,
		)
	finally:
		for h in handles:
			h.remove()
		if checkpoint and not was_checkpointing:
			model.gradient_checkpointing_disable()

	hidden: dict[tuple[Site, int], torch.Tensor] = {}
	if need_hidden:
		# hidden_states = (embeddings, layer_0_out, ..., layer_{L-1}_out); layer li output is index li+1.
		hs = out.hidden_states
		want = capture.layers if capture.layers is not None else tuple(range(len(hs) - 1))
		for li in want:
			hidden[("residual", li)] = hs[li + 1]
	hidden.update(hook_store)
	return GradForwardOutput(logits=out.logits, hidden=hidden)

def gumbel_softmax_tokens(logits: torch.Tensor, tau: float = 1.0, hard: bool = False) -> torch.Tensor:
	"""Relaxed (Gumbel-softmax) sample over the last (vocab) dim of ``logits``, differentiable in ``logits``.

	``tau`` is the temperature: high -> smooth mixture (biased but low-variance gradients), low -> near-one-hot
	(faithful to discrete sampling but higher-variance). Anneal ``tau`` down over training to move from an easy soft
	optimization toward the discrete tokens B will actually see. ``hard=True`` uses the straight-through estimator: a
	true one-hot on the forward pass, softmax gradient on the backward pass — so the value B consumes is a real token
	while gradients still flow. Returns ``[..., V]`` to feed straight into ``soft_embed``.
	"""
	return F.gumbel_softmax(logits, tau=tau, hard=hard, dim=-1)

def js_divergence(p: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
	"""Per-layer Jensen–Shannon divergence (symmetric, bounded by ln 2) → ``[n_layers]``."""
	m = 0.5 * (p + q)
	return 0.5 * kl_divergence(p, m) + 0.5 * kl_divergence(q, m)

def kl_divergence(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
	"""Per-layer KL(p || q) between expert distributions ``[n_layers, n_experts]`` → ``[n_layers]``."""
	p = p + eps
	q = q + eps
	p = p / p.sum(-1, keepdim=True)
	q = q / q.sum(-1, keepdim=True)
	return (p * (p / q).log()).sum(-1)

def message_token_spans(tokenizer: "PreTrainedTokenizerBase", view: list[dict]) -> list[tuple[int, int]]:
	"""Token span ``(start, end)`` of each message of ``view`` in the fully-rendered chat-template sequence.

	Method: render the *string* prefixes ``apply_chat_template(view[:i], tokenize=False)`` for each ``i`` and
	require each to be a string-prefix of the next (raised on non-prefix-stable templates). The full string is
	then tokenized **once** with ``return_offsets_mapping=True`` and each char boundary is mapped to the first
	token whose offset starts at/after it — prefix *strings* are never tokenized independently, because a
	tokenization of a prefix string is not in general a prefix of the full tokenization.

	Note this describes the final replayed view (the whole conversation templated once). Live generation builds
	the sequence incrementally, but for standard chat templates the rendered text is identical, so spans match.
	The returned spans cover each message's rendered chunk *including* its role header/footer markup.
	"""
	renders = [tokenizer.apply_chat_template(view[:i], tokenize=False) for i in range(1, len(view) + 1)]
	for a, b in zip(renders, renders[1:]):
		if not b.startswith(a):
			raise ValueError("chat template is not prefix-stable; cannot compute message spans by prefix diffing")
	full = renders[-1]
	enc = tokenizer(full, return_offsets_mapping=True, add_special_tokens=False)
	starts = [off[0] for off in enc["offset_mapping"]]
	n = len(starts)

	def tok_at(char_pos: int) -> int:
		for ti in range(n):
			if starts[ti] >= char_pos:
				return ti
		return n

	bounds = [0] + [len(r) for r in renders]
	return [(tok_at(bounds[i]), tok_at(bounds[i + 1])) for i in range(len(view))]

def moe_layer_indices(model: "PreTrainedModel") -> tuple[int, ...]:
	"""Decoder-layer indices that carry a sparse MoE block.

	HF MoE models return ``router_logits`` only for the *sparse* layers, in layer order, with no index
	attached. This maps that tuple back to real decoder-layer indices by checking each layer's ``mlp`` for a
	router ``gate`` submodule — which handles mixed stacks (e.g. Qwen-MoE ``decoder_sparse_step`` /
	``mlp_only_layers`` leaving some layers dense) as well as fully-sparse stacks like OLMoE.
	"""
	idx = []
	for i, layer in enumerate(decoder_layers(model)):
		mlp = getattr(layer, "mlp", None)
		if mlp is not None and hasattr(mlp, "gate") and hasattr(mlp, "experts"):
			idx.append(i)
	if not idx:
		raise ValueError(f"{type(model).__name__} has no sparse MoE layers (no mlp.gate/mlp.experts found)")
	return tuple(idx)

def moe_num_experts(model: "PreTrainedModel") -> int:
	"""Number of routed experts per MoE layer (``config.num_experts`` — same field name in OLMoE/Qwen-MoE)."""
	return int(model.config.num_experts)

def moe_topk(model: "PreTrainedModel") -> int:
	"""Experts selected per token (``config.num_experts_per_tok``)."""
	return int(model.config.num_experts_per_tok)

def routing_stats(captures: list[RoutingCapture], n_experts: int,
                  spans: list[tuple[int, int]] | None = None) -> RoutingStats:
	"""Pool per-token routing into per-layer expert-usage distributions.

	Parameters:
		captures: output of ``capture_router_logits`` (all layers share one token sequence).
		n_experts: total routed experts (``moe_num_experts(model)``) — needed to size the histogram since a
			span may never touch some experts.
		spans: optional ``[(start, end), ...]`` token windows to restrict to (e.g. only the MoE's own generated
			messages, from ``message_token_spans``). Default: all positions.
	"""
	if not captures:
		raise ValueError("no captures given")
	seq = captures[0].topk_experts.shape[0]
	mask = torch.zeros(seq, dtype=torch.bool)
	if spans is None:
		mask[:] = True
	else:
		for s, e in spans:
			mask[s:e] = True
	n_tok = int(mask.sum())
	if n_tok == 0:
		raise ValueError(f"spans select zero tokens (seq={seq}, spans={spans})")

	k = captures[0].topk_experts.shape[1]
	load = torch.zeros(len(captures), n_experts)
	mass = torch.zeros(len(captures), n_experts) if captures[0].router_logits is not None else None
	for li, cap in enumerate(captures):
		ids = cap.topk_experts[mask].long().reshape(-1)                       # [n_tok * k]
		load[li] = torch.bincount(ids, minlength=n_experts).float() / ids.numel()
		if mass is not None:
			mass[li] = torch.softmax(cap.router_logits[mask], dim=-1).mean(0)
	return RoutingStats(expert_load=load, expert_mass=mass, layers=tuple(c.layer for c in captures),
	                    n_tokens=n_tok, top_k=k)

def soft_embed(model: "PreTrainedModel", probs: torch.Tensor) -> torch.Tensor:
	"""Mix ``model``'s input-embedding rows by a per-position distribution: ``[..., V] @ E[V, d] -> [..., d]``.

	``probs`` is a (relaxed or one-hot) distribution over ``model``'s vocab for each position — e.g. a softmax /
	``gumbel_softmax_tokens`` output of an upstream model. Returns grad-connected soft embeddings suitable to pass as
	``inputs_embeds`` to ``forward_with_grad`` / ``continuation_logprob``. This is exact for hard one-hots: feeding
	``one_hot(ids)`` reproduces ``model.get_input_embeddings()(ids)``.
	"""
	weight = model.get_input_embeddings().weight  # [V, d]
	return probs.to(dtype=weight.dtype, device=weight.device) @ weight

def token_logprobs(scores: tuple[torch.Tensor, ...], generated_ids: torch.Tensor) -> dict:
	"""Compute per-token logprobs / surprisal / entropy for a generation.

	``scores`` is the tuple of per-step logit tensors from ``model.generate(..., output_scores=True,
	return_dict_in_generate=True)`` (one ``[vocab]`` per generated token); ``generated_ids`` are the sampled
	token ids. Returns lists suitable to drop into ``Message.metadata`` — scalar-per-token, so they stay small
	and don't violate the "no heavy tensors in metadata" invariant.

	Surprisal = -logprob (nats); entropy is the full next-token distribution entropy at each step (a readout of
	the model's uncertainty, distinct from the surprisal of the token it actually emitted).
	"""
	if len(scores) == 0:
		return {"logprobs": [], "surprisal": [], "entropy": []}
	# Stack once -> [steps, vocab]; batch all ops so there is a single GPU->CPU sync at the end.
	logits = torch.stack([s[0] for s in scores]).float()
	logp = torch.log_softmax(logits, dim=-1)
	toks = torch.as_tensor(generated_ids, device=logp.device, dtype=torch.long)[: logp.shape[0]]
	lp = logp.gather(1, toks.unsqueeze(1)).squeeze(1)
	ent = -(logp.exp() * logp).sum(dim=-1)
	logprobs = lp.tolist()
	return {
		"logprobs": logprobs,
		"surprisal": (-lp).tolist(),
		"entropy": ent.tolist(),
	}

def topk_expert_overlap(p: torch.Tensor, q: torch.Tensor, k: int = 8) -> torch.Tensor:
	"""Per-layer fraction of overlap between the ``k`` most-used experts of ``p`` and of ``q`` → ``[n_layers]``."""
	pi = p.topk(k, dim=-1).indices
	qi = q.topk(k, dim=-1).indices
	out = torch.zeros(p.shape[0])
	for li in range(p.shape[0]):
		out[li] = len(set(pi[li].tolist()) & set(qi[li].tolist())) / k
	return out