Steps — Metrics

Metric computation steps for Murano pipelines.

Converts the legacy BaseComputationLens classes into Step subclasses that operate on Results objects.

CrossEntropyLossStep

Compute cross-entropy loss from logits and target IDs.

Reads logits_key and targets_key from results, computes the loss, and writes the result under output_key.

Args:

• logits_key: Key in results containing the logits tensor [B, S, V].
• targets_key: Key in results containing the target IDs [B, S].
• output_key: Key under which to store the computed loss.
• reduction: Loss reduction: "mean", "sum", or "none".

__init__

def __init__(self, logits_key: str = keys.FINAL_LOGITS, targets_key: str = keys.TARGET_IDS, output_key: str = keys.LOSS, reduction: Literal['mean', 'sum', 'none'] = 'mean'):

AccuracyStep

Compute token-level accuracy from logits and target IDs.

Reads logits_key and targets_key from results, computes the fraction of argmax predictions that match the target, and writes the result (a plain Python float) under output_key.

Args:

• logits_key: Key in results containing the logits tensor [B, S, V].
• targets_key: Key in results containing the target IDs [B, S].
• output_key: Key under which to store the computed accuracy.

__init__

def __init__(self, logits_key: str = keys.FINAL_LOGITS, targets_key: str = keys.TARGET_IDS, output_key: str = keys.ACCURACY):

ComparisonComputationStep

Compute element-wise difference or row-wise cosine similarity.

Reads two tensors from results, computes the comparison, and writes the result under output_key.

Args:

• key_a: Key in results for the first tensor.
• key_b: Key in results for the second tensor.
• output_key: Key under which to store the result.
• comparison_type: "difference" or "cosine_similarity".

__init__

def __init__(self, key_a: str, key_b: str, output_key: str, comparison_type: str = 'difference'):

LogitDiffStep

Logit difference between a correct and an incorrect answer token.

At the answer position of each example, subtracts the incorrect token’s logit from the correct token’s logit; a positive value means the model favors the correct answer. This is the standard comparable metric for activation-patching and circuit experiments.

Reads from results:

• results[logits_key]: Tensor [B, S, V] of output logits.
• results[mask_key]: optional Tensor [B, S], used to locate the answer
• position when positions is not given.

Writes to results:

• results[output_key]: MetricScore.

Args:

• correct: Correct-answer token spec: a token id, an answer string, a
• per-example sequence of either, or a tensor ([B] or [B, k]
• token set).
• incorrect: Incorrect-answer token spec, in the same forms as
• correct.
• logits_key: Results key holding the logits.
• mask_key: Results key holding the attention mask.
• output_key: Results key to write the MetricScore under.
• positions: Optional explicit answer position(s); overrides the mask.
• model: Model backend, required only to tokenize string answers.

__init__

def __init__(self, correct: int | str | Sequence[int] | Sequence[str] | torch.Tensor, incorrect: int | str | Sequence[int] | Sequence[str] | torch.Tensor, logits_key: str = keys.FINAL_LOGITS, mask_key: str = keys.ATTENTION_MASK, output_key: str = keys.LOGIT_DIFF, positions: int | Sequence[int] | torch.Tensor | None = None, model: ModelBackend | None = None):

KLDivergenceStep

KL divergence between two logit distributions at the answer position.

Computes KL(softmax(P) || softmax(Q)) per example at the answer position and reduces by mean. Useful for measuring how far a corrupted or patched run’s next-token distribution moves from a clean reference.

Reads from results:

• results[p_key], results[q_key]: Tensors [B, S, V].
• results[mask_key]: optional Tensor [B, S] for the answer position.

Writes to results:

• results[output_key]: MetricScore.

Note:

• The answer position is taken from P; P and Q are assumed to come from
• the same prompts (same tokenization), as in a clean-vs-patched run.

Args:

• p_key: Results key for the reference distribution P.
• q_key: Results key for the comparison distribution Q.
• mask_key: Results key holding the attention mask.
• output_key: Results key to write the MetricScore under.
• positions: Optional explicit answer position(s); overrides the mask.

