AutoCircuit

https://ufo-101.github.io/auto-circuit/
Admin
AutoCircuit is a Python library designed for efficient edge patching and automatic circuit discovery...
Tokens:176,732
Snippets:673
Trust Score:9.7
Update:1 month ago
Show doc for...
Tokens:
Raw
### Install auto_circuit with pip

Source: https://ufo-101.github.io/auto-circuit/guides/1%29_Getting_Started

Install the library via pip.

```bash
pip install auto-circuit

```

--------------------------------

### Install pre-commit git hook

Source: https://ufo-101.github.io/auto-circuit/guides/contributing

Installs the git hook for pre-commit, enabling automatic checks on commits.

```bash
pre-commit install

```

--------------------------------

### Serve documentation locally with mkdocs

Source: https://ufo-101.github.io/auto-circuit/guides/contributing

Builds and serves the documentation locally with Material for MkDocs. Requires the Python environment to be activated.

```bash
mkdocs serve

```

--------------------------------

### batch_avg_answer_val(vals: Tensor, batch: PromptPairBatch, wrong_answer: bool = False) -> Tensor

Source: https://ufo-101.github.io/auto-circuit/reference/utils/tensor_ops

Get the average value of the logits (or some function of them) for the correct answers in the batch. Optionally, can get the average for wrong answers instead.

```APIDOC
## batch_avg_answer_val(vals: Tensor, batch: PromptPairBatch, wrong_answer: bool = False) -> Tensor

### Description
Get the average value of the logits (or some function of them) for the correct answers in the batch. If `wrong_answer` is True, it gets the average for the wrong answers instead.

### Method
Function call

### Endpoint
`batch_avg_answer_val(vals, batch, wrong_answer=False)`

### Parameters
- **vals** (Tensor) - Required - The logits values or some tensor of the same shape.
- **batch** (PromptPairBatch) - Required - The batch of prompts and answers.
- **wrong_answer** (bool) - Optional - Whether to get the average value of the wrong answers instead of the correct answers. Default: `False`.

### Returns
- **Tensor** - The average value of the logits for the correct answers in the batch.

### Source Code
```python
def batch_avg_answer_val(
    vals: t.Tensor, batch: PromptPairBatch, wrong_answer: bool = False
) -> t.Tensor:
    """
    Get the average value of the logits (or some function of them) for the correct
    answers in the batch.

    Args:
        vals: The logits values or some tensor of the same shape.
        batch: The batch of prompts and answers.
        wrong_answer: Whether to get the average value of the wrong answers instead of
            the correct answers.

    Returns:
        The average value of the logits for the correct answers in the batch.
    """
    answers = batch.answers if not wrong_answer else batch.wrong_answers
    if isinstance(answers, t.Tensor):
        return vocab_avg_val(vals, answers)
    else:
        # If each prompt has a different number of answers we have a list of tensor
        assert isinstance(answers, list)
        return t.stack([vocab_avg_val(v, a) for v, a in zip(vals, answers)]).mean()
```
```

--------------------------------

### Prepare model and data

Source: https://ufo-101.github.io/auto-circuit/guides/2%29_Circuit_Discovery

Loads a GPT-2 model, loads the IOI dataset, and patches the model for circuit discovery. The model is patched with factorized=True, slice_output='last_seq', and separate_qkv=True.

```python
import torch as t

from auto_circuit.data import load_datasets_from_json
from auto_circuit.experiment_utils import load_tl_model
from auto_circuit.prune_algos.mask_gradient import mask_gradient_prune_scores
from auto_circuit.types import PruneScores
from auto_circuit.utils.graph_utils import patchable_model
from auto_circuit.utils.misc import repo_path_to_abs_path
from auto_circuit.visualize import draw_seq_graph

device = t.device("cuda" if t.cuda.is_available() else "cpu")
model = load_tl_model("gpt2", device)

path = repo_path_to_abs_path("datasets/ioi/ioi_vanilla_template_prompts.json")
train_loader, test_loader = load_datasets_from_json(
    model=model,
    path=path,
    device=device,
    prepend_bos=True,
    batch_size=16,
    train_test_size=(128, 128),
)

model = patchable_model(
    model,
    factorized=True,
    slice_output="last_seq",
    separate_qkv=True,
    device=device,
)
```

--------------------------------

### EdgeCounts Enum

Source: https://ufo-101.github.io/auto-circuit/reference/types

Enum with special values for TestEdges that get computed at runtime.

```APIDOC
## EdgeCounts

### Description
Enum with special values for TestEdges that get computed at runtime.

### Members
- **ALL** = 1 - Test `0, 1, 2, ..., n_edges` edges.
- **GROUPS** = 3 - Group edges by PruneScores and cumulatively add the number of edges in each group in descending order by score.
- **LOGARITHMIC** = 2 - Test `0, 1, 2, ..., 10, 20, ..., 100, 200, ..., 1000, 2000, ...` edges.
```

--------------------------------

### load_datasets_from_json function

Source: https://ufo-101.github.io/auto-circuit/reference/data

Initializes prompt datasets and data loaders, optionally handling KV caching for models with tail divergence.

