feat: add the Sweep step and rebuild the demo surface as notebooks

.github/workflows/docs.yml

@@ -35,6 +35,9 @@ jobs:
       - name: Generate API reference
         run: uv run python docs/scripts/gen_api_docs.py
 
+      - name: Generate notebook pages
+        run: uv run python docs/scripts/gen_notebook_docs.py
+
       - uses: actions/setup-node@v6
         with:
           node-version: "22"

.gitignore

@@ -140,7 +140,9 @@ sketch*
 murano/sketch*
 .vscode*
 .draft/
-examples/artifacts
+
+# Default destination of the Save and Plot steps, written by the notebooks.
+murano_outputs/
 
 # OS files
 .DS_Store

CHANGELOG.md

@@ -11,6 +11,31 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `SelectComponents` step and `ComponentSelection` artifact: rank an attribution result (for example `LogitAttribution`) by magnitude, signed value, or most-negative, keep the top `top_k` or everything past a `threshold`, and write the chosen addresses for a downstream step to read. `Patch` / `PathPatch` / `Ablate` accept a `targets_key` / `senders_key` naming that selection, so attribute-then-patch runs as one pipeline instead of two with a hand-copied node list.
 - `Intervene` gains `direction_layers` (`"all"`, `"best"`, or an explicit layer list) for the `direction_key` steering path, so a one-pipeline steer can apply only the best-separating layer's direction instead of every recorded layer, which keeps deep models coherent.
+- `Sweep` step and `SweepResult` artifact: run a step chain once per item and harvest one or more metric keys. Every component study has this shape ("patch each head and measure what it restores", "zero each head and measure the damage", "steer at each layer"), and every notebook was hand-rolling it as a closure over the model, the task, and a baseline `Results`. `Sweep` forks the incoming `Results` per item, so the shared prefix runs once and the swept steps' writes stay out of the pipeline, and it derives its own read contract from the chain, so a missing upstream key fails pre-flight validation. A sweep over `Node` addresses publishes the same `{Node: float}` map an attribution does, so it feeds `SelectComponents` and `plot_head_matrix` with no adapter: attribute, sweep, select and path-patch now compose in one pipeline.
+- `plot_sweep` in `murano.plotting`, auto-dispatched by `Plot`: a sweep over attention heads renders as the layer-by-head heatmap, anything else as one bar per item. The color scale diverges only when the values actually straddle zero.
+- `AnswerRankStep`: how many tokens outscore the correct answer at the answer position. Rank 0 means the model would emit it greedily. Like its sibling metric steps it resolves the answer position from the attention mask, so it stays correct under either padding side.
+- `Logits` accepts `fn`, `layers`, `modules` and `per_head`, forwarding them to the backend's `forward_logits`, which already took them. It is the forward-pass analogue of `Intervene`: a twelve-layer steering sweep now costs twelve forward passes instead of twelve decodes.
+- `SAEModel.decoder`: the whole decoder matrix, for callers projecting every feature through the unembedding. Previously reachable only as `sae_model._sae.W_dec`.
+- `murano.tasks`: the two toy tasks the tutorials share, defined once and tested. `ioi()` builds the indirect-object-identification task as a `CleanCorruptDataset`; `sentiment()` returns contrastive sentences; `positive_word_rate()` is the crude scorer the steering notebooks use.
+- `plot_activation_projection` in `murano.plotting`: reduce one component's activations with any scikit-learn-style reducer (PCA, LDA, t-SNE, UMAP) and scatter them by class. It accepts both a contrastive `ActivationStore` and a `LabeledActivationStore`, so a single `Record` can feed both `Probe` and the plot.
+- `zmid` on `plot_heatmap` and `plot_head_matrix`, anchoring a diverging colorscale at zero so a signed statistic no longer shades zero as if it had a sign.
+- `notebooks/getting_started.ipynb`, plus fourteen runnable notebooks under `notebooks/applications/`: steering, probing, logit lens, logit attribution, attention, ablation, activation patching, circuit discovery, metrics, custom pipeline, weight ablation, and the three sparse-autoencoder notebooks. All fifteen share one template, enforced by `tests/test_notebook_structure.py`, which also rejects a step constructed inside a loop (that is a hand-rolled `Sweep`) and a pipeline built inside a function.
+- The notebooks now render on the documentation site, generated from the executed `.ipynb` files by `docs/scripts/gen_notebook_docs.py` at deploy time.
+
+### Changed
+
+- The SAE notebooks move from `notebooks/sae/` into `notebooks/applications/` as `sae_features`, `sae_steering` and `sae_enrichment`, on the shared template. `sae_enrichment` replaces a 150-line bespoke ranking heuristic with the standardized mean difference between classes, which is four lines and surfaces the same sentiment features, along with the negation and junk features that ranking alone cannot tell apart.
+
+### Fixed
+
+- `Intervene`'s docstring recommended `direction_layers="best"` while the code defaulted to `"all"`, and never stated the default. It now documents the default and explains the trade-off: the best-*separating* layer is not necessarily an effective place to intervene, because a concept is often most separable in an early layer whose contribution later layers overwrite.
+- `sae_steer`'s docstring quoted `alpha=200` as if it were a usable default. `alpha` is an absolute magnitude added at every decoded token, so it means nothing without the residual norm at the steering site, which spans tens to thousands across models. The docstring now says so, and points at `notebooks/applications/sae_steering.ipynb`, which measures that norm and finds no strength that both invokes the concept and preserves the text. The notebook's old `alpha=2000` output predates the fix that made generation-time interventions apply on every decoded token.
+- `tests/test_notebook_structure.py` compared function *names* across notebooks, so `patch_head` and `recovered_fraction`, the same seventeen lines twice, lived in two notebooks unnoticed. It now compares normalized syntax trees.
+
+### Removed
+
+- The `examples/` directory. Its scripts became notebooks under `notebooks/`, and the one reusable class it held (`PlotterLens`) moved into the library as `murano.plotting.plot_activation_projection`.
+- `.coveragerc`, which duplicated the `[tool.coverage]` configuration in `pyproject.toml` and took precedence over it.
 
 ## [0.1.0a2] - 2026-07-08
 

