circuitkit 0.1.2__tar.gz → 0.1.4__tar.gz

{circuitkit-0.1.2 → circuitkit-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: circuitkit
-Version: 0.1.2
+Version: 0.1.4
 Summary: Unified Discover, Evaluate, Intervene toolkit for mechanistic interpretability: 13 circuit discovery algorithms, 6-pillar faithfulness evaluation, and pruning / quantization / knowledge editing / unlearning / activation steering / contrastive weight steering interventions behind one API.
 Author: Hem Gosalia, Aditya Kasliwal, Vinay Kumar Sankarapu
 Author-email: Pratinav Seth <pratinav.seth@lexsi.ai>
@@ -84,7 +84,10 @@ Requires-Dist: torchvision==0.21.0+cu126; extra == "gpu-cu126"
 Dynamic: license-file
 
 <p align="center">
-  <img src="docs/assets/circuitkit-mark-v3.svg" width="140" alt="CircuitKit">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="docs/assets/circuitkit-logo-white.png">
+    <img src="docs/assets/circuitkit-logo-black.png" width="360" alt="CircuitKit">
+  </picture>
 </p>
 
 <h1 align="center">CircuitKit</h1>

{circuitkit-0.1.2 → circuitkit-0.1.4}/README.md

@@ -1,5 +1,8 @@
 <p align="center">
-  <img src="docs/assets/circuitkit-mark-v3.svg" width="140" alt="CircuitKit">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="docs/assets/circuitkit-logo-white.png">
+    <img src="docs/assets/circuitkit-logo-black.png" width="360" alt="CircuitKit">
+  </picture>
 </p>
 
 <h1 align="center">CircuitKit</h1>

{circuitkit-0.1.2 → circuitkit-0.1.4}/docs/index.md

@@ -6,7 +6,7 @@ hide:
 
 <div class="ck-hero" markdown>
 
-<div class="ck-hero-logo-wrap"><img class="ck-hero-mark" src="assets/circuitkit-mark-v3.svg" alt=""></div>
+<div class="ck-hero-logo-wrap"><img class="ck-hero-mark ck-lexsi-on-light" src="assets/circuitkit-mark-black.svg" alt=""><img class="ck-hero-mark ck-lexsi-on-dark" src="assets/circuitkit-mark-white.svg" alt=""></div>
 
 <h1 class="ck-hero-title">Discover a circuit.<br><em>Then act on it.</em></h1>
 

{circuitkit-0.1.2 → circuitkit-0.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "circuitkit"
-version = "0.1.2"
+version = "0.1.4"
 # Pratinav Seth, Hem Gosalia, and Aditya Kasliwal are equal co-first authors.
 # Pratinav Seth is the corresponding author. Author order is kept consistent
 # with src/circuitkit/__init__.py and CITATION.cff.

{circuitkit-0.1.2 → circuitkit-0.1.4}/src/circuitkit/__init__.py

@@ -3,7 +3,7 @@ CircuitKit: A comprehensive toolkit for circuit discovery in transformer models.
 """
 
 # Version information
-__version__ = "0.1.2"
+__version__ = "0.1.4"
 __author__ = "Pratinav Seth, Hem Gosalia, Aditya Kasliwal, Vinay Kumar Sankarapu"
 __description__ = "Unified Discover, Evaluate, Intervene toolkit for mechanistic interpretability"
 _API_EXPORTS = {"discover_circuit", "evaluate_circuit", "load_circuit"}

{circuitkit-0.1.2 → circuitkit-0.1.4}/src/circuitkit/api.py

@@ -1399,6 +1399,24 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                                     (abs(score.item()), (node.name, neuron_idx))
                                 )
 
+                    # Per-node aggregate (max abs score among a node's neurons) so
+                    # `Circuit.from_artifact()` -- which only ever looks for a
+                    # "node_scores" key in the scores side-car, the same shape
+                    # node-level discovery produces -- gets a non-empty
+                    # `Circuit.scores` for neuron-level circuits too. Without
+                    # this, `Circuit.scores` silently ends up `{}` for every
+                    # neuron-level run (including Pipeline.discover()'s own
+                    # first, in-process Circuit construction, not just a later
+                    # reload), which is why neuron-level circuit graphs render
+                    # with every node/edge looking unscored. `neurons_scores`
+                    # (the full per-neuron tensor) is kept in the side-car
+                    # alongside this for anyone who needs the finer-grained
+                    # detail.
+                    node_scores: Dict[str, float] = {}
+                    for abs_score, (node_name, _neuron_idx) in all_neuron_scores:
+                        if abs_score > node_scores.get(node_name, float("-inf")):
+                            node_scores[node_name] = abs_score
+
                     all_neuron_scores.sort(key=lambda x: x[0])  # Sort by absolute score, ascending
                     num_to_prune = int(len(all_neuron_scores) * pruning_cfg["target_sparsity"])
 
@@ -1433,6 +1451,7 @@ def discover_circuit(  # noqa: C901 - complex function, refactor out of scope fo
                         {
                             "algo": algo,
                             "level": "neuron",
+                            "node_scores": node_scores,
                             "neurons_scores": graph.neurons_scores.cpu(),
                             "total_neurons": len(all_neuron_scores),
                         },

{circuitkit-0.1.2 → circuitkit-0.1.4}/src/circuitkit/circuit.py

@@ -262,7 +262,14 @@ class Circuit:
     # ------------------------------------------------------------------ #
     # Visualisation                                                      #
     # ------------------------------------------------------------------ #
-    def plot(self, output_path: Optional[Union[str, Path]] = None) -> Optional[str]:
+    def plot(
+        self,
+        output_path: Optional[Union[str, Path]] = None,
+        *,
+        max_nodes: Optional[int] = None,
+        max_edges: Optional[int] = None,
+        edge_threshold: float = 0.0,
+    ) -> Union[str, Dict[str, Any], None]:
         """Render the circuit as an interactive graph.
 
         Delegates to :class:`circuitkit.visualize.graph_viz.CircuitGraphVisualizer`.
@@ -271,15 +278,26 @@ class Circuit:
         returns ``None`` instead of raising.
 
         Args:
-            output_path: Optional path for an HTML export. When omitted, an
-                in-Jupyter interactive widget is returned where supported.
+            output_path: Optional export path. A ``.json`` extension writes a
+                bounded graph payload via
+                :meth:`CircuitGraphVisualizer.to_json` (see ``max_nodes`` /
+                ``max_edges`` / ``edge_threshold`` below); any other
+                extension (or ``None``) writes/returns the interactive HTML
+                export as before. Omitted entirely, an in-Jupyter interactive
+                widget is returned where supported.
+            max_nodes: For ``.json`` output only — cap on nodes kept.
+            max_edges: For ``.json`` output only — cap on edges kept.
+            edge_threshold: For ``.json`` output only — minimum normalized
+                edge weight to keep.
 
         Returns:
-            The HTML string when ``output_path`` is given and rendering
-            succeeded, otherwise ``None``.
+            The HTML string when ``output_path`` is an HTML/other path and
+            rendering succeeded; the exported dict when ``output_path`` is a
+            ``.json`` path; otherwise ``None``.
 
         Example:
             >>> circuit.plot("ioi_circuit.html")   # writes an interactive HTML
+            >>> circuit.plot("ioi_circuit.json", max_nodes=300, max_edges=3000)
         """
         if not self.scores:
             logger.info(
@@ -291,7 +309,6 @@ class Circuit:
             from .artifacts.scores import CircuitScores
             from .visualize.graph_viz import CircuitGraphVisualizer
 
-
             cs = self.circuit_scores
             if cs is None:
                 cs = CircuitScores(
@@ -312,9 +329,8 @@ class Circuit:
                 pruned_nodes = self.nodes
             elif isinstance(self.nodes, dict):
                 # neuron-level artifact: flatten mlp + heads keys
-                pruned_nodes = (
-                    list(self.nodes.get("mlp", {}).keys())
-                    + list(self.nodes.get("heads", {}).keys())
+                pruned_nodes = list(self.nodes.get("mlp", {}).keys()) + list(
+                    self.nodes.get("heads", {}).keys()
                 )
 
             graph_dict = {"nodes": {name: {} for name in self.scores}, "edges": []}
@@ -323,6 +339,13 @@ class Circuit:
             if output_path is not None:
                 out = Path(output_path)
                 out.parent.mkdir(parents=True, exist_ok=True)
+                if out.suffix.lower() == ".json":
+                    return viz.to_json(
+                        str(out),
+                        max_nodes=max_nodes,
+                        max_edges=max_edges,
+                        edge_threshold=edge_threshold,
+                    )
                 title = f"Circuit ({self.task or 'unknown'} / {self.algorithm or 'eap-ig'})"
                 return viz.to_html(str(out), title=title)
             return viz.interactive_widget()

{circuitkit-0.1.2 → circuitkit-0.1.4}/src/circuitkit/cli/main.py

@@ -1549,16 +1549,31 @@ def data_prepare(source, shape, strategy, max_records, output, hf_subset, hf_spl
         f"strategy={chosen_strategy or 'native'})"
     )
 
+
 @data.command("template")
 @click.argument("source", type=click.Path(exists=True))
-@click.option("--clean-prompt", required=True, help="Clean prompt template, e.g. 'The capital of {country} is'")
-@click.option("--corrupt-prompt", default=None, help="Corrupt prompt template, e.g. 'The capital of {other_country} is'. Omit for ibcircuit/cdt.")
+@click.option(
+    "--clean-prompt",
+    required=True,
+    help="Clean prompt template, e.g. 'The capital of {country} is'",
+)
+@click.option(
+    "--corrupt-prompt",
+    default=None,
+    help="Corrupt prompt template, e.g. 'The capital of {other_country} is'. Omit for ibcircuit/cdt.",
+)
 @click.option("--clean-answer", required=True, help="Clean answer template, e.g. '{capital}'")
-@click.option("--corrupt-answer", default=None, help="Corrupt answer template, e.g. '{other_capital}'. Omit for ibcircuit/cdt.")
+@click.option(
+    "--corrupt-answer",
+    default=None,
+    help="Corrupt answer template, e.g. '{other_capital}'. Omit for ibcircuit/cdt.",
+)
 @click.option("--pairing-mode", type=click.Choice(["explicit", "auto_peer"]), default="explicit")
 @click.option("--max-records", type=int, default=256)
 @click.option("--output", required=True, help="Output NormalizedDataset JSON path.")
-@click.option("--validate", "run_validation", is_flag=True, default=False, help="Run worthiness checks.")
+@click.option(
+    "--validate", "run_validation", is_flag=True, default=False, help="Run worthiness checks."
+)
 @click.option(
     "--align-strategy",
     type=click.Choice(["filter", "pad_question", "none"]),
@@ -1643,6 +1658,7 @@ def data_template(
             else:
                 try:
                     from transformers import AutoTokenizer
+
                     tokenizer = AutoTokenizer.from_pretrained(model)
                     click.echo(f"Loaded tokenizer from {model}")
                 except Exception as e:
@@ -1676,7 +1692,9 @@ def data_template(
             pad_region_end=pad_region_end,
         )
         if pair_padding_side != "left":
-            ds.meta.setdefault("_alignment", {})["recommended_pair_padding_side"] = pair_padding_side
+            ds.meta.setdefault("_alignment", {})[
+                "recommended_pair_padding_side"
+            ] = pair_padding_side
 
         _alignment = ds.meta.get("_alignment", {})
         if _alignment:
@@ -1692,6 +1710,7 @@ def data_template(
 
         if run_validation:
             from ..data.worthiness import evaluate_worthiness
+
             report = evaluate_worthiness(ds)
             click.echo(report.render_terminal())
 
@@ -1717,6 +1736,7 @@ def data_template(
 
         if run_validation:
             from ..data.worthiness import evaluate_worthiness
+
             report = evaluate_worthiness(ds)
             click.echo(report.render_terminal())
 
@@ -1728,23 +1748,37 @@ def data_template(
 
     if run_validation:
         from ..data.worthiness import evaluate_worthiness
+
         report = evaluate_worthiness(ds)
         click.echo(report.render_terminal())
 
     ds.save_json(output)
     click.echo(f"Saved {len(ds)} template-paired records to {output} (paired={ds.n_paired})")
 
+
 @data.command("clean-only")
 @click.argument("source", type=click.Path(exists=True))
-@click.option("--prompt-column", default="prompt", show_default=True,
-              help="CSV column name for the clean prompt.")
-@click.option("--answer-column", default="answer", show_default=True,
-              help="CSV column name for the answer. Pass 'none' to skip (valid for CD-T).")
-@click.option("--max-records", type=int, default=None,
-              help="Truncate to this many records.")
+@click.option(
+    "--prompt-column",
+    default="prompt",
+    show_default=True,
+    help="CSV column name for the clean prompt.",
+)
+@click.option(
+    "--answer-column",
+    default="answer",
+    show_default=True,
+    help="CSV column name for the answer. Pass 'none' to skip (valid for CD-T).",
+)
+@click.option("--max-records", type=int, default=None, help="Truncate to this many records.")
 @click.option("--output", required=True, help="Output NormalizedDataset JSON path.")
-@click.option("--validate", "run_validation", is_flag=True, default=False,
-              help="Run worthiness checks after loading.")
+@click.option(
+    "--validate",
+    "run_validation",
+    is_flag=True,
+    default=False,
+    help="Run worthiness checks after loading.",
+)
 def data_clean_only(source, prompt_column, answer_column, max_records, output, run_validation):
     """Load SOURCE CSV as a clean-only dataset (no corrupt partner).
 
@@ -1766,6 +1800,7 @@ def data_clean_only(source, prompt_column, answer_column, max_records, output, r
 
     if run_validation:
         from ..data.worthiness import evaluate_worthiness
+
         report = evaluate_worthiness(ds)
         click.echo(report.render_terminal())
 
@@ -1775,6 +1810,7 @@ def data_clean_only(source, prompt_column, answer_column, max_records, output, r
         f"(compatible algorithms: ibcircuit, cdt)"
     )
 
+
 @data.command("shapes")
 def data_shapes():
     """List supported dataset shapes + default strategies."""
@@ -1837,6 +1873,7 @@ def _resolve_source(source, hf_subset, hf_split, max_records):
 # inspect
 # ---------------------------------------------------------------------------
 
+
 @cli.command()
 @click.argument("artifact_path", type=click.Path(exists=True))
 def inspect(artifact_path):
@@ -1874,22 +1911,29 @@ def inspect(artifact_path):
 # prune
 # ---------------------------------------------------------------------------
 
+
 @cli.command()
 @click.option("--model", "-m", required=True, help="Model name or path (e.g., gpt2)")
 @click.option(
-    "--artifact", "-a", required=True, type=click.Path(exists=True),
-    help="Path to circuit artifact (.pt)"
+    "--artifact",
+    "-a",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to circuit artifact (.pt)",
+)
+@click.option(
+    "--sparsity", "-s", type=float, default=0.3, show_default=True, help="Target sparsity (0.0–1.0)"
 )
-@click.option("--sparsity", "-s", type=float, default=0.3, show_default=True,
-              help="Target sparsity (0.0–1.0)")
 @click.option(
-    "--scope", type=click.Choice(["heads", "mlp", "both"]), default="both", show_default=True,
-    help="Which component type to prune"
+    "--scope",
+    type=click.Choice(["heads", "mlp", "both"]),
+    default="both",
+    show_default=True,
+    help="Which component type to prune",
 )
 @click.option("--output", "-o", required=True, help="Output checkpoint directory path")
 @click.option(
-    "--precision", default="bfloat16", show_default=True,
-    help="Torch dtype for model loading"
+    "--precision", default="bfloat16", show_default=True, help="Torch dtype for model loading"
 )
 def prune(model, artifact, sparsity, scope, output, precision):
     """Prune a model using a discovered circuit and export a HF checkpoint."""
@@ -1903,6 +1947,7 @@ def prune(model, artifact, sparsity, scope, output, precision):
         raise click.Abort()
 
     from ..circuit import Circuit
+
     try:
         circuit = Circuit.from_artifact(artifact)
     except Exception as exc:
@@ -1926,24 +1971,34 @@ def prune(model, artifact, sparsity, scope, output, precision):
 # quantize
 # ---------------------------------------------------------------------------
 
+
 @cli.command()
 @click.option("--model", "-m", required=True, help="Model name or path")
 @click.option(
-    "--artifact", "-a", required=True, type=click.Path(exists=True),
-    help="Path to circuit artifact (.pt)"
+    "--artifact",
+    "-a",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to circuit artifact (.pt)",
+)
+@click.option("--bits", type=int, default=4, show_default=True, help="Quantization bit-width")
+@click.option(
+    "--high-fraction",
+    type=float,
+    default=0.3,
+    show_default=True,
+    help="Fraction of top layers kept at high precision",
 )
-@click.option("--bits", type=int, default=4, show_default=True,
-              help="Quantization bit-width")
-@click.option("--high-fraction", type=float, default=0.3, show_default=True,
-              help="Fraction of top layers kept at high precision")
 @click.option(
-    "--backend", type=click.Choice(["quanto", "llmcompressor"]), default="quanto",
-    show_default=True, help="Quantization backend"
+    "--backend",
+    type=click.Choice(["quanto", "llmcompressor"]),
+    default="quanto",
+    show_default=True,
+    help="Quantization backend",
 )
 @click.option("--output", "-o", required=True, help="Output checkpoint directory path")
 @click.option(
-    "--precision", default="bfloat16", show_default=True,
-    help="Torch dtype for model loading"
+    "--precision", default="bfloat16", show_default=True, help="Torch dtype for model loading"
 )
 def quantize(model, artifact, bits, high_fraction, backend, output, precision):
     """Apply circuit-guided mixed-precision quantization and export a HF checkpoint."""
@@ -1957,6 +2012,7 @@ def quantize(model, artifact, bits, high_fraction, backend, output, precision):
         raise click.Abort()
 
     from ..circuit import Circuit
+
     try:
         circuit = Circuit.from_artifact(artifact)
     except Exception as exc:
@@ -1969,7 +2025,8 @@ def quantize(model, artifact, bits, high_fraction, backend, output, precision):
     )
     try:
         quantized = quick.quantize(
-            tl_model, circuit,
+            tl_model,
+            circuit,
             bits=bits,
             high_fraction=high_fraction,
             backend=backend,
@@ -1985,26 +2042,41 @@ def quantize(model, artifact, bits, high_fraction, backend, output, precision):
 # export
 # ---------------------------------------------------------------------------
 
+
 @cli.command()
 @click.option("--model", "-m", required=True, help="Model name or path")
 @click.option(
-    "--artifact", "-a", required=True, type=click.Path(exists=True),
-    help="Path to circuit artifact (.pt)"
+    "--artifact",
+    "-a",
+    required=True,
+    type=click.Path(exists=True),
+    help="Path to circuit artifact (.pt)",
 )
 @click.option("--output", "-o", required=True, help="Output checkpoint directory path")
 @click.option(
-    "--intervention", type=click.Choice(["pruning", "quantization"]), default="pruning",
-    show_default=True, help="Intervention type applied before export"
+    "--intervention",
+    type=click.Choice(["pruning", "quantization"]),
+    default="pruning",
+    show_default=True,
+    help="Intervention type applied before export",
 )
-@click.option("--sparsity", "-s", type=float, default=0.3, show_default=True,
-              help="Sparsity for pruning intervention")
 @click.option(
-    "--scope", type=click.Choice(["heads", "mlp", "both"]), default="both", show_default=True,
-    help="Pruning scope (only used when intervention=pruning)"
+    "--sparsity",
+    "-s",
+    type=float,
+    default=0.3,
+    show_default=True,
+    help="Sparsity for pruning intervention",
+)
+@click.option(
+    "--scope",
+    type=click.Choice(["heads", "mlp", "both"]),
+    default="both",
+    show_default=True,
+    help="Pruning scope (only used when intervention=pruning)",
 )
 @click.option(
-    "--precision", default="bfloat16", show_default=True,
-    help="Torch dtype for model loading"
+    "--precision", default="bfloat16", show_default=True, help="Torch dtype for model loading"
 )
 def export(model, artifact, output, intervention, sparsity, scope, precision):
     """Apply an intervention to a model and export a HuggingFace checkpoint.
@@ -2049,6 +2121,7 @@ def export(model, artifact, output, intervention, sparsity, scope, precision):
 # run (YAML pipeline)
 # ---------------------------------------------------------------------------
 
+
 @cli.command()
 @click.argument("config_path", type=click.Path(exists=True))
 def run(config_path):
@@ -2106,6 +2179,14 @@ def run(config_path):
     if disc:
         console.print("[cyan]Step: discovery[/cyan]")
         try:
+            # ig_steps/mlp_hook/chat_template_mode reach Pipeline.discover()'s
+            # **kw only when the YAML sets them, so unset keys behave exactly
+            # as before (no seed override, no extra discovery-config keys).
+            extra_disc_kw = {
+                key: disc[key]
+                for key in ("ig_steps", "mlp_hook", "chat_template_mode")
+                if key in disc
+            }
             pipe.discover(
                 algorithm=disc.get("algorithm", "eap-ig"),
                 level=disc.get("level", "node"),
@@ -2113,6 +2194,8 @@ def run(config_path):
                 n_examples=disc.get("n_examples", 128),
                 batch_size=disc.get("batch_size", 4),
                 scope=disc.get("scope", "both"),
+                seed=disc.get("seed"),
+                **extra_disc_kw,
             )
             console.print(f"  Circuit: {pipe._circuit!r}")
         except Exception as exc:
@@ -2127,6 +2210,8 @@ def run(config_path):
             pipe.evaluate(
                 pillars=eval_cfg.get("pillars"),
                 n_examples=eval_cfg.get("n_examples", 256),
+                n_stability_runs=eval_cfg.get("n_stability_runs", 5),
+                target_task=eval_cfg.get("target_task"),
             )
         except Exception as exc:
             console.print(f"[yellow]Warning: evaluation failed:[/yellow] {exc}")
@@ -2188,9 +2273,18 @@ def run(config_path):
     if viz_cfg and viz_cfg.get("enabled", True) and pipe._circuit is not None:
         console.print("[cyan]Step: visualize[/cyan]")
         try:
+            # max_nodes/max_edges/edge_threshold only apply to a `.json` output
+            # (a bounded graph payload); only forward them when the YAML sets
+            # them, so other visualize modes (comparison/dashboard) are unaffected.
+            extra_viz_kw = {
+                key: viz_cfg[key]
+                for key in ("max_nodes", "max_edges", "edge_threshold")
+                if key in viz_cfg
+            }
             pipe.visualize(
                 mode=viz_cfg.get("mode", "graph"),
                 output=viz_cfg.get("output"),
+                **extra_viz_kw,
             )
         except Exception as exc:
             console.print(f"[yellow]Warning: visualize failed:[/yellow] {exc}")

{circuitkit-0.1.2 → circuitkit-0.1.4}/src/circuitkit/quick.py

@@ -39,7 +39,7 @@ if TYPE_CHECKING:  # pragma: no cover - typing only
     from transformer_lens import HookedTransformer
 
     from .evaluation.report import FaithfulnessReport
-    
+
     from pathlib import Path
 
 __all__ = [
@@ -656,9 +656,7 @@ def quantize(
     import re
 
     if backend not in ("quanto", "llmcompressor"):
-        raise ValueError(
-            f"backend must be 'quanto' or 'llmcompressor', got {backend!r}"
-        )
+        raise ValueError(f"backend must be 'quanto' or 'llmcompressor', got {backend!r}")
 
     if not circuit.scores:
         raise ValueError(
@@ -688,9 +686,7 @@ def quantize(
         n_layers = getattr(cfg, "num_hidden_layers", None)
         if n_layers is None and cfg is not None:
             # Gemma-3 / multimodal configs nest the decoder under text_config.
-            n_layers = getattr(
-                getattr(cfg, "text_config", None), "num_hidden_layers", None
-            )
+            n_layers = getattr(getattr(cfg, "text_config", None), "num_hidden_layers", None)
         if n_layers is None and max_layer >= 0:
             n_layers = max_layer + 1
     if not n_layers:
@@ -702,9 +698,9 @@ def quantize(
         if tokenizer is None:
             from transformers import AutoTokenizer
 
-            repo_id = getattr(
-                getattr(model, "config", None), "_name_or_path", None
-            ) or circuit.model_name
+            repo_id = (
+                getattr(getattr(model, "config", None), "_name_or_path", None) or circuit.model_name
+            )
             if not repo_id:
                 raise ValueError(
                     "quantize(backend='llmcompressor') needs a tokenizer; pass "
@@ -812,9 +808,7 @@ def export_checkpoint(
     return path
 
 
-def _push_checkpoint_to_hub(
-    path: str, hub_repo: Optional[str], hub_private: bool
-) -> None:
+def _push_checkpoint_to_hub(path: str, hub_repo: Optional[str], hub_private: bool) -> None:
     """Upload a written checkpoint directory to the HuggingFace Hub.
 
     Opt-in archival path for :func:`export_checkpoint`. The local checkpoint is
@@ -822,9 +816,7 @@ def _push_checkpoint_to_hub(
     authenticated token (``huggingface-cli login`` or ``HF_TOKEN``).
     """
     if not hub_repo:
-        raise ValueError(
-            "export_checkpoint(push_to_hub=True) needs hub_repo='org/name'."
-        )
+        raise ValueError("export_checkpoint(push_to_hub=True) needs hub_repo='org/name'.")
     try:
         from huggingface_hub import HfApi
     except ImportError as exc:  # pragma: no cover - optional dependency
@@ -898,7 +890,8 @@ def benchmark(
         dtype=dtype,
         **kw,
     )
-    
+
+
 # --------------------------------------------------------------------------- #
 # load_scores                                                                  #
 # --------------------------------------------------------------------------- #
@@ -1026,6 +1019,7 @@ def selective_finetune(
             )
         try:
             from transformers import AutoConfig
+
             hf_cfg = AutoConfig.from_pretrained(resolved_model)
             if n_layers is None:
                 n_layers = hf_cfg.num_hidden_layers
@@ -1105,7 +1099,7 @@ def visualize_circuit(
         ...              second_circuit=circuit2, output="compare.html")
     """
     if mode == "graph":
-        return circuit.plot(output)
+        return circuit.plot(output, **kw)
 
     if mode == "comparison":
         from .visualize.comparison import ComparisonDashboard
@@ -1152,6 +1146,5 @@ def visualize_circuit(
         return None
 
     raise ValueError(
-        f"Unknown visualization mode: {mode!r}. "
-        "Use 'graph', 'comparison', or 'dashboard'."
+        f"Unknown visualization mode: {mode!r}. " "Use 'graph', 'comparison', or 'dashboard'."
     )