```python
    kvs = []
    diverge_idx: int = 0
    if model is None:
        clean_prompts = [t.tensor(p).to(device) for p in clean_prompts]
        corrupt_prompts = [t.tensor(p).to(device) for p in corrupt_prompts]
        answers = [t.tensor(a).to(device) for a in answer_strs]
        wrong_answers = [t.tensor(a).to(device) for a in wrong_answer_strs]
        seq_len = clean_prompts[0].shape[0]
        assert not tail_divergence
    else:
        tokenizer: Any = model.tokenizer
        if prepend_bos:
            clean_prompts = [tokenizer.bos_token + p for p in clean_prompts]
            corrupt_prompts = [tokenizer.bos_token + p for p in corrupt_prompts]
        tokenizer.padding_side = "left"
        clean_prompts = tokenizer(clean_prompts, padding=pad, return_tensors="pt")
        corrupt_prompts = tokenizer(corrupt_prompts, padding=pad, return_tensors="pt")
        seq_len = None
        if return_seq_length:
            assert t.all(clean_prompts["attention_mask"] == 1)
            assert t.all(corrupt_prompts["attention_mask"] == 1)
            seq_len = clean_prompts["input_ids"].shape[1]
        ans_dicts: List[Dict] = [tokenizer(a, return_tensors="pt") for a in answer_strs]
        wrong_ans_dicts: List[Dict] = [
            tokenizer(a, return_tensors="pt") for a in wrong_answer_strs
        ]
        clean_prompts = clean_prompts["input_ids"].to(device)
        corrupt_prompts = corrupt_prompts["input_ids"].to(device)
        answers = [a["input_ids"].squeeze(-1).to(device) for a in ans_dicts]
        wrong_answers = [a["input_ids"].squeeze(-1).to(device) for a in wrong_ans_dicts]

        if tail_divergence:
            diverge_idxs = (~(clean_prompts == corrupt_prompts)).int().argmax(dim=1)
            diverge_idx = int(diverge_idxs.min().item())
        if diverge_idx > 0:
            seq_labels = seq_labels[diverge_idx:] if seq_labels is not None else None
            prefixs, cfg, device = [], model.cfg, model.cfg.device
            if isinstance(batch_size, tuple):
                prefixs.append(clean_prompts[: (bs0 := batch_size[0]), :diverge_idx])
                prefixs.append(clean_prompts[: (bs1 := batch_size[1]), :diverge_idx])
                kvs.append(HookedTransformerKeyValueCache.init_cache(cfg, device, bs0))
                kvs.append(HookedTransformerKeyValueCache.init_cache(cfg, device, bs1))
            else:
                prefixs.append(clean_prompts[:batch_size, :diverge_idx])
                kvs.append(
                    HookedTransformerKeyValueCache.init_cache(cfg, device, batch_size)
                )

            for prefix, kv_cache in zip(prefixs, kvs):
                with t.inference_mode():
                    model(prefix, past_kv_cache=kv_cache)
                kv_cache.freeze()

            print("seq_len before divergence", seq_len)
            if return_seq_length:
                assert seq_len is not None
                seq_len -= diverge_idx
            print("seq_len after divergence", seq_len)

            # This must be done AFTER gathering the kv caches
            clean_prompts = clean_prompts[:, diverge_idx:]
            corrupt_prompts = corrupt_prompts[:, diverge_idx:]

    dataset = PromptDataset(clean_prompts, corrupt_prompts, answers, wrong_answers)
    train_set = Subset(dataset, list(range(train_test_size[0])))
    test_set = Subset(dataset, list(range(train_test_size[0], n_train_and_test)))
    train_loader = PromptDataLoader(
        train_set,
        seq_len=seq_len,
        diverge_idx=diverge_idx,
        kv_cache=kvs[0] if len(kvs) > 0 else None,
        seq_labels=seq_labels,
        word_idxs=word_idxs,
        batch_size=batch_size[0] if isinstance(batch_size, tuple) else batch_size,
        shuffle=False,
    )
    test_loader = PromptDataLoader(
        test_set,
        seq_len=seq_len,
        diverge_idx=diverge_idx,
        kv_cache=kvs[-1] if len(kvs) > 0 else None,
        seq_labels=seq_labels,
        word_idxs=word_idxs,
        batch_size=batch_size[1] if isinstance(batch_size, tuple) else batch_size,
        shuffle=False,
    )
    return train_loader, test_loader
```

--------------------------------

### PromptDataLoader Initialization

Source: https://ufo-101.github.io/auto-circuit/reference/data

The constructor for PromptDataLoader, which forces drop_last=True to ensure consistent batch sizes.

```python
super().__init__(
    prompt_dataset, **kwargs, drop_last=True, collate_fn=collate_fn
)
self.seq_len = seq_len
self.diverge_idx = diverge_idx
self.seq_labels = seq_labels
assert kv_cache is None or diverge_idx > 0
self.kv_cache = kv_cache
self.word_idxs = word_idxs
```

--------------------------------

### Get prompt sequence length

Source: https://ufo-101.github.io/auto-circuit/guides/2%29_Circuit_Discovery

Loads the dataset with return_seq_length=True to obtain the sequence length, which is required for token-specific circuit discovery.

```python
train_loader, test_loader = load_datasets_from_json(
    model=model,
    path=path,
    device=device,
    prepend_bos=True,
    batch_size=16,
    train_test_size=(128, 128),
    return_seq_length=True,
)
```

--------------------------------

### Block.__init__ - trivial linear layer

Source: https://ufo-101.github.io/auto-circuit/reference/model_utils/micro_model_utils

Initializes a Block with identity input and output and a 2x2 weight parameter.

```python
def __init__(self):
    super().__init__()
    self.head_inputs = t.nn.Identity()
    self.head_outputs = t.nn.Identity()
    self.weights = t.nn.Parameter(t.tensor([[1.0, 2.0], [3.0, 4.0]]))
```

--------------------------------

