circuitkit 0.1.0__py3-none-any.whl

circuitkit/__init__.py

@@ -0,0 +1,128 @@
+"""
+CircuitKit: A comprehensive toolkit for circuit discovery in transformer models.
+"""
+
+# Version information
+__version__ = "0.1.0"
+__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"}
+
+# Flat front-door API (circuitkit.quick) — lazily imported so `import circuitkit`
+# stays fast and torch-free until one of these is actually accessed.
+_QUICK_EXPORTS = {
+    "load_model",
+    "discover",
+    "faithfulness",
+    "prune",
+    "quantize",
+    "export_checkpoint",
+    "benchmark",
+    "load_scores",
+    "selective_finetune",
+    "visualize_circuit",
+}
+_CIRCUIT_EXPORTS = {"Circuit"}
+_PIPELINE_EXPORTS = {"Pipeline"}
+
+# Main exports
+__all__ = [
+    # Core dict-config API
+    "discover_circuit",
+    "evaluate_circuit",
+    "load_circuit",
+    # Flat front-door API
+    "load_model",
+    "discover",
+    "faithfulness",
+    "prune",
+    "quantize",
+    "export_checkpoint",
+    "benchmark",
+    "Circuit",
+    "load_scores",
+    "selective_finetune",
+    "visualize_circuit",
+    "Pipeline",
+    # Task management
+    "get_task",
+    "list_tasks",
+    "register_task",
+    # Version info
+    "__version__",
+    "__author__",
+    "__description__",
+]
+
+
+def __getattr__(name):
+    """Lazily import heavy API helpers when accessed from the package root."""
+    if name in _API_EXPORTS:
+        from . import api
+
+        value = getattr(api, name)
+        globals()[name] = value
+        return value
+    if name in _QUICK_EXPORTS:
+        from . import quick
+
+        value = getattr(quick, name)
+        globals()[name] = value
+        return value
+    if name in _CIRCUIT_EXPORTS:
+        from .circuit import Circuit
+
+        globals()["Circuit"] = Circuit
+        return Circuit
+    if name in _PIPELINE_EXPORTS:
+        from .pipeline import Pipeline
+
+        globals()["Pipeline"] = Pipeline
+        return Pipeline
+    if name == "visualize":
+        # Deprecated pre-1.0 alias (H1 in the 1.0.0 audit): ``visualize`` was
+        # renamed to ``visualize_circuit`` in 1.0.0 and is now a subpackage.
+        # Return the (callable) subpackage so ``ck.visualize(...)`` keeps working
+        # and warns on use — see circuitkit/visualize/__init__.py. The warning
+        # fires on the call, not on attribute access, so it survives the
+        # submodule being imported (which shadows any parent __getattr__ shim).
+        from . import visualize as _visualize
+
+        return _visualize
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    """Expose lazily-imported names for autocomplete / dir()."""
+    return sorted(set(globals()) | set(__all__))
+
+
+def _ensure_builtin_tasks():
+    """Register built-in tasks only when task helpers are used."""
+    from .tasks.bootstrap import _bootstrap_builtin_tasks
+
+    return _bootstrap_builtin_tasks()
+
+
+def get_task(name):
+    """Get a built-in or registered task by name."""
+    _ensure_builtin_tasks()
+    from .tasks.registry import get_task as _get_task
+
+    return _get_task(name)
+
+
+def list_tasks():
+    """List registered task names."""
+    _ensure_builtin_tasks()
+    from .tasks.registry import list_tasks as _list_tasks
+
+    return _list_tasks()
+
+
+def register_task(spec):
+    """Register a custom task specification."""
+    _ensure_builtin_tasks()
+    from .tasks.registry import register_task as _register_task
+
+    return _register_task(spec)

circuitkit/__main__.py

@@ -0,0 +1,9 @@
+"""
+CircuitKit module execution entry point.
+Allows running: python -m circuitkit
+"""
+
+from .cli.main import main
+
+if __name__ == "__main__":
+    main()

circuitkit/analysis/__init__.py

