Fix/audit corrections

CHANGELOG.md

@@ -81,11 +81,11 @@ Initial release.
 
 ### Added
 
-- Quick API on `MuranoModel`: `find_direction()`, `generate(intervention=...)`, direct activation recording.
+- Quick API on `MuranoModel`: `find_direction()`, `generate(ablate=...)` / `generate(steer=...)`, direct activation recording.
 - Pipeline API: composable `Step` + `Pipeline` with pre-flight validation of `reads`/`writes` contracts.
 - Steps: `Load`, `Record`, `SteeringVector`, `Intervene`, `Probe`, `GenerationMetric`.
 - Logit lens (`LogitLens` step).
-- Datasets: `MuranoDataset` (contrastive) and `LabeledDataset`, with `from_hub()` and `from_template()` factories.
+- Datasets: `MuranoDataset` (contrastive) and `LabeledDataset`, with `contrastive()` / `from_hub()` factories.
 - Direction-based interventions: `ablate_direction`, `steer_direction`.
 - I/O: `save_results()` with structured output layout, `load_steering()`, `save_ablated_model()`.
 - Top-level `__version__` via `importlib.metadata`.

CONTRIBUTING.md

@@ -103,8 +103,8 @@ No bespoke try/except guards: every optional import goes through
 
 ## Code of Conduct
 
-Please note that this project is released with a [Code of Conduct](CODE_OF_CONDUCT.md). By participating in this project you agree to abide by its terms.
+ToDo: adopt and link a Code of Conduct. We ask contributors to be respectful and constructive in the meantime.
 
 ## License
 
-By contributing, you agree that your contributions will be licensed under the MIT License.
+By contributing, you agree that your contributions will be licensed under the Apache License 2.0, the license this project is released under.

README.md

@@ -28,9 +28,10 @@ libraries ship as extras, so you install per use case:
 | (base) | recording, steering, intervention, logits, ablation, metrics, paired datasets | nnsight, nnterp, torch, transformers |
 | `probe` | linear probing | scikit-learn |
 | `data` | loading datasets by name from the Hub | datasets |
-| `plot` | figures and visualizations | matplotlib, seaborn, plotly |
+| `plot` | figures and visualizations | plotly, kaleido |
 | `sae` | sparse autoencoder features | sae-lens |
-| `all` | everything above | all of the above |
+| `notebook` | running the example notebooks | jupyter, ipykernel, nltk |
+| `all` | everything above | probe, data, plot, sae, notebook |
 
 ```bash
 pip install "murano-interp[probe,plot]"   # combine as needed
@@ -133,6 +134,13 @@ caught up-front.
 | `WeightAblation`   | `prompts`, `steering` | `intervene`, `weight_ablation` | Project a direction out of model weights, then generate. |
 | `Logits` ‡         | `prompts`             | `final_logits`, `attention_mask`, `target_ids` | Run a forward pass and expose output logits plus next-token targets. With `fn=`, apply an intervention during that pass. |
 | `Ablate` ‡         | `prompts`             | `ablated_logits`, `attention_mask` | Zero, mean, or resample a component and return the logits. |
+| `Patch` ‡          | `corrupt_prompts`, `prompts` | `patched_logits`, `patched_mask` | Interchange (activation) patching: inject source-run activations into the base run. |
+| `PathPatch` ‡      | `prompts`, `corrupt_prompts` | `path_patched_logits`, `path_patched_mask` | Isolate a sender-to-receiver path's effect. |
+| `RecordAttention` ‡ | `prompts`            | `attention_pattern`            | Capture per-head attention weights (needs `enable_attention_probs`). |
+| `AblateAttention` ‡ | `prompts`            | `attn_ablated_logits`, `attn_ablated_mask` | Overwrite per-head attention weights and return the logits. |
+| `LogitAttribution` ‡ | `prompts`           | `logit_attribution`            | Decompose a logit (difference) into per-component contributions (DLA). |
+| `SelectComponents` ‡ | `logit_attribution` (or `sweep`) | `selection`          | Rank per-component scores and keep the strongest as a target set. |
+| `LogitLens`        | `prompts`             | `logit_lens`                   | Project each layer's residual onto the vocabulary. |
 | `Sweep` ‡          | (the swept chain's)   | `sweep`                        | Run a step chain once per item and harvest a metric into a `SweepResult`. |
 | `Probe` §          | `record`              | `probe`                        | Train a linear probe per layer via cross-validation.     |
 | `GenerationMetric` | `intervene`           | `metric`                       | Score baseline vs modified outputs with a user metric.   |
@@ -172,15 +180,18 @@ Pipeline API like the other steps.
 
 ```text
 src/murano/