### batch_src_ablations

Source: https://ufo-101.github.io/auto-circuit/reference/utils/ablation_activations

Function to get the activations used to ablate each Edge in a model, given a particular set of model inputs and an ablation type.

```APIDOC
## batch_src_ablations

### Description
Get the activations used to ablate each `Edge` in a model, given a particular set of model inputs and an ablation type. See `AblationType` for the different types of ablations that can be computed.

### Function Signature
`batch_src_ablations(model: PatchableModel, sample: Tensor | PromptDataLoader, ablation_type: AblationType = AblationType.RESAMPLE) -> Tensor`

### Parameters
- **model** (`PatchableModel`) - Required - The model to get the ablations for.
- **sample** (`Tensor | PromptDataLoader`) - Required - The data sample to get the ablations for. This is not used for all `ablation_type`s. Either a single batch of inputs or a DataLoader.
- **ablation_type** (`AblationType`) - Optional - The type of ablation to perform. Default: `RESAMPLE`.

### Returns
- **Tensor** - A tensor of activations used to ablate each `Edge` model on the given input. Shape is `[Srcs, ...]` where `Srcs` is the number of `SrcNode`s in the model and `...` is the shape of the activations of the model. In a transformer this will be `[Srcs, batch, seq, d_model]`.

### Source Code
```python
def src_ablations(
    model: PatchableModel,
    sample: t.Tensor | PromptDataLoader,
    ablation_type: AblationType = AblationType.RESAMPLE,
) -> t.Tensor:
    """
    Get the activations used to ablate each [`Edge`][auto_circuit.types.Edge] in a
    model, given a particular set of model inputs and an ablation type. See
    [`AblationType`][auto_circuit.types.AblationType] for the different types of
    ablations that can be computed.

    Args:
        model: The model to get the ablations for.
        sample: The data sample to get the ablations for. This is not used for all
            `ablation_type`s. Either a single batch of inputs or a DataLoader.
        ablation_type: The type of ablation to perform.

    Returns:
        A tensor of activations used to ablated each [`Edge`][auto_circuit.types.Edge]
            model on the given input.  Shape is `[Srcs, ...]` where `Srcs` is the number
            of [`SrcNode`][auto_circuit.types.SrcNode]s in the model and `...` is the
            shape of the activations of the model. In a transformer this will be
            `[Srcs, batch, seq, d_model]`.
    """
    src_outs: Dict[SrcNode, t.Tensor] = {}
    src_modules: Dict[t.nn.Module, List[SrcNode]] = defaultdict(list)
    [src_modules[src.module(model)].append(src) for src in model.srcs]
    with remove_hooks() as handles, inference_mode():
        # Install hooks to collect activations at each src module
        for mod, src_nodes in src_modules.items():
            hook_fn = partial(
                mean_src_out_hook if ablation_type.mean_over_dataset else src_out_hook,
                src_nodes=src_nodes,
                src_outs=src_outs,
                ablation_type=ablation_type,
            )
            handles.add(mod.register_forward_hook(hook_fn))

        if ablation_type.mean_over_dataset:
            # Collect activations over the entire dataset and take the mean
            assert isinstance(sample, PromptDataLoader)
            for batch in sample:
                if ablation_type.clean_dataset:
                    model(batch.clean)
                if ablation_type.corrupt_dataset:
                    model(batch.corrupt)
            # PromptDataLoader has equal size batches, so we can take the mean of means
            mult = int(ablation_type.clean_dataset) + int(ablation_type.corrupt_dataset)
            assert mult == 2 or mult == 1
            for src, src_out in src_outs.items():
                src_outs[src] = src_out / (len(sample) * mult)
        else:
            # Collect activations for a single batch
            assert isinstance(sample, t.Tensor)
            model(sample)

    # Sort the src_outs dict by node idx
    src_outs = dict(sorted(src_outs.items(), key=lambda x: x[0].src_idx))
    assert [src.src_idx for src in src_outs.keys()] == list(range(len(src_outs)))
    return t.stack(list(src_outs.values())).detach()
```
```

--------------------------------

### make_model_patchable: Prepare model for patching (function signature)

Source: https://ufo-101.github.io/auto-circuit/reference/utils/graph_utils

Signatures show what make_model_patchable expects: model, factorized, src_nodes, nodes, device, seq_len, and seq_dim. It takes a Module input and returns a tuple of PatchWrapperImpl sets.

```python
make_model_patchable(model: Module, factorized: bool, src_nodes: Set[SrcNode], nodes: Set[Node], device: device, seq_len: Optional[int] = None, seq_dim: Optional[int] = None) -> Tuple[Set[PatchWrapperImpl], Set[PatchWrapperImpl], Set[PatchWrapperImpl]]
```

--------------------------------

### subnetwork_probing_prune_scores implementation

Source: https://ufo-101.github.io/auto-circuit/reference/prune_algos/subnetwork_probing

The core logic for calculating edge importance, including data loading, inference mode setup, and batch ablation processing.