@@ -0,0 +1,19 @@
+"""
+CircuitKit Analysis Module
+
+Provides analysis tools for circuits including metrics, scoring, and statistical analysis.
+"""
+
+from .cross_method_jaccard import CrossMethodJaccardResult, cross_method_jaccard
+from .metrics import *  # noqa: F401,F403 - intentional API re-export
+from .scores import *  # noqa: F401,F403 - intentional API re-export
+
+__all__ = [  # noqa: F405 - names provided via star imports above
+    # Metrics
+    "compute_metrics",
+    # Scores
+    "compute_scores",
+    # Cross-method comparison (EMNLP 2026, Section 5)
+    "cross_method_jaccard",
+    "CrossMethodJaccardResult",
+]

circuitkit/analysis/cross_method_jaccard.py

@@ -0,0 +1,116 @@
+"""Cross-method circuit comparison via top-K Jaccard.
+
+Used by the EMNLP 2026 submission's Finding B (Section 5) to
+quantify component-level disagreement across discovery methods
+discovering circuits for the same task on the same model.
+
+Example
+-------
+>>> from circuitkit.analysis.cross_method_jaccard import (
+...     cross_method_jaccard,
+... )
+>>> # circuits: dict mapping method name to a CircuitScores artifact
+>>> # (or a list of node names, or a dict node_name -> score)
+>>> result = cross_method_jaccard(circuits, top_k=10)
+>>> # result.matrix:   numpy-style nested list of pairwise Jaccards
+>>> # result.top_nodes: dict method -> top_k node names
+>>> # result.range:    (min, max) pairwise Jaccard over off-diagonal cells
+
+The Jaccard at top-K is the standard descriptive statistic for the
+"do methods agree on the circuit?" question reported in the EMNLP
+paper, derived from the workshop result on GPT-2 IOI where the same
+analysis showed Jaccard 0.11 to 1.00 across 8 EAP-family variants.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Tuple, Union
+
+CircuitLike = Union[Dict[str, float], List[str], object]
+
+
+def _extract_top_k_nodes(circuit: CircuitLike, top_k: int) -> List[str]:
+    """Coerce a circuit-like input into a top-K node-name list.
+
+    Accepts:
+      - dict[str, float]: node_name -> score; top-K by absolute score
+      - list[str]: assumed to be a pre-ranked list; truncated to top-K
+      - object with `.node_scores` attribute (e.g. CircuitScores): dict-like
+    """
+    if hasattr(circuit, "node_scores"):
+        node_scores = circuit.node_scores
+    else:
+        node_scores = circuit
+    if isinstance(node_scores, dict):
+        ranked = sorted(node_scores.items(), key=lambda kv: abs(float(kv[1])), reverse=True)
+        return [n for n, _ in ranked[:top_k]]
+    if isinstance(node_scores, list):
+        return list(node_scores[:top_k])
+    raise TypeError(f"Unsupported circuit type for jaccard: {type(circuit)!r}")
+
+
+def _jaccard(a: List[str], b: List[str]) -> float:
+    sa, sb = set(a), set(b)
+    if not sa and not sb:
+        return 1.0
+    return len(sa & sb) / len(sa | sb)
+
+
+@dataclass
+class CrossMethodJaccardResult:
+    methods: List[str]
+    top_k: int
+    matrix: List[List[float]]
+    top_nodes: Dict[str, List[str]]
+    range: Tuple[float, float]
+    n_pairs: int
+
+
+def cross_method_jaccard(
+    circuits: Dict[str, CircuitLike],
+    top_k: int = 10,
+) -> CrossMethodJaccardResult:
+    """Compute the symmetric pairwise Jaccard matrix between method
+    circuits at top-K nodes.
+
+    Parameters
+    ----------
+    circuits : dict
+        Maps method name to either a CircuitScores artifact, a dict of
+        node_name to score, or a pre-ranked list of node names.
+    top_k : int
+        Truncation depth for each method's circuit. Default 10 to match
+        the EMNLP paper's protocol.
+
+    Returns
+    -------
+    CrossMethodJaccardResult
+        Includes the pairwise Jaccard matrix, the top-K node sets per
+        method, and the off-diagonal range (min, max).
+    """
+    methods = sorted(circuits.keys())
+    n = len(methods)
+    top_nodes = {m: _extract_top_k_nodes(circuits[m], top_k) for m in methods}
+    matrix = [[1.0] * n for _ in range(n)]
+    off_diag: List[float] = []
+    for i, mi in enumerate(methods):
+        for j, mj in enumerate(methods):
+            if i <= j:
+                continue
+            j_val = _jaccard(top_nodes[mi], top_nodes[mj])
+            matrix[i][j] = j_val
+            matrix[j][i] = j_val
+            off_diag.append(j_val)
+    rng = (min(off_diag), max(off_diag)) if off_diag else (0.0, 0.0)
+    return CrossMethodJaccardResult(
+        methods=methods,
+        top_k=top_k,
+        matrix=matrix,
+        top_nodes=top_nodes,
+        range=rng,
+        n_pairs=len(off_diag),
+    )
+
+
+__all__ = ["cross_method_jaccard", "CrossMethodJaccardResult"]

circuitkit/analysis/metrics.py

@@ -0,0 +1,54 @@
+# FILE: circuitkit/analysis/metrics.py
+
+
+def calculate_faithfulness(original_output, pruned_output):
+    """
+    Calculate faithfulness metric between original and pruned outputs.
+    Uses multiple metrics for comprehensive evaluation.
+    """
+    import torch
+    import torch.nn.functional as F
+
+    # Ensure tensors are the same shape
+    if original_output.shape != pruned_output.shape:
+        min_size = min(original_output.numel(), pruned_output.numel())
+        original_output = original_output.flatten()[:min_size]
+        pruned_output = pruned_output.flatten()[:min_size]
+
+    # L2 norm difference
+    l2_diff = (original_output - pruned_output).pow(2).sum().item()
+
+    # KL divergence
+    try:
+        kl_div = F.kl_div(
+            F.log_softmax(pruned_output, dim=-1),
+            F.softmax(original_output, dim=-1),
+            reduction="sum",
+        ).item()
+    except (RuntimeError, ValueError):
+        kl_div = float("inf")
+
+    # Cosine similarity
+    cos_sim = F.cosine_similarity(
+        original_output.flatten().unsqueeze(0), pruned_output.flatten().unsqueeze(0)
+    ).item()
+
+    # Relative difference
+    rel_diff = torch.abs(original_output - pruned_output).sum().item() / (
+        torch.abs(original_output).sum().item() + 1e-8
+    )
+
+    return {
+        "l2_difference": l2_diff,
+        "kl_divergence": kl_div,
+        "cosine_similarity": cos_sim,
+        "relative_difference": rel_diff,
+        "faithfulness_score": 1.0 - min(rel_diff, 1.0),  # Higher is better
+    }
+
+
+def calculate_complexity(graph):
+    """
+    Calculates the complexity of a circuit, e.g., by node or edge count.
+    """
+    return {"node_count": graph.number_of_nodes(), "edge_count": graph.number_of_edges()}

circuitkit/analysis/scores.py

@@ -0,0 +1,44 @@
+# FILE: circuitkit/analysis/scores.py
+
+from collections import defaultdict
+
+from ..backends.acdc.types import PruneScores
+from ..backends.acdc.utils.patchable_model import PatchableModel
+
+
+def calculate_node_scores_from_edges(
+    p_model: PatchableModel, edge_prune_scores: PruneScores
+) -> dict[str, float]:
+    """
+    Calculates an importance score for each source node based on its outgoing edges.
+
+    The score for a node is the average of the absolute scores of all its
+    outgoing edges. Nodes with no outgoing edges will not be included.
+
+    Args:
+        p_model: The patchable model, used to access the graph of nodes and edges.
+        edge_prune_scores: A dictionary mapping destination modules to tensors of
+                           edge importance scores.
+
+    Returns:
+        A dictionary mapping source node names to their calculated importance score.
+    """
+    # Use absolute values of scores as importance can be positive or negative in EAP
+    abs_edge_scores = {mod: scores.abs() for mod, scores in edge_prune_scores.items()}
+
+    # Group edge scores by their source node
+    outgoing_scores_by_node = defaultdict(list)
+    for edge in p_model.edges:
+        # edge.prune_score looks up the score for this specific edge
+        score = edge.prune_score(abs_edge_scores).item()
+        outgoing_scores_by_node[edge.src.name].append(score)
+
+    # Calculate the average score for each node. Use Python's built-in sum/len
+    # rather than np.mean() so values stay plain Python floats — numpy scalars
+    # (numpy.float64) are rejected by torch.load with weights_only=True.
+    node_scores = {}
+    for node_name, scores in outgoing_scores_by_node.items():
+        if scores:
+            node_scores[node_name] = sum(scores) / len(scores)
+
+    return node_scores