Skip to content
Murano

Steps — Logits

Logits step: expose the model’s output logits and next-token LM targets.

Runs a single forward pass over the batched prompts and stores the model’s raw [B, S, V] output logits, plus (by default) the left-shifted next-token target IDs the cross-entropy and accuracy metric steps consume. The shift and padding mask live only in the targets, never in the logits, so final_logits stays the honest model output for downstream analyses (logit-diff, DLA) while the generic LM metrics stay correct without slicing.

Logits

Section titled “Logits”

Run the model and write its output logits (plus next-token targets).

Runs one forward pass over the batched prompts via the backend’s forward_logits primitive and stores the raw [B, S, V] logits. When targets is set (the default), it also writes left-shifted next-token target IDs so CrossEntropyLossStep/AccuracyStep can run directly.

With fn given, the pass is intervened: fn rewrites each hooked module’s activation before the logits are read. This is the forward-pass analogue of :class:~murano.steps.intervene.Intervene, which applies the same edit during generation. Reach for it to measure an intervention with a metric rather than read it off a completion: a twelve-layer steering sweep costs twelve forward passes instead of twelve decodes.

Reads from results:

  • results[prompts_key]: PromptBatch (default prompts). Set this to run
  • the step on a different prompt set, e.g. the corrupt side of a
  • paired run.

Writes to results:

  • results['final_logits']: Tensor [B, S, vocab] of output logits.
  • results['attention_mask']: Tensor [B, S] marking real (1) vs padding
  • (0) positions, so downstream metric steps can locate the answer
  • position without re-tokenizing.
  • results['target_ids']: Tensor [B, S] of next-token targets (only when
  • targets is set).

Args:

  • model: Model backend to run.
  • prompts_key: Results key to read the prompts from. When a second Logits
  • runs in the same pipeline (e.g. the corrupt side of a paired run),
  • also override logits_key and mask_key (and set targets accordingly)
  • so it does not overwrite the first run’s outputs.
  • logits_key: Results key to write the logits under.
  • targets_key: Results key to write the next-token targets under.
  • mask_key: Results key to write the attention mask under.
  • targets: Target-generation mode. "next_token" writes left-shifted
  • next-token targets; None writes logits only, for callers that
  • supply their own task targets (e.g. logit-diff).
  • fn: (activation, node) -> activation, applied to each hooked module’s
  • activation before the logits are read; for example
  • steer_direction(direction, alpha=4) or ablate_direction(...).
  • None (the default) runs the plain forward pass, in which case
  • layers, modules, and per_head are ignored.
  • layers: Layer indices to hook, or "all".
  • modules: Module name(s) to hook at each layer.
  • per_head: Hook the attention output projection’s input and hand fn the
  • per-head activation, so it can rewrite one head’s slice.

Raises:

  • ValueError: If targets is neither "next_token" nor None.

__init__

Section titled “__init__”
def __init__(self, model: ModelBackend, prompts_key: str = keys.PROMPTS, logits_key: str = keys.FINAL_LOGITS, targets_key: str = keys.TARGET_IDS, mask_key: str = keys.ATTENTION_MASK, targets: str | None = 'next_token', fn: Callable[[Tensor, Node], Tensor] | None = None, layers: list[int] | str = 'all', modules: str | list[str] = 'residual', per_head: bool = False):