```python
    assert len(dataloader) > 0, "Dataloader is empty"

    out_slice = model.out_slice
    n_edges = model.n_edges
    n_avoid = len(avoid_edges or [])

    clean_logits: Dict[BatchKey, t.Tensor] = {}
    with t.inference_mode():
        for batch in dataloader:
            clean_logits[batch.key] = model(batch.clean)[out_slice]

    val_clean_logits: Optional[Dict[BatchKey, t.Tensor]] = None
    if validation_dataloader is not None:
        val_clean_logits = {}
        with t.inference_mode():
            for batch in validation_dataloader:
                val_clean_out = model(batch.clean)[out_slice]
                val_clean_logits[batch.key] = log_softmax(val_clean_out, dim=-1)

    src_outs: Dict[BatchKey, t.Tensor] = batch_src_ablations(
        model,
        dataloader,
        # ablation_type=AblationType.RESAMPLE,
        ablation_type=AblationType.TOKENWISE_MEAN_CORRUPT,
        # clean_corrupt="corrupt" if tree_optimisation else "clean",
    )

    val_src_outs: Optional[Dict[BatchKey, t.Tensor]] = None
    if validation_dataloader is not None:
        val_src_outs = batch_src_ablations(
            model,
            validation_dataloader,
            # ablation_type=AblationType.RESAMPLE,
            ablation_type=AblationType.TOKENWISE_MEAN_CORRUPT,
            # clean_corrupt="corrupt" if tree_optimisation else "clean",
        )
```

--------------------------------

### module_by_name(model: Any, module_name: str) -> Module

Source: https://ufo-101.github.io/auto-circuit/reference/utils/misc

Gets a module from a model by its name. The module name can be a dotted path, and the function handles wrapped models.

```APIDOC
## module_by_name(model: Any, module_name: str) -> Module

### Description
Gets a module from a model by name. The module name can be a dotted path, and the function handles wrapped models.

### Method
Function

### Signature
module_by_name(model: Any, module_name: str) -> Module

### Parameters
- **model** (Any) - Required - The model to get the module from.
- **module_name** (str) - Required - The name of the module to get.

### Returns
- **Module** - The module.
```

--------------------------------

### Clone the auto-circuit repository

Source: https://ufo-101.github.io/auto-circuit/guides/contributing

Clones the auto-circuit repository from GitHub. This is the first step in setting up the development environment.

```bash
git clone https://github.com/UFO-101/auto-circuit.git

```

--------------------------------

### Dataset Class Initialization

Source: https://ufo-101.github.io/auto-circuit/reference/data

Initializes a dataset containing clean and corrupt prompt pairs along with their corresponding correct and incorrect answers.

```python
def __init__(
    self,
    clean_prompts: List[t.Tensor] | t.Tensor,
    corrupt_prompts: List[t.Tensor] | t.Tensor,
    answers: List[t.Tensor],
    wrong_answers: List[t.Tensor],
):
    """
    A dataset of clean/corrupt prompt pairs with correct/incorrect answers.

    Args:
        clean_prompts: The 'clean' prompts. These are typically examples of the
            behavior we want to isolate where the model performs well.
            If a list, each element is a 1D prompt tensor.
            If a tensor, it should be 2D with shape (n_prompts, prompt_length).
        corrupt_prompts: The 'corrupt' prompts. These are typically similar to the
            'clean' prompts, but with some crucial difference that changes the model
            output.
            If a list, each element is a 1D prompt tensor.
            If a tensor, it should be 2D with shape (n_prompts, prompt_length).
        answers: A list of correct answers.
            Each element is a 1D tensor with the answer tokens.
        wrong_answers: A list of incorrect answers.
            Each element is a 1D tensor with the wrong answer tokens.
    """

    self.clean_prompts = clean_prompts
    self.corrupt_prompts = corrupt_prompts
    self.answers = answers
    self.wrong_answers = wrong_answers
```

--------------------------------

### docstring_node_based_official_edges

Source: https://ufo-101.github.io/auto-circuit/reference/metrics/official_circuits/circuits/docstring_official

Returns the set of edges for the attention heads in the Docstring circuit. This endpoint provides an option to distinguish token positions and a sequence start offset for KV caching.

```APIDOC
## Function: docstring_node_based_official_edges

### Description
Returns the set of edges for the attention heads in the Docstring circuit from the blog post by Stefan Heimersheim and Jett Janiak (2023). The function supports node-based patching to replicate the original measurement of the circuit's overall performance.

### Signature
```
docstring_node_based_official_edges(model: PatchableModel, token_positions: bool = False, word_idxs: Dict[str, int] = {}, seq_start_idx: int = 0) -> Set[Edge]
```

### Parameters
- **model** (`PatchableModel`) - Required - A patchable TransformerLens GPT-2 `HookedTransformer` model.
- **token_positions** (`bool`) - Optional - Whether to distinguish between token positions when returning the set of circuit edges. If `True`, the model must have `seq_len` not `None` (i.e., separate edges for each token position) and `word_idxs` must be provided. Default is `False`.
- **word_idxs** (`Dict[str, int]`) - Optional - A dictionary defining the index of specific named tokens in the circuit definition. Required tokens: `A_def`, `B_def`, `,_B`, `C_def`, `A_doc`, `B_doc`, `param_3`. Default is `{}`.
- **seq_start_idx** (`int`) - Optional - Offset to add to all token positions in `word_idxs`. Useful when using KV caching to skip the common prefix of the prompt. Default is `0`.

### Returns
- **Set[Edge]** - The set of edges in the circuit.

### Source Code
Located in `auto_circuit/metrics/official_circuits/circuits/docstring_official.py`.

```

--------------------------------

### PromptDataLoader constructor

Source: https://ufo-101.github.io/auto-circuit/reference/data

Initializes a PromptDataLoader for clean/corrupt prompt pairs. Requires prompt_dataset, seq_len, and diverge_idx. The kv_cache is only used when diverge_idx > 0. Note that drop_last=True is always passed to the parent DataLoader, ensuring all batches are the same size.