README.md

@@ -131,11 +131,12 @@ caught up-front.
 | `SteeringVector`   | `record`              | `steering`                     | Find a contrastive steering direction (mean diff).       |
 | `Intervene`        | `prompts`             | `intervene`                    | Generate baseline + intervened outputs side-by-side.     |
 | `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. |
+| `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. |
+| `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.   |
-| Metric steps ‡     | logits keys           | a metric key                   | Score a run into a comparable number: `LogitDiffStep`, `KLDivergenceStep`, `AnswerLogProbStep`, `RecoveredMetricStep`. |
+| Metric steps ‡     | logits keys           | a metric key                   | Score a run into a comparable number: `LogitDiffStep`, `KLDivergenceStep`, `AnswerLogProbStep`, `AnswerRankStep`, `RecoveredMetricStep`. |
 | `Save`             | (any present)         | `output_dir`                   | Persist all results to organized subdirectories.         |
 | `SAEEncode` †      | `prompts`             | `sae_record`                   | Encode residuals through an SAE loaded from HuggingFace. |
 | `SAETopActivations` | `sae_record`         | `feature_examples`             | Rank the top-K activating contexts per SAE feature.      |
@@ -182,10 +183,20 @@ src/murano/
   plotting/
 ```
 
-## Examples
+## Notebooks
 