__init__

def __init__(self, p_key: str, q_key: str, mask_key: str = keys.ATTENTION_MASK, output_key: str = keys.KL_DIV, positions: int | Sequence[int] | torch.Tensor | None = None):

AnswerLogProbStep

Log-probability of the correct answer token at the answer position.

Softmaxes the logits at the answer position and reports the mean log-prob of the correct token. With as_loss=True it reports the mean negative log-prob instead (the per-example answer cross-entropy), the base quantity a loss-recovered metric normalizes.

Reads from results:

• results[logits_key]: Tensor [B, S, V].
• results[mask_key]: optional Tensor [B, S] for the answer position.

Writes to results:

• results[output_key]: MetricScore.

Args:

• correct: Correct-answer token spec (id, string, per-example sequence,
• or tensor), as in: class:LogitDiffStep.
• logits_key: Results key holding the logits.
• mask_key: Results key holding the attention mask.
• output_key: Results key to write the MetricScore under.
• positions: Optional explicit answer position(s); overrides the mask.
• as_loss: If True, report mean negative log-prob instead of log-prob.
• model: Model backend, required only to tokenize string answers.

__init__

def __init__(self, correct: int | str | Sequence[int] | Sequence[str] | torch.Tensor, logits_key: str = keys.FINAL_LOGITS, mask_key: str = keys.ATTENTION_MASK, output_key: str = keys.ANSWER_LOGPROB, positions: int | Sequence[int] | torch.Tensor | None = None, as_loss: bool = False, model: ModelBackend | None = None):

AnswerRankStep

Rank of the correct answer token among the whole vocabulary.

Counts how many tokens the model scores strictly above the correct answer at the answer position. Rank 0 means the model would emit the answer greedily. Unlike :class:LogitDiffStep, the rank ignores how far ahead the answer is, which makes it a blunt but very readable check: it bottoms out at 0 long before the logit difference stops moving.

Ties are resolved optimistically (a tied token does not count against the answer), so the reported rank is the best position the answer could occupy. When the answer spec names a token set, each token’s rank is computed and the mean reported, matching how :class:AnswerLogProbStep averages a set.

Reads from results:

• results[logits_key]: Tensor [B, S, V].
• results[mask_key]: optional Tensor [B, S] for the answer position.

Writes to results:

• results[output_key]: MetricScore.

Args:

• correct: Correct-answer token spec (id, string, per-example sequence, or
• tensor), as in: class:LogitDiffStep.
• logits_key: Results key holding the logits.
• mask_key: Results key holding the attention mask.
• output_key: Results key to write the MetricScore under.
• positions: Optional explicit answer position(s); overrides the mask.
• model: Model backend, required only to tokenize string answers.

__init__

def __init__(self, correct: int | str | Sequence[int] | Sequence[str] | torch.Tensor, logits_key: str = keys.FINAL_LOGITS, mask_key: str = keys.ATTENTION_MASK, output_key: str = keys.ANSWER_RANK, positions: int | Sequence[int] | torch.Tensor | None = None, model: ModelBackend | None = None):

RecoveredMetricStep

Normalized recovered effect from clean, corrupted, and patched scores.

Computes (patched - corrupted) / (clean - corrupted): 0 means the patched run matches the corrupted baseline, 1 means it matches the clean run. The three inputs are scalars produced by any base metric (for example :class:LogitDiffStep) under different conditions. Producing the corrupted and patched logits is the activation-patching step’s job and is out of scope here; this only combines already-computed scores.

Reads from results:

• results[clean_key], results[corrupted_key], results[patched_key]:
• a MetricScore, a 0-dim tensor, or a plain float.

Writes to results:

• results[output_key]: MetricScore.

Args:

• clean_key: Results key for the clean-run score.
• corrupted_key: Results key for the corrupted-run score.
• patched_key: Results key for the patched-run score.
• output_key: Results key to write the MetricScore under.

__init__

def __init__(self, clean_key: str, corrupted_key: str, patched_key: str, output_key: str = keys.RECOVERED):