```python
def __init__(
    self,
    prompt_dataset: Any,
    seq_len: Optional[int],
    diverge_idx: int,
    kv_cache: Optional[HookedTransformerKeyValueCache] = None,
    seq_labels: Optional[List[str]] = None,
    word_idxs: Dict[str, int] = {},
    **kwargs: Any,
):
    """
    A `DataLoader` for clean/corrupt prompt pairs with correct/incorrect answers.

    Args:
        prompt_dataset: A [`PromptDataset`][auto_circuit.data.PromptDataset] with
            clean and corrupt prompts.
        seq_len: The token length of the prompts (if fixed length). This prompt
            length can be passed to `patchable_model` to enable patching specific
            token positions.
        diverge_idx: The index at which the clean and corrupt prompts diverge. (See
            [`load_datasets_from_json`][auto_circuit.data.load_datasets_from_json]
            for more information.)
        kv_cache: A cache of past key-value pairs for the transformer. Only used if
            `diverge_idx` is greater than 0. (See
            [`load_datasets_from_json`][auto_circuit.data.load_datasets_from_json]
            for more information.)
        seq_labels: A list of strings that label each token for fixed length
            prompts. Used by
            [`draw_seq_graph`][auto_circuit.visualize.draw_seq_graph] to label the
            circuit diagram.
        word_idxs: A dictionary with the token indexes of specific words. Used by
            official circuit functions.
        kwargs: Additional arguments to pass to `DataLoader`.
    """
```

--------------------------------

### Visualize the patched edges

Source: https://ufo-101.github.io/auto-circuit/guides/1%29_Getting_Started

Converts the current patch masks into prune scores and generates a sequence graph visualization. Note the final write_image line is present only in the combined example.

```python
prune_scores = model.current_patch_masks_as_prune_scores()
fig = draw_seq_graph(model, prune_scores)

```

--------------------------------

### Configure tqdm with custom defaults

Source: https://ufo-101.github.io/auto-circuit/reference/utils/custom_tqdm

This snippet defines the tqdm wrapper using functools.partial with specific default arguments for dynamic column width, bar formatting, and delay behavior.

```python
tqdm = partial(tqdm, dynamic_ncols=True, bar_format='{desc}{bar}{r_bar}', leave=None, delay=0)
```

--------------------------------

### tracr_reverse_true_edges function signature

Source: https://ufo-101.github.io/auto-circuit/reference/metrics/official_circuits/circuits/tracr/reverse_official

The function signature for tracr_reverse_true_edges. It takes a PatchableModel, a flag for token positions, a dictionary of word indices, and a sequence start index, returning a set of Edge objects.

```python
tracr_reverse_true_edges(model: PatchableModel, token_positions: bool = False, word_idxs: Dict[str, int] = {}, seq_start_idx: int = 0) -> Set[Edge]
```

--------------------------------

### Complete patching workflow end-to-end

Source: https://ufo-101.github.io/auto-circuit/guides/1%29_Getting_Started

Combines all previous steps into a single script, including loading, patching, running the model, and saving a visualization to a file.

```python
import torch as t

from auto_circuit.data import load_datasets_from_json
from auto_circuit.experiment_utils import load_tl_model
from auto_circuit.types import AblationType
from auto_circuit.utils.ablation_activations import src_ablations
from auto_circuit.utils.graph_utils import patch_mode, patchable_model
from auto_circuit.utils.misc import repo_path_to_abs_path
from auto_circuit.visualize import draw_seq_graph

device = t.device("cuda" if t.cuda.is_available() else "cpu")
model = load_tl_model("gpt2", device)

path = repo_path_to_abs_path("datasets/ioi/ioi_vanilla_template_prompts.json")
train_loader, test_loader = load_datasets_from_json(
    model=model,
    path=path,
    device=device,
    prepend_bos=True,
    batch_size=16,
    train_test_size=(128, 128),
)

model = patchable_model(
    model,
    factorized=True,
    slice_output="last_seq",
    separate_qkv=True,
    device=device,
)

ablations = src_ablations(model, test_loader, AblationType.TOKENWISE_MEAN_CORRUPT)

patch_edges = [
    "Resid Start->MLP 1",
    "MLP 1->MLP 2",
    "MLP 1->MLP 3",
    "MLP 2->A5.2.Q",
    "MLP 3->A5.2.Q",
    "A5.2->Resid End",
]

with patch_mode(model, ablations, patch_edges):
    for batch in test_loader:
        patched_out = model(batch.clean)

prune_scores = model.current_patch_masks_as_prune_scores()
fig = draw_seq_graph(model, prune_scores)
fig.write_image(repo_path_to_abs_path("docs/assets/Small_Circuit_Viz.png"), scale=4)

```

--------------------------------

### graph_edges: Get computation graph nodes and edges

Source: https://ufo-101.github.io/auto-circuit/reference/utils/graph_utils

Function signature for graph_edges, which returns a tuple containing node sets, edge sets, sequence dimension, and sequence length. Parameters include model, factorized, separate_qkv, and seq_len.

```python
graph_edges(model: Module, factorized: bool, separate_qkv: Optional[bool] = None, seq_len: Optional[int] = None) -> Tuple[Set[Node], Set[SrcNode], Set[DestNode], Dict[int | None, List[Edge]], Set[Edge], int, Optional[int]]
```

--------------------------------

### measure_task_roc function implementation

Source: https://ufo-101.github.io/auto-circuit/reference/metrics/official_circuits/measure_roc

Computes ROC curve points for given prune_scores against official_edges. Sweeps through edge counts, calculating true positive and false positive rates. Returns a list of (FPR, TPR) tuples starting with (0,0) and ending with (1,1).

```python
def measure_task_roc(
    model: PatchableModel,
    official_edges: Set[Edge],
    prune_scores: PruneScores,
    all_edges: bool = False,
) -> Measurements:
    """
    Finds points for the Receiver Operating Characteristic (ROC) curve that
    measures the performance of the given `prune_scores` at classifying which edges
    are in or out of the circuit defined by `official_edges`.

    Args:
        model: The model to measure the ROC curve for.
        official_edges: The edges that define the correct circuit.
        prune_scores: The pruning scores to measure the ROC curve for. The scores
            define an ordering of the edges in the model. We sweep through the scores
            in descending order, including the top-k edges in the circuit.
        all_edges: By default we calculate the True Positive Rate (TRP) and False
            Positive Rate (FPR) for the set of [number of edges] determined by passing
            `prune_scores` to
            [`edge_counts_util`][auto_circuit.utils.graph_utils.edge_counts_util]. If
            `all_edges` is `True`, we instead calculate the TRP and FPR for every number
            `0, 1, 2, ..., model.n_edges` of edges.

    Returns:
        A list of points that define the ROC curve. Each point is a tuple of the form
            `(FPR, TPR)`. The first point is always `(0, 0)` and the last point is
            always `(1, 1)`.
    """
    n_official = len(official_edges)
    n_complement = model.n_edges - n_official
    if all_edges:
        test_edge_counts = edge_counts_util(model.edges, test_counts=EdgeCounts.ALL)
    else:
        test_edge_counts = edge_counts_util(model.edges, prune_scores=prune_scores)
    desc_ps = desc_prune_scores(prune_scores)

    correct_ps: PruneScores = model.circuit_prune_scores(official_edges, bool=True)
    incorrect_edges: PruneScores = dict([(mod, ~ps) for mod, ps in correct_ps.items()])

    points: List[Tuple[float, float]] = [(0.0, 0.0)]
    for edge_count in (edge_count_pbar := tqdm(test_edge_counts)):
        edge_count_pbar.set_description_str(f"ROC for {edge_count} edges")
        threshold = prune_scores_threshold(desc_ps, edge_count)
        true_positives, false_positives = 0, 0
        for mod, ps in prune_scores.items():
            ps_circuit = ps.abs() >= threshold
            correct_circuit, incorrect_circuit = correct_ps[mod], incorrect_edges[mod]
            true_positives += (ps_circuit & correct_circuit).sum().item()
            false_positives += (ps_circuit & incorrect_circuit).sum().item()
        true_positive_rate = true_positives / n_official
        false_positive_rate = false_positives / n_complement
        points.append((false_positive_rate, true_positive_rate))
    points.append((1.0, 1.0))
    return points

```

--------------------------------

### PatchableModel.__init__

Source: https://ufo-101.github.io/auto-circuit/reference/utils/patchable_model

Initializes a PatchableModel instance with the provided computation graph components and configuration. This constructor sets up the model for patching and ablation experiments.

```APIDOC
## PatchableModel.__init__

### Description
Initializes a PatchableModel instance with the provided computation graph components and configuration. This constructor sets up the model for patching and ablation experiments.

### Method
Constructor

### Parameters
- **nodes** (Set[Node]) - Required - The set of all nodes in the computation graph.
- **srcs** (Set[SrcNode]) - Required - The set of all source nodes in the computation graph.
- **dests** (Set[DestNode]) - Required - The set of all destination nodes in the computation graph.
- **edge_dict** (Dict[int | None, List[Edge]]) - Required - A dictionary mapping sequence positions to the edges at that position.
- **edges** (Set[Edge]) - Required - The set of all edges in the computation graph.
- **seq_dim** (int) - Required - The sequence dimension of the model. This is the dimension on which new inputs are concatenated. In transformers, this is 1 because the activations are of shape [batch_size, seq_len, hidden_dim].
- **seq_len** (Optional[int]) - Required - The sequence length of the model inputs. If None, all token positions are simultaneously ablated.
- **wrappers** (Set[PatchWrapperImpl]) - Required - The set of all PatchWrappers in the model.
- **src_wrappers** (Set[PatchWrapperImpl]) - Required - The set of all PatchWrappers that are source nodes.
- **dest_wrappers** (Set[PatchWrapperImpl]) - Required - The set of all PatchWrappers that are destination nodes.
- **out_slice** (Tuple[slice | int, ...]) - Required - Specifies the index/slice of the output of the model to be considered for the task.
- **is_factorized** (bool) - Required - Whether the model is factorized, for Edge Ablation. Otherwise, only Node Ablation is possible.
- **is_transformer** (bool) - Required - Whether the model is a transformer.
- **separate_qkv** (Optional[bool]) - Required - Whether the model has separate query, key, and value inputs. Only used for transformers.
- **kv_caches** (Tuple[Optional[HookedTransformerKeyValueCache], ...]) - Required - A dictionary mapping batch sizes to the past key-value caches of the transformer. Only used for transformers.
- **wrapped_model** (Module) - Required - The model to wrap which is being made patchable.

### Returns
None
```

--------------------------------

### factorized_src_nodes function definition

Source: https://ufo-101.github.io/auto-circuit/reference/model_utils/transformer_lens_utils

Defines factorized_src_nodes, which returns a set of SrcNode objects for the source part of each edge in the factorized graph. It asserts model.cfg.use_attn_result and, if not attention-only, model.cfg.use_hook_mlp_in. It adds a 'Resid Start' node and then iterates over layers and heads, adding attention head result nodes and MLP output nodes when applicable.