-  model.py
-  pipeline.py
-  results.py
-  artifacts.py
-  dataset.py
-  io.py
-  evaluation.py
-  steps/
-  plotting/
+  model.py       # MuranoModel: the nnterp-backed model wrapper
+  backend.py     # ModelBackend: the surface steps depend on
+  pipeline.py    # Pipeline / Step orchestration
+  results.py     # Results: the shared state passed between steps
+  keys.py        # canonical Results keys
+  nodes.py       # Node / Edge / NodeSet: component addressing
+  artifacts.py   # PromptBatch, SteeringResult, MetricScore, ...
+  dataset.py     # MuranoDataset, LabeledDataset, CleanCorruptDataset
+  tasks.py       # small canonical tasks (ioi, sentiment)
+  io.py          # save/load of results
+  steps/         # pipeline steps
+  plotting/      # Plotly visualizations
 ```
 
 ## Notebooks
@@ -195,8 +206,11 @@ then pick the application you need:
 - [`notebooks/reproductions/`](notebooks/reproductions/) — published results
   reproduced with Murano.
 
-Every notebook is executed before it is committed, and all of them render on the
+Notebooks are executed before they are committed so their stored outputs are
+current, and all of them render on the
 [documentation site](https://ukplab.github.io/murano/docs/notebooks/getting_started/).
+(This is a maintainer pre-commit step; CI checks that the notebooks import and
+parse, but does not re-execute them, since several need a GPU.)
 
 ## Development
 

docs/src/content/docs/docs/getting-started/installation.md

@@ -24,9 +24,10 @@ libraries ship as extras, so you install for what you actually do:
 | (base) | recording, steering, intervention, logits, ablation, metrics, paired datasets | nnsight, nnterp, torch, transformers |
 | `probe` | linear probing | scikit-learn |
 | `data` | loading datasets by name from the Hub | datasets |
-| `plot` | figures and visualizations | matplotlib, seaborn, plotly |
+| `plot` | figures and visualizations | plotly, kaleido |
 | `sae` | sparse autoencoder features | sae-lens |
-| `all` | everything above | all of the above |
+| `notebook` | running the example notebooks | jupyter, ipykernel, nltk |
+| `all` | everything above | probe, data, plot, sae, notebook |
 
 Combine extras as needed:
 

docs/src/content/docs/docs/getting-started/quickstart.mdx

@@ -40,8 +40,8 @@ direction = model.find_direction(
     negative=["What a miserable, dreadful day", "This is awful and depressing"],
     layers=[10, 11, 12, 13],
 )
-print(direction.best_layer)         # e.g. 12
-print(direction.separation_scores)  # {10: 1.3, 11: 2.1, 12: 3.4, 13: 2.8}
+print(direction.best_layer)         # a Node, e.g. L12.resid_post
+print(direction.separation_scores)  # {L10.resid_post: 1.3, L11.resid_post: 2.1, ...}
 ```
 
 ## Steer generation

docs/src/content/docs/docs/guides/patch.md