-- `examples/quick_prototype.py`
-- `examples/sae_example.py`
+Start with [`notebooks/getting_started.ipynb`](notebooks/getting_started.ipynb),
+then pick the application you need:
+
+- [`notebooks/applications/`](notebooks/applications/) — one self-contained
+  notebook per application: steering, probing, logit lens, logit attribution,
+  attention, ablation, activation patching, circuit discovery, metrics, custom
+  pipelines, weight ablation, and sparse autoencoders.
+- [`notebooks/reproductions/`](notebooks/reproductions/) — published results
+  reproduced with Murano.
+
+Every notebook is executed before it is committed, and all of them render on the
+[documentation site](https://ukplab.github.io/murano/docs/notebooks/getting_started/).
 
 ## Development
 

docs/.gitignore

@@ -3,3 +3,5 @@ node_modules/
 .astro/
 src/content/docs/docs/reference/
 doc_report.md
+src/content/docs/docs/notebooks/
+public/notebook-figures/

docs/astro.config.mjs

@@ -55,6 +55,10 @@ export default defineConfig({
 						{ label: 'Pipelines', slug: 'docs/guides/pipeline' },
 					],
 				},
+				{
+					label: 'Notebooks',
+					autogenerate: { directory: 'docs/notebooks' },
+				},
 				{
 					label: 'Reproductions',
 					items: [

docs/scripts/gen_api_docs.py

@@ -21,6 +21,7 @@
     ("murano.model", "model"),
     ("murano.pipeline", "pipeline"),
     ("murano.dataset", "dataset"),
+    ("murano.tasks", "tasks"),
     ("murano.artifacts", "artifacts"),
     ("murano.results", "results"),
     ("murano.io", "io"),
@@ -45,6 +46,7 @@
     ("murano.steps.logit_attribution", "steps/logit_attribution"),
     ("murano.steps.sae", "steps/sae"),
     ("murano.steps.metrics", "steps/metrics"),
+    ("murano.plotting.activations", "plotting/activations"),
     ("murano.plotting.steering", "plotting/steering"),
     ("murano.plotting.probing", "plotting/probing"),
     ("murano.plotting.logit_lens", "plotting/logit_lens"),

docs/scripts/gen_notebook_docs.py

@@ -0,0 +1,206 @@
+"""Render the executed notebooks into Starlight pages.
+
+Run from the repo root:
+    python docs/scripts/gen_notebook_docs.py
+
+Outputs Markdown to docs/src/content/docs/docs/notebooks/ and figure images to
+docs/public/notebook-figures/. Both are generated, gitignored, and rebuilt on
+every docs deploy, exactly like the API reference.
+
+The notebooks are committed with their outputs, so nothing here runs a model.
+Figures are stored only as Plotly JSON, which no Markdown renderer understands,
+so each figure is rebuilt into a static PNG with kaleido (the same headless
+exporter the `plot` extra pins for Slurm nodes).
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import re
+import warnings
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).parent.parent.parent
+NOTEBOOKS = REPO_ROOT / "notebooks"
+OUT = REPO_ROOT / "docs" / "src" / "content" / "docs" / "docs" / "notebooks"
+FIGURES = REPO_ROOT / "docs" / "public" / "notebook-figures"
+FIGURE_URL = "/murano/notebook-figures"
+GITHUB = "https://github.com/UKPLab/murano/blob/main/notebooks"
+
+# Ordered so the sidebar reads as a curriculum rather than an alphabet.
+ORDER = [
+    "getting_started.ipynb",
+    "applications/steering.ipynb",
+    "applications/probing.ipynb",
+    "applications/logit_lens.ipynb",
+    "applications/logit_attribution.ipynb",
+    "applications/attention.ipynb",
+    "applications/ablation.ipynb",
+    "applications/activation_patching.ipynb",
+    "applications/circuit_discovery.ipynb",
+    "applications/metrics.ipynb",
+    "applications/custom_pipeline.ipynb",
+    "applications/weight_ablation.ipynb",
+    "applications/sae_features.ipynb",
+    "applications/sae_steering.ipynb",
+    "applications/sae_enrichment.ipynb",
+]
+
+_MAX_OUTPUT_LINES = 40
+_MAX_DESCRIPTION = 150
+
+
+def _title_and_description(cells: list[dict]) -> tuple[str, str]:
+    """Take the page title from the H1 and the description from the paragraph under it."""
+    header = "".join(cells[0]["source"])
+    title = header.splitlines()[0].lstrip("# ").strip()
+    body = header.split("\n", 1)[1].strip()
+    first_paragraph = body.split("\n\n", 1)[0].replace("\n", " ").strip()
+    description = re.sub(r"[*`]", "", first_paragraph)
+    if len(description) > _MAX_DESCRIPTION:
+        # Cut at a word boundary; a description sliced mid-word reads as a bug.
+        description = description[:_MAX_DESCRIPTION].rsplit(" ", 1)[0] + "..."
+    return title, description
+
+
+def _stream_text(cell: dict) -> str:
+    text = "".join(
+        "".join(out.get("text", ""))
+        for out in cell.get("outputs", [])
+        if out.get("output_type") == "stream" and out.get("name") == "stdout"
+    )
+    for out in cell.get("outputs", []):
+        if out.get("output_type") == "execute_result":
+            text += "".join(out.get("data", {}).get("text/plain", ""))
+    return text.rstrip()
+
+
+def _write_figures(cell: dict, slug: str, counter: list[int]) -> list[str]:
+    """Write each figure in a cell as a PNG; return the image paths.
+
+    A Plotly figure lives in the notebook only as its JSON spec, which no Markdown
+    renderer understands, so it is rebuilt through kaleido. A notebook that already
+    rasterized its figures carries them as base64 ``image/png`` instead, and those
+    are written straight out: dropping them would leave the page silently figureless.
+    """
+    import plotly.graph_objects as go
+    import plotly.io as pio
+
+    paths = []
+    for out in cell.get("outputs", []):
+        data = out.get("data", {})
+        spec = data.get("application/vnd.plotly.v1+json")
+        png = data.get("image/png")
+        if not spec and not png:
+            continue
+
+        counter[0] += 1
+        name = f"{slug}-{counter[0]}.png"
+        if spec:
+            # The mime bundle carries a "config" key that go.Figure does not accept,
+            # so rebuild from data and layout only.
+            figure = go.Figure(data=spec.get("data", []), layout=spec.get("layout", {}))
+            with warnings.catch_warnings():
+                # kaleido<1 ships a self-contained Chromium and is the only exporter
+                # that works headless; its deprecation notice is not actionable here.
+                warnings.simplefilter("ignore", DeprecationWarning)
+                pio.write_image(
+                    figure, FIGURES / name, format="png", width=900, scale=2
+                )
+        else:
+            payload = png if isinstance(png, str) else "".join(png)
+            (FIGURES / name).write_bytes(base64.b64decode(payload))
+        paths.append(f"{FIGURE_URL}/{name}")
+    return paths
+
+
+def _render(path: Path, order: int) -> str:
+    notebook = json.loads(path.read_text())
+    cells = notebook["cells"]
+    title, description = _title_and_description(cells)
+    relative = path.relative_to(NOTEBOOKS).as_posix()
+    slug = path.stem
+
+    # The header cell becomes the frontmatter plus the intro block: title, then
+    # everything under it (overview, questions, outline, model, extras).
+    intro = "".join(cells[0]["source"]).split("\n", 1)[1].strip()
+
+    lines = [
+        "---",
+        # json.dumps yields a double-quoted scalar, which YAML accepts and which
+        # survives the colons these titles contain.
+        f"title: {json.dumps(title)}",
+        f"description: {json.dumps(description)}",
+        "sidebar:",
+        f"  order: {order}",
+        "---",
+        "",
+        ":::note",
+        f"Generated from [`notebooks/{relative}`]({GITHUB}/{relative}), which you "
+        "can run yourself. The outputs below are the ones it produced.",
+        ":::",
+        "",
+        intro,
+        "",
+    ]
+
+    counter = [0]
+    for cell in cells[1:]:
+        if cell["cell_type"] == "markdown":
+            lines += ["".join(cell["source"]), ""]
+            continue
+
+        source = "".join(cell["source"]).strip()
+        if source:
+            lines += ["```python", source, "```", ""]
+
+        text = _stream_text(cell)
+        if text:
+            shown = text.splitlines()
+            if len(shown) > _MAX_OUTPUT_LINES:
+                dropped = len(shown) - _MAX_OUTPUT_LINES
+                shown = shown[:_MAX_OUTPUT_LINES] + [f"... ({dropped} more lines)"]
+            lines += ["```text", *shown, "```", ""]
+
+        for image in _write_figures(cell, slug, counter):
+            lines += [f"![{title} figure {counter[0]}]({image})", ""]
+
+    return "\n".join(lines).rstrip() + "\n"
+
+
+def main() -> None:
+    try:
+        import plotly  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - deploy-time guard
+        raise SystemExit(
+            "gen_notebook_docs needs the plot extra: uv sync --all-extras"
+        ) from exc
+
+    OUT.mkdir(parents=True, exist_ok=True)
+    FIGURES.mkdir(parents=True, exist_ok=True)
+
+    missing = [name for name in ORDER if not (NOTEBOOKS / name).exists()]
+    if missing:
+        raise SystemExit(f"listed in ORDER but not on disk: {missing}")
+
+    found = {
+        p.relative_to(NOTEBOOKS).as_posix()
+        for p in [NOTEBOOKS / "getting_started.ipynb"]
+        + sorted((NOTEBOOKS / "applications").glob("*.ipynb"))
+    }
+    unlisted = found - set(ORDER)
+    if unlisted:
+        raise SystemExit(f"notebooks missing from ORDER: {sorted(unlisted)}")
+
+    for position, name in enumerate(ORDER, 1):
+        page = _render(NOTEBOOKS / name, order=position)
+        (OUT / f"{Path(name).stem}.md").write_text(page)
+        print(f"  {name} -> {OUT.relative_to(REPO_ROOT)}/{Path(name).stem}.md")
+
+    figures = len(list(FIGURES.glob("*.png")))
+    print(f"\n{len(ORDER)} notebook pages, {figures} figures")
+
+
+if __name__ == "__main__":
+    main()

docs/src/pages/reproductions/function_vectors.mdx

@@ -1,7 +1,7 @@
 ---
 layout: ../../layouts/ReproLayout.astro
 title: "Function Vectors: a portable task representation"
-description: Reproducing Todd et al. (ICLR 2024): a small set of attention heads carry a compact, portable representation of an in-context task, summing them gives a function vector, and adding it to a zero-shot prompt makes GPT-J perform the task, all with Murano recording, causal mediation, and interventions.
+description: "Reproducing Todd et al. (ICLR 2024): a small set of attention heads carry a compact, portable representation of an in-context task, summing them gives a function vector, and adding it to a zero-shot prompt makes GPT-J perform the task, all with Murano recording, causal mediation, and interventions."
 ---
 
 <div class="mc-header">