```python
def factorized_src_nodes(model: tl.HookedTransformer) -> Set[SrcNode]:
    """
    Get the source part of each edge in the factorized graph, grouped by layer.

    Factorization introduced by
    [Elhage et al. (2021)](https://transformer-circuits.pub/2021/framework/index.html).
    See also [Molina (2023)](https://arxiv.org/pdf/2309.07315.pdf) for a good
    explanation.

    ![](../../assets/Factorized_Transformer.png)
    """
    assert model.cfg.use_attn_result  # Get attention head outputs separately
    if not model.cfg.attn_only:
        assert model.cfg.use_hook_mlp_in  # Get MLP input BEFORE layernorm
    layers, idxs = count(), count()
    nodes = set()
    nodes.add(
        SrcNode(
            name="Resid Start",
            module_name="blocks.0.hook_resid_pre",
            layer=next(layers),
            src_idx=next(idxs),
            weight="embed.W_E",
        )
    )

    for block_idx in range(model.cfg.n_layers):
        layer = next(layers)
        for head_idx in range(model.cfg.n_heads):
            nodes.add(
                SrcNode(
                    name=f"A{block_idx}.{head_idx}",
                    module_name=f"blocks.{block_idx}.attn.hook_result",
                    layer=layer,
                    src_idx=next(idxs),
                    head_dim=2,
                    head_idx=head_idx,
                    weight=f"blocks.{block_idx}.attn.W_O",
                    weight_head_dim=0,
                )
            )
        if not model.cfg.attn_only:
            nodes.add(
                SrcNode(
                    name=f"MLP {block_idx}",
                    module_name=f"blocks.{block_idx}.hook_mlp_out",
                    layer=layer if model.cfg.parallel_attn_mlp else next(layers),
                    src_idx=next(idxs),
                    weight=f"blocks.{block_idx}.mlp.W_out",
                )
            )
    return nodes
```

--------------------------------

### PromptDataset Constructor

Source: https://ufo-101.github.io/auto-circuit/reference/data

The signature for the PromptDataset class.

```python
PromptDataset(clean_prompts: List[Tensor] | Tensor, corrupt_prompts: List[Tensor] | Tensor, answers: List[Tensor], wrong_answers: List[Tensor])
```

--------------------------------

### PatchWrapperImpl.__init__

Source: https://ufo-101.github.io/auto-circuit/reference/utils/patch_wrapper

Constructor for the PatchWrapperImpl class. Initializes the wrapper with module metadata, source/destination flags, and patch configuration parameters.

```APIDOC
## PatchWrapperImpl.__init__

### Description
Initializes a PatchWrapperImpl instance that wraps a PyTorch module for activation patching and ablation. Sets up source/destination node behavior, head and sequence dimensions, and patch mask parameters.

### Method
Constructor

### Signature
PatchWrapperImpl(module_name: str, module: Module, head_dim: Optional[int] = None, seq_dim: Optional[int] = None, is_src: bool = False, src_idxs: Optional[slice] = None, is_dest: bool = False, patch_mask: Optional[Tensor] = None, in_srcs: Optional[slice] = None)

### Parameters
- **module_name** (str) - Required - Name of the wrapped module.
- **module** (Module) - Required - The module to wrap.
- **head_dim** (Optional[int]) - Optional - The dimension along which to split the heads. In TransformerLens HookedTransformers this is 2 because activations have shape [batch, seq_len, n_heads, head_dim].
- **seq_dim** (Optional[int]) - Optional - The sequence dimension of the model. This is the dimension on which new inputs are concatenated. In transformers, this is 1 because activations are of shape [batch_size, seq_len, hidden_dim].
- **is_src** (bool) - Optional - Whether the wrapped module is a SrcNode. Default: False.
- **src_idxs** (Optional[slice]) - Optional - The slice of the list of indices of SrcNodes which output from this module. Used to slice the shared curr_src_outs tensor when updating activations of the current forward pass.
- **is_dest** (bool) - Optional - Whether the wrapped module is a DestNode. Default: False.
- **patch_mask** (Optional[Tensor]) - Optional - The mask that interpolates between default activations (curr_src_outs) and ablation activations (patch_src_outs).
- **in_srcs** (Optional[slice]) - Optional - The slice of the list of indices of SrcNodes which input to this module. Used to slice the shared curr_src_outs and patch_src_outs tensors when interpolating activations of incoming edges.

### Source Code
```python
def __init__(
    self,
    module_name: str,
    module: t.nn.Module,
    head_dim: Optional[int] = None,
    seq_dim: Optional[int] = None,
    is_src: bool = False,
    src_idxs: Optional[slice] = None,
    is_dest: bool = False,
    patch_mask: Optional[t.Tensor] = None,
    in_srcs: Optional[slice] = None,
):
    super().__init__()
    self.module_name: str = module_name
    self.module: t.nn.Module = module
    self.head_dim: Optional[int] = head_dim
    self.seq_dim: Optional[int] = seq_dim
    self.curr_src_outs: Optional[t.Tensor] = None
    self.in_srcs: Optional[slice] = in_srcs

    self.is_src = is_src
    if self.is_src:
        assert src_idxs is not None
        self.src_idxs: slice = src_idxs

    self.is_dest = is_dest
    if self.is_dest:
        assert patch_mask is not None
        self.patch_mask: t.nn.Parameter = t.nn.Parameter(patch_mask)
        self.patch_src_outs: Optional[t.Tensor] = None
        self.mask_fn: MaskFn = None
        self.dropout_layer: t.nn.Module = t.nn.Dropout(p=0.0)
    self.patch_mode = False
    self.batch_size = None

    assert head_dim is None or seq_dim is None or head_dim > seq_dim
    dims = range(1, max(head_dim if head_dim else 2, seq_dim if seq_dim else 2))
    self.dims = " ".join(["seq" if i == seq_dim else f"d{i}" for i in dims])