@@ -67,13 +67,13 @@ results = Pipeline([
 
 ### Custom intervention functions
 
-`Intervene` accepts any callable with signature `(activation: Tensor, layer_idx: int) -> Tensor`:
+`Intervene` accepts any callable with signature `(activation: Tensor, node: Node) -> Tensor`. The second argument is a `murano.Node` naming the hooked site (`node.layer`, `node.module`, ...), not a bare layer index, so a function that targets specific sites branches on it:
 
 ```python
 import torch
 
-def zero_out(activation: torch.Tensor, layer: int) -> torch.Tensor:
-    """Zero all activations — a sanity check."""
+def zero_out(activation: torch.Tensor, node) -> torch.Tensor:
+    """Zero all activations — a sanity check. `node` names the hooked site."""
     return torch.zeros_like(activation)
 
 results = Pipeline([
@@ -84,7 +84,7 @@ results = Pipeline([
 
 ## Weight-level ablation
 
-The `WeightAblation` step applies an orthogonal projection P = I - dd^T to the model's weight matrices directly. This removes the direction from **all** computations, not just the residual stream at hook points.
+The `WeightAblation` step applies an orthogonal projection P = I - dd^T to the model's weight matrices directly (write matrices as `P @ W`, read matrices as `W @ P`), so the ablation is baked into the weights rather than applied at hook points during the forward pass. It is a stronger, more global edit than activation-level ablation; note that the read-matrix projection interacts with the intervening layer norm, so it does not remove the direction perfectly.
 
 ```python
 from murano.steps import WeightAblation

docs/src/pages/reproductions/ioi.mdx

@@ -63,7 +63,11 @@ from murano.steps.metrics import LogitDiffStep
 from murano.steps.paired import LoadPaired
 from murano.results import Results
 
-model = MuranoModel("gpt2", enable_attention_probs=True)
+import torch
+
+# gpt2 is small, so float32 is cheap and keeps the logit differences clean
+# (bfloat16 measurably reorders close logits at this scale).
+model = MuranoModel("gpt2", dtype=torch.float32, enable_attention_probs=True)
 ds = CleanCorruptDataset(clean=clean, corrupt=corrupt, correct=io_ids, incorrect=s_ids)
 
 base = LoadPaired(ds)(Results())

docs/src/pages/reproductions/tmlr2026.mdx

@@ -161,6 +161,14 @@ plt.tight_layout(); plt.show()
 
 ## Key results
 
+<p class="aside">
+  Preview: this reproduction is a work in progress. It depends on an external
+  <code>SupervisedMDS</code> package and a placeholder dataset (see above), so the
+  numbers below are the paper's reported figures, not yet independently
+  reproduced with Murano. The other gallery entries (IOI, Geometry of Truth,
+  Function Vectors) are full reproductions.
+</p>
+
 | Concept type      | Model        | Best-fit manifold | Probe accuracy |
 | ----------------- | ------------ | ----------------- | -------------- |
 | Month of year     | Llama-3.2-1B | Circle            | 0.91           |

src/murano/dataset.py

@@ -162,6 +162,16 @@ def from_hub(
         pos_texts = _load_hub_column(positive, n=n_total)
         neg_texts = _load_hub_column(negative, n=n_total)
 
+        # A source with fewer than n_train + n_eval rows would otherwise yield a
+        # short (or empty) eval split silently; fail loudly instead.
+        for label, texts in (("positive", pos_texts), ("negative", neg_texts)):
+            if len(texts) < n_total:
+                raise ValueError(
+                    f"{label} source yielded {len(texts)} examples, but n_train + "
+                    f"n_eval = {n_total} are needed; lower n_train/n_eval or use a "
+                    f"larger split."
+                )
+
         raw_pos = list(pos_texts)
         raw_neg = list(neg_texts)
 
@@ -298,7 +308,14 @@ def from_hub(
             )
         """
         if isinstance(source, tuple):
-            name, config = source if len(source) == 2 else (source[0], None)
+            if len(source) == 2:
+                name, config = source
+            elif len(source) == 1:
+                name, config = source[0], None
+            else:
+                raise ValueError(
+                    f"source tuple must be (name,) or (name, config), got {source!r}"
+                )
         else:
             name, config = source, None
 

src/murano/io.py

@@ -1,4 +1,12 @@
-"""I/O utilities: saving and loading Murano results."""
+"""I/O utilities: saving and loading Murano results.
+
+Security note: the ``.pt`` loaders (``load_steering``, ``load_logit_lens``,
+``load_attention``, ``load_activation_store``, and the labeled variant) call
+``torch.load(..., weights_only=False)`` because the payloads hold custom objects
+(``Node`` keys, dataclasses) that the safe loader cannot reconstruct. That path
+executes arbitrary pickle, so only load artifacts you produced yourself or
+otherwise trust; do not load ``.pt`` files from an untrusted source.
+"""
 
 from __future__ import annotations
 
@@ -408,14 +416,28 @@ def _restore(value: Any) -> Any:
         # null is how a non-finite scalar was stored; restore it as nan.
         return float("nan") if value is None else value
 
+    # The per-component maps can carry non-finite values too (a head whose frozen
+    # projection produced nan), stored as null; restore them the same way as the
+    # scalar fields so the round-trip is symmetric.
+    contributions = {node: _restore(v) for node, v in data["contributions"].items()}
+    raw_per_example = data.get("per_example")
+    per_example = (
+        None
+        if raw_per_example is None
+        else {
+            node: [_restore(x) for x in values]
+            for node, values in raw_per_example.items()
+        }
+    )
+
     return LogitAttributionResult(
-        contributions=data["contributions"],
+        contributions=contributions,
         embed_contribution=_restore(data["embed_contribution"]),
         other_contribution=_restore(data["other_contribution"]),
         target=data["target"],
         total=_restore(data["total"]),
         completeness_error=_restore(data["completeness_error"]),
-        per_example=data.get("per_example"),
+        per_example=per_example,
         metadata=data.get("metadata", {}),
     )
 
@@ -951,10 +973,12 @@ def serialize_logit_lens(
     ) -> None:
         filename = "logit_lens.pt" if key == keys.LOGIT_LENS else f"{key}.pt"
         save_logit_lens(logit_lens, out / "logit_lens" / filename)
+        # Read shapes off max_probs, not all_probs: the full-vocab tensor is
+        # None unless the step ran with store_full_probs=True.
         metadata[key] = {
             "addresses": [str(a) for a in logit_lens.addresses],
-            "n_layers": logit_lens.all_probs.shape[0],
-            "n_inputs": logit_lens.all_probs.shape[1],
+            "n_layers": logit_lens.max_probs.shape[0],
+            "n_inputs": logit_lens.max_probs.shape[1],
         }
 
     def serialize_attention(