```
```

--------------------------------

### run_constrained_prune_funcs implementation

Source: https://ufo-101.github.io/auto-circuit/reference/metrics/avoid_edges

Full implementation of run_constrained_prune_funcs. It iterates over tasks and algorithms, skips algorithms not in ['Official Circuit', 'Tree Probing'] or starting with 'Constrained', sorts edges by absolute score, and runs circuit_probing_prune_scores with avoid_edges set to the top true_edge_count edges. It registers a new constrained algorithm in PRUNE_ALGO_DICT and runs it if not already present.

```python
def run_constrained_prune_funcs(task_prune_scores: TaskPruneScores) -> TaskPruneScores:
    """
    For each task and each algorithm's [`PruneScores`][auto_circuit.types.PruneScores],
    run
    [`circuit_probing_prune_scores`][auto_circuit.prune_algos.circuit_probing.circuit_probing_prune_scores]
    with the `avoid_edges` parameter set to the top `true_edge_count` edges of the
    [`PruneScores`][auto_circuit.types.PruneScores] of the given task and algorithm.

    This is intended to test if we can find a circuit very different from the original
    solutions found by the given
    [`PruneAlgo`s][auto_circuit.prune_algos.prune_algos.PruneAlgo], that still perform
    well on the task.

    Args:
        task_prune_scores: Prune scores for each task and algorithm.

    Returns:
        A new set of prune scores for each task and algorithm, that attempts to be as
            different as possible from the original solutions found by the given
            algorithms while still performing well on the task.

    Note:
        This is an experimental function and the internal parameters used to find the
        new edges may not be well tuned.
    """
    constrained_task_prune_scores: TaskPruneScores = {}
    for task_key in (experiment_pbar := tqdm(task_prune_scores.keys())):
        task = TASK_DICT[task_key]
        experiment_pbar.set_description_str(f"Task: {task.name}")
        constrained_ps: AlgoPruneScores = {}
        algo_prune_scores = task_prune_scores[task_key]
        for algo_key, algo_ps in (prune_score_pbar := tqdm(algo_prune_scores.items())):
            if algo_key.startswith("Constrained") or algo_key not in [
                "Official Circuit",
                "Tree Probing",
            ]:
                continue
            sorted_edges: List[Edge] = list(
                sorted(algo_ps.keys(), key=lambda x: abs(algo_ps[x]), reverse=True)
            )
            algo_circuit = set([e for e in sorted_edges[: task.true_edge_count]])
            prune_score_pbar.set_description_str(f"Constrained Pruning: {algo_key}")
            constrained_algo = PruneAlgo(
                key="Constrained Circuit Probing " + algo_key,
                name=f"Not {PRUNE_ALGO_DICT[algo_key].name} Circuit Probing",
                _short_name=f"¬{PRUNE_ALGO_DICT[algo_key].short_name} TP",
                func=partial(
                    circuit_probing_prune_scores,
                    learning_rate=0.1,
                    epochs=2000,
                    regularize_lambda=0.1,
                    mask_fn="hard_concrete",
                    show_train_graph=True,
                    circuit_sizes=["true_size"],
                    tree_optimisation=True,
                    avoid_edges=algo_circuit,
                    avoid_lambda=0.3,
                ),
            )
            PRUNE_ALGO_DICT[constrained_algo.key] = constrained_algo
            if constrained_algo.key not in algo_prune_scores:
                print(f"Running {constrained_algo.name}")
                constrained_ps[constrained_algo.key] = constrained_algo.func(
                    task.model, task.train_loader, task.true_edges
                )
            else:
                print(f"Already ran {constrained_algo.name}")
        constrained_task_prune_scores[task_key] = constrained_ps
    return constrained_task_prune_scores
```

--------------------------------

### run_prompt(model, prompt, answer, top_k, prepend_bos)

Source: https://ufo-101.github.io/auto-circuit/reference/utils/misc

Helper function to run a string prompt through a TransformerLens HookedTransformer model and print the top top_k output logits.

```APIDOC
## run_prompt(model, prompt, answer, top_k, prepend_bos)

### Description
Helper function to run a string prompt through a TransformerLens `HookedTransformer` model and print the top `top_k` output logits.

### Method
Function call

### Endpoint
`run_prompt(model: Module, prompt: str, answer: Optional[str] = None, top_k: int = 10, prepend_bos: bool = False)`

### Parameters
- **model** (Module) - Required - The model to run the prompt through.
- **prompt** (str) - Required - The prompt to run through the model.
- **answer** (Optional[str]) - Optional - Token to show the rank of in the output logits. Default: `None`.
- **top_k** (int) - Optional - The number of top output logits to show. Default: `10`.
- **prepend_bos** (bool) - Optional - Whether to prepend the `BOS` token to the prompt. Default: `False`.

### Request Example
```python
run_prompt(model, "The capital of France is", answer=" Paris", top_k=10, prepend_bos=False)
```

### Response
Prints the prompt tokens and the top `top_k` output logits with their ranks.
```
Tokens:
Raw