### Setup and Imports

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/tutorial_2_0.ipynb

Installs necessary libraries and sets up the environment for the tutorial. Detects and prints the available device (MPS, CUDA, or CPU).

```python
try:
    import google.colab  # type: ignore
    from google.colab import output

    COLAB = True
    %pip install sae-lens transformer-lens sae-dashboard
except:
    COLAB = False
    from IPython import get_ipython  # type: ignore

    ipython = get_ipython()
    assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")

# Standard imports
import os
import torch
from tqdm.auto import tqdm
import plotly.express as px
import pandas as pd

# Imports for displaying vis in Colab / notebook

torch.set_grad_enabled(False)

# For the most part I'll try to import functions and classes near where they are used
# to make it clear where they come from.

if torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Device: {device}")
```

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

### Colab and IPython Environment Setup

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/test_loading_llama_8b_distill.ipynb

Sets up the environment for either Google Colab or a local IPython instance. Installs necessary packages and loads IPython extensions for autoreload.

```python
try:
    import google.colab  # type: ignore
    from google.colab import output

    COLAB = True
    %pip install sae-lens transformer-lens sae-dashboard
except:
    COLAB = False
    from IPython import get_ipython  # type: ignore

    ipython = get_ipython()
    assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")

# Standard imports
import os
import torch
from tqdm.auto import tqdm
import plotly.express as px

# Imports for displaying vis in Colab / notebook
import webbrowser
import http.server
import socketserver
import threading

PORT = 8000

torch.set_grad_enabled(False);
```

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

### Environment Setup and Imports

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/basic_loading_and_analysing.ipynb

Installs required libraries and sets up the environment for use in Colab or a local Jupyter notebook. It also imports standard libraries and modules for visualization.

```python
try:
    import google.colab  # type: ignore
    from google.colab import output

    COLAB = True
    %pip install sae-lens transformer-lens sae-dashboard
except:
    COLAB = False
    from IPython import get_ipython  # type: ignore

    ipython = get_ipython()
    assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")

# Standard imports
import os
import torch
from tqdm import tqdm
import plotly.express as px

# Imports for displaying vis in Colab / notebook
import webbrowser
import http.server
import socketserver
import threading

PORT = 8000

torch.set_grad_enabled(False)
```

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

### Setup and Device Configuration for SynthSAEBench

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/synth_sae_bench.ipynb

Initializes the environment by installing the 'sae-lens' library if in Google Colab and determines the appropriate device (GPU or CPU) for running the benchmark. It includes a warning if CUDA is not available, as a GPU is highly recommended for reasonable execution times.

```python
import warnings

import torch

try:
    import google.colab  # type: ignore

    COLAB = True
    %pip install sae-lens
except Exception:
    COLAB = False

device = "cuda"
if not torch.cuda.is_available():
    warnings.warn(
        "CUDA is not available. This notebook requires a GPU to run in a reasonable time. "
        "Training on SynthSAEBench-16k takes ~15-20 minutes on an H100 but will be "
        "extremely slow on CPU.",
        stacklevel=1,
    )
    device = "cpu"
```

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

### Setup Environment for Colab or Jupyter

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/Hooked_SAE_Transformer_Demo.ipynb

Installs the SAELens library and sets up the development environment. It includes conditional logic to handle different setups for Google Colab and local Jupyter notebooks, including autoreload functionality for development.

```python
DEVELOPMENT_MODE = False
try:
    import google.colab

    IN_COLAB = True
    print("Running as a Colab notebook")
    %pip install git+https://github.com/decoderesearch/SAELens

except:
    IN_COLAB = False
    print("Running as a Jupyter notebook - intended for development only!")
    from IPython import get_ipython

    ipython = get_ipython()
    # Code to automatically update the HookedTransformer code as its edited without restarting the kernel
    ipython.magic("load_ext autoreload")
    ipython.magic("autoreload 2")
```

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

### Setup Device and Import Utilities

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/using_an_sae_as_a_steering_vector.ipynb

Configures the computation device (MPS, CUDA, or CPU) and imports necessary utilities from PyTorch and TransformerLens.

```python
# package import
from torch import Tensor
from transformer_lens import utils
from functools import partial

# device setup
if torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Device: {device}")
```

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

### Setup Device and Environment Variables

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_a_gated_sae.ipynb

Configures the training device (GPU, MPS, or CPU) and sets up tokenization parallelism. Ensures compatibility with different hardware.

```python
import torch
import os

from sae_lens import LanguageModelSAERunnerConfig, SAETrainingRunner

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

print("Using device:", device)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
```

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

### Install SAELens

Source: https://github.com/decoderesearch/saelens/blob/main/docs/index.md

Install the SAELens library using pip. This is the first step to using the library's functionalities.

```bash
pip install sae-lens
```

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

### Setup and Imports for SAE Evaluation

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/evaluating_saes_with_sae_lens_evals.ipynb

Installs necessary libraries (sae-lens, transformer-lens, sae-dashboard) and loads extensions for interactive development. This snippet handles both Google Colab and local IPython environments.

```python
try:
    import google.colab  # type: ignore
    from google.colab import output

    COLAB = True
    %pip install sae-lens transformer-lens sae-dashboard
except:
    COLAB = False
    from IPython import get_ipython  # type: ignore

    ipython = get_ipython()
    assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")

# Standard imports
import os
import torch
from tqdm import tqdm
import pandas as pd
import plotly.express as px
import json
import numpy as np

torch.set_grad_enabled(False);
```

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

### Install Dependencies with Poetry

Source: https://github.com/decoderesearch/saelens/blob/main/CLAUDE.md

Use this command to install project dependencies. Ensure Poetry is installed and configured for the project.

```bash
poetry install
```

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

### Install Pre-commit Hook

Source: https://github.com/decoderesearch/saelens/blob/main/CLAUDE.md

Set up the pre-commit hook for the repository. This automates checks before each commit.

```bash
poetry run pre-commit install
```

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

### Install and Configure Packages

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/using_an_sae_as_a_steering_vector.ipynb

Installs necessary packages like sae-lens and transformer-lens. It also sets up IPython for local use or Colab environment, and imports modules for visualization and general utilities.

```python
try:
    # for google colab users
    import google.colab  # type: ignore
    from google.colab import output

    COLAB = True
    %pip install sae-lens transformer-lens
except:
    # for local setup
    COLAB = False
    from IPython import get_ipython  # type: ignore

    ipython = get_ipython()
    assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")

# Imports for displaying vis in Colab / notebook
import webbrowser
import http.server
import socketserver
import threading

PORT = 8000

# general imports
import os
import torch
from tqdm.auto import tqdm
import plotly.express as px

torch.set_grad_enabled(False);
```

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

### Setup and Train a Toy SAE

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_saes_on_synthetic_data.ipynb

Initializes FeatureDictionary and ActivationGenerator, configures a StandardTrainingSAE, and trains it on synthetic data, capturing snapshots of the SAE's state during training.

```python
from copy import deepcopy
from sae_lens.synthetic import FeatureDictionary, train_toy_sae, ActivationGenerator
from sae_lens import StandardTrainingSAE, StandardTrainingSAEConfig

feature_dict = FeatureDictionary(
    num_features=10,
    hidden_dim=20,
)

activation_gen = ActivationGenerator(
    num_features=feature_dict.num_features,
    firing_probabilities=0.2,
    std_firing_magnitudes=0.01,
)


# Create a standard SAE matching the feature dictionary dimensions
sae = StandardTrainingSAE(
    StandardTrainingSAEConfig(
        d_in=feature_dict.hidden_dim,
        d_sae=feature_dict.num_features,
        l1_coefficient=5e-2,
    )
)

snapshots = []

train_toy_sae(
    sae=sae,
    feature_dict=feature_dict,
    activations_generator=activation_gen,
    n_snapshots=20,
    snapshot_fn=lambda trainer: snapshots.append(
        (deepcopy(sae), f"Snapshot: train step {trainer.n_training_steps}")
    ),
)

plot_sae_feature_similarity(sae, feature_dict, reorder_features=True)
```

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

### Import Libraries and Setup Device

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/run.ipynb

Imports necessary PyTorch and SAE Lens libraries, sets up the computation device (CUDA, MPS, or CPU), and configures tokenization parallelism.

```python
import torch
import os
import sys

sys.path.append("..")

from sae_lens import (
    LanguageModelSAERunnerConfig,
    LoggingConfig,
    StandardTrainingSAEConfig,
)
from sae_lens.llm_sae_training_runner import LanguageModelSAETrainingRunner

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

print("Using device:", device)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
```

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

### Install Ansible and Collections

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Install Ansible using pip and then install necessary Ansible Galaxy collections for the project. Navigate to the Ansible scripts directory first.

```bash
pip install ansible
ansible --version

cd scripts/ansible
ansible-galaxy collection install -r util/requirements.yml
```

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

### Run Example Job with Ansible

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Execute the main Ansible playbook to set up prerequisites, run Cache Activations, and then run Train SAE jobs. Ensure your Wandb API key is exported.

```bash
cd scripts/ansible
export WANDB_API_KEY=[WANDB API key here]
ansible-playbook run-configs.yml
```

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

### Wandb API Key Setup

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Load your Wandb API key automatically by exporting it in your shell profile. Replace `[Paste Wandb API Key]` with your actual key.

```bash
export WANDB_API_KEY=[Paste Wandb API Key]
```

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

### Clone Repository and Install Dependencies

Source: https://github.com/decoderesearch/saelens/blob/main/docs/contributing.md

Clone the SAELens repository and install project dependencies using poetry. It's recommended to fork the repository before cloning for easier pull request submissions.

```bash
git clone https://github.com/decoderesearch/SAELens.git # we recommend you make a fork for submitting PR's and clone that!
poetry lock # can take a while.
poetry install
make check-ci # validate the install
```

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

### Install SAE-Lens and Dependencies

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_a_gated_sae.ipynb

Installs the required libraries for SAE training. This snippet handles both Google Colab and local environments.

```python
try:
    # import google.colab # type: ignore
    # from google.colab import output
    %pip install sae-lens transformer-lens circuitsvis
except:
    from IPython import get_ipython  # type: ignore

    ipython = get_ipython()
    assert ipython is not None
    ipython.run_line_magic("load_ext", "autoreload")
    ipython.run_line_magic("autoreload", "2")
```

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

### Instantiate ActivationsStore

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/tutorial_2_0.ipynb

Instantiates an ActivationsStore to hold activations from a dataset. Use the `from_sae` method for a convenient setup. Ensure parameters are conservative for larger models to prevent memory issues.

```python
from sae_lens import ActivationsStore

activation_store = ActivationsStore.from_sae(
    model=model,
    dataset=sae.cfg.metadata.dataset_path,
    sae=sae,
    streaming=True,
    store_batch_size_prompts=8,
    train_batch_size_tokens=4096,
    n_batches_in_buffer=32,
    device=device,
)
```

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

### Example Usage: Steering Model Generation

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/tutorial_2_0.ipynb

Demonstrates how to use the `find_max_activation` and `generate_with_steering` functions to steer model generation towards a specific SAE feature. Includes comparison with normal text generation.

```python
# Choose a feature to steer
steering_feature = steering_feature = 20115  # Choose a feature to steer towards

# Find the maximum activation for this feature
max_act = find_max_activation(model, sae, activation_store, steering_feature)
print(f"Maximum activation for feature {steering_feature}: {max_act:.4f}")

# note we could also get the max activation from Neuronpedia (https://www.neuronpedia.org/api-doc#tag/lookup/GET/api/feature/{modelId}/{layer}/{index})

# Generate text without steering for comparison
prompt = "Once upon a time"
normal_text = model.generate(
    prompt,
    max_new_tokens=95,
    stop_at_eos=False if device == "mps" else True,
    prepend_bos=sae.cfg.metadata.prepend_bos,
)

print("\nNormal text (without steering):")
print(normal_text)

# Generate text with steering
steered_text = generate_with_steering(
    model, sae, prompt, steering_feature, max_act, steering_strength=2.0
)
print("Steered text:")
print(steered_text)
```

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

### Configure Job with Shared Settings

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Copy the example configurations and modify `shared.yml` to set a unique S3 bucket name. Bucket names must be globally unique.

```bash
cd scripts/ansible
cp -r configs_example configs
```

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

### Setup Challenging Synthetic Task

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_saes_on_synthetic_data.ipynb

Configures a complex synthetic dataset with a high superposition ratio and a feature hierarchy. This setup is used to test SAE performance under more realistic conditions.

```python
# A more challenging setup: 32 features in 28 dimensions
challenge_num_features = 32
challenge_hidden_dim = 28

# Create feature dictionary (features will overlap due to superposition)
challenge_feature_dict = FeatureDictionary(
    num_features=challenge_num_features,
    hidden_dim=challenge_hidden_dim,
)

# Realistic Zipfian firing probabilities
challenge_probs = zipfian_firing_probabilities(
    num_features=challenge_num_features,
    exponent=1.2,
    max_prob=0.25,
    min_prob=0.01,
)

# Create a hierarchy: feature 0 is a parent with features 1-7 as children
# When feature 0 is inactive, all children will be deactivated
challenge_hierarchy = HierarchyNode(
    feature_index=0,
    children=[HierarchyNode(feature_index=i) for i in range(1, 8)],
)

# Create activation generator with hierarchy
challenge_gen = ActivationGenerator(
    num_features=challenge_num_features,
    firing_probabilities=challenge_probs,
    std_firing_magnitudes=0.3,
    modify_activations=hierarchy_modifier(challenge_hierarchy),
)

# Create a larger SAE
challenge_sae = StandardTrainingSAE(
    StandardTrainingSAEConfig(
        d_in=challenge_hidden_dim,
        d_sae=challenge_num_features,
        l1_coefficient=0.1,
    )
)

print(f"Challenge setup:")
print(f"  Features: {challenge_num_features}")
print(f"  Hidden dim: {challenge_hidden_dim}")
print(f"  Superposition ratio: {challenge_num_features / challenge_hidden_dim:.1f}x")
print(f"  Expected L0: {challenge_probs.sum():.2f}")
print(f"  Hierarchy: Feature 0 is parent of features 1-7")
```

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

### Steer SAE Features with API (Different Parameters)

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/tutorial_2_0.ipynb

Another example of using the /api/steer endpoint with different prompt and feature parameters. Adjust temperature and n_tokens for varied generation.

```python
import requests
import numpy as np

url = "https://www.neuronpedia.org/api/steer"

payload = {
    "prompt": 'I wrote a letter to my girlfiend. It said "',
    "modelId": "gpt2-small",
    "features": [
        {"modelId": "gpt2-small", "layer": "7-res-jb", "index": 20115, "strength": 4}
    ],
    "temperature": 0.7,
    "n_tokens": 120,
    "freq_penalty": 1,
    "seed": np.random.randint(100),
    "strength_multiplier": 4,
}
headers = {"Content-Type": "application/json"}

response = requests.post(url, json=payload, headers=headers)

print(response.json())
```

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

### Train SAE on SynthSAEBench-16k

Source: https://github.com/decoderesearch/saelens/blob/main/docs/synth_sae_bench.md

Configures and runs an SAE training session on the SynthSAEBench-16k model. This example uses recommended settings for width, training samples, learning rate, and batch size, and includes optional logging to Weights & Biases.

```python
from sae_lens.synthetic import SyntheticSAERunner, SyntheticSAERunnerConfig
from sae_lens import BatchTopKTrainingSAEConfig, LoggingConfig

runner_cfg = SyntheticSAERunnerConfig(
    # Load the pretrained benchmark model
    synthetic_model="decoderesearch/synth-sae-bench-16k-v1",

    # Configure the SAE
    sae=BatchTopKTrainingSAEConfig(
        d_in=768,
        d_sae=4096,
        k=25,
    ),

    # Training parameters
    training_samples=200_000_000,
    batch_size=1024,
    lr=3e-4,
    device="cuda",

    # Output path
    output_path="output",

    # Evaluation
    eval_frequency=1000,  # Evaluate metrics every N steps
    eval_samples=500_000,

    # Performance (recommended for modern GPUs)
    autocast_sae=True,
    autocast_data=True,

    # Optional: Logging to Weights & Biases
    logger=LoggingConfig(
        log_to_wandb=True,
        wandb_project="my_project",
        wandb_entity="my_team",  # Optional
        run_name="my-run",       # Auto-generated if not set
        wandb_log_frequency=100,  # Log metrics every N training steps
    ),
)

runner = SyntheticSAERunner(runner_cfg)
result = runner.run()

# Evaluate with ground-truth metrics
print(f"MCC: {result.final_eval.mcc:.3f}")
print(f"Explained variance: {result.final_eval.explained_variance:.3f}")
print(f"Uniqueness: {result.final_eval.uniqueness:.3f}")
print(f"F1: {result.final_eval.classification.f1_score:.3f}")
print(f"Precision: {result.final_eval.classification.precision:.3f}")
print(f"Recall: {result.final_eval.classification.recall:.3f}")
print(f"L0: {result.final_eval.sae_l0:.1f}")
print(f"Dead latents: {result.final_eval.dead_latents}")
```

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

### Run Prompt and Get Feature Activations

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/using_an_sae_as_a_steering_vector.ipynb

This snippet runs a given prompt through the model to obtain cache and then uses the SAE to calculate feature activations. It prints the tokens and the top 3 feature activations.

```python
sv_prompt = " The Golden Gate Bridge"
sv_logits, cache = model.run_with_cache(sv_prompt, prepend_bos=True)
tokens = model.to_tokens(sv_prompt)
print(tokens)

# get the feature activations from our SAE
sv_feature_acts = sae.encode(cache[hook_point])

# get sae_out
sae_out = sae.decode(sv_feature_acts)

# print out the top activations, focus on the indices
print(torch.topk(sv_feature_acts, 3))
```

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

### Display Top 10 Activating Examples

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/tutorial_2_0.ipynb

Retrieves and displays the data corresponding to the top 10 highest activations for a given feature. Ensure 'feature_acts_df', 'feature_list', and 'all_token_dfs' are available.

```python
top_10_activations = feature_acts_df.sort_values(
    f"feature_{feature_list[0]}", ascending=False
).head(10)
all_token_dfs.iloc[
    top_10_activations.index
]  # TODO: double check this is working correctly
```

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

### Setup SAETransformerBridge and Load Model

Source: https://github.com/decoderesearch/saelens/blob/main/docs/usage.md

Initialize SAETransformerBridge and load a HuggingFace model. This is necessary for models not natively supported by HookedTransformer. Ensure TransformerLens v3 is installed.

```python
from sae_lens import SAE
from sae_lens.analysis.sae_transformer_bridge import SAETransformerBridge

# Load model using TransformerBridge
model = SAETransformerBridge.boot_transformers("google/gemma-3-4b-it", device="cuda")

# Load SAE (Gemma Scope 2 SAEs work with Gemma 3 models)
sae = SAE.from_pretrained(
    release="gemma-scope-2-4b-it-res",
    sae_id="layer_17_width_16k_l0_medium",
    device="cuda"
)
```

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

### Vision Model SAE Training Example

Source: https://github.com/decoderesearch/saelens/blob/main/docs/extending_saelens.md

This code sketch illustrates the setup and training of a Sparse Autoencoder (SAE) on activations extracted from a vision model like ViT. It includes data loading, model configuration, and the SAE training loop.

```python
from dataclasses import dataclass
from collections.abc import Iterator
from typing import Any
import torch
from torchvision import transforms
from torchvision.datasets import ImageNet
from torch.utils.data import DataLoader

from sae_lens.training.sae_trainer import SAETrainer
from sae_lens.training.mixing_buffer import mixing_buffer
from sae_lens.config import SAETrainerConfig, LoggingConfig
from sae_lens.saes.sae import TrainingSAE
from sae_lens import StandardTrainingSAEConfig


@dataclass
class VisionSAEConfig:
    model_name: str = "vit_base_patch16_224"
    layer_name: str = "blocks.6"  # Target layer
    dataset_path: str = "/path/to/imagenet"
    batch_size: int = 32
    training_tokens: int = 1_000_000
    device: str = "cuda"
    output_path: str | None = None


class VisionSAERunner:
    def __init__(self, cfg: VisionSAEConfig):
        self.cfg = cfg

        # Load vision model (using timm for example)
        import timm
        self.model = timm.create_model(
            cfg.model_name,
            pretrained=True,
        ).to(cfg.device)
        self.model.eval()

        # Get the hidden dimension
        # This depends on your model - ViT base has d=768
        self.d_in = 768

        # Create SAE
        sae_cfg = StandardTrainingSAEConfig(
            d_in=self.d_in,
            d_sae=self.d_in * 8,
        )
        self.sae = TrainingSAE.from_dict(sae_cfg.to_dict()).to(cfg.device)

        # Setup data loading
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
            ),
        ])

        self.dataset = ImageNet(cfg.dataset_path, split="train", transform=transform)
        self.dataloader = DataLoader(
            self.dataset,
            batch_size=cfg.batch_size,
            shuffle=True,
            num_workers=4,
        )

    def _get_layer(self, name: str) -> torch.nn.Module:
        """Get a layer by name like 'blocks.6'."""
        parts = name.split(".")
        module = self.model
        for part in parts:
            if part.isdigit():
                module = module[int(part)]
            else:
                module = getattr(module, part)
        return module

    def _iterate_activations(self) -> Iterator[torch.Tensor]:
        """Extract activations from vision model."""
        while True:
            for images, _ in self.dataloader:
                images = images.to(self.cfg.device)
                activations: list[torch.Tensor] = []

                def hook_fn(m: Any, i: Any, o: Any) -> None:
                    activations.append(o.detach())

                layer = self._get_layer(self.cfg.layer_name)
                handle = layer.register_forward_hook(hook_fn)

                try:
                    with torch.no_grad():
                        self.model(images)
                finally:
                    handle.remove()

                # ViT outputs: (batch, num_patches + 1, hidden_dim)
                # Flatten patches: (batch * num_patches, hidden_dim)
                act = activations[0]
                flat_act = act.view(-1, act.size(-1))
                yield flat_act

    def run(self) -> TrainingSAE[Any]:
        # Use mixing buffer for shuffling patch activations
        data_provider = mixing_buffer(
            buffer_size=50_000,
            batch_size=4096,
            activations_loader=self._iterate_activations(),
        )

        trainer_cfg = SAETrainerConfig(
            total_training_samples=self.cfg.training_tokens,
            train_batch_size_samples=4096,
            device=self.cfg.device,
            lr=3e-4,
            lr_end=3e-5,
            lr_scheduler_name="constant",
            lr_warm_up_steps=1000,
            adam_beta1=0.9,
            adam_beta2=0.999,
            lr_decay_steps=0,
            n_restart_cycles=1,
            autocast=True,
            dead_feature_window=1000,
            feature_sampling_window=2000,
            n_checkpoints=0,
            checkpoint_path=None,
            save_final_checkpoint=False,
            logger=LoggingConfig(log_to_wandb=True),
        )

        trainer = SAETrainer(
            cfg=trainer_cfg,
            sae=self.sae,
            data_provider=data_provider,
        )

        trained_sae = trainer.fit()

        # Save the trained SAE
        if self.cfg.output_path is not None:
            trained_sae.save_inference_model(self.cfg.output_path)

        return trained_sae

```

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

### Serve Documentation Locally

Source: https://github.com/decoderesearch/saelens/blob/main/docs/contributing.md

Build and serve the project documentation locally using MkDocs. This allows you to preview changes to the documentation.

```bash
make docs-serve
```

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

### Install SAE-Lens and Dependencies

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_saes_on_synthetic_data.ipynb

Installs the `sae-lens` library and `kaleido` and `plotly` for visualization. This snippet checks if the code is running in Google Colab and installs the necessary packages.

```python
try:
    import google.colab  # type: ignore

    COLAB = True
    %pip install -U sae-lens kaleido==0.2.1 plotly
except Exception:
    COLAB = False
```

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

### Initialize SAELens Training Environment

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb

Sets up the PyTorch environment, imports necessary SAELens components, and configures the device (CUDA, MPS, or CPU) for training. It also disables parallel tokenization for consistency.

```python
import torch
import os

from sae_lens import (
    LanguageModelSAERunnerConfig,
    SAETrainingRunner,
    StandardTrainingSAEConfig,
    LoggingConfig,
)

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

print("Using device:", device)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
```

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

### Configure W&B and HuggingFace Credentials

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/wandb_to_hf.ipynb

Set your Weights & Biases project name, HuggingFace repository ID, and HuggingFace write access token. Ensure your HuggingFace token has write access.

```python
wandb_project_name = "YOUR-WANDB-PROJECT"
hf_repo_id = "YOUR-HF-REPO"
hf_token = "YOUR-HF-TOKEN"  # do not upload to github!
```

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

### Scatter Plot: Starts with Space vs. Starts with Capital Enrichment

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/logits_lens_with_features.ipynb

Generates a scatter plot comparing the enrichment scores of 'starts_with_space' and 'starts_with_capital'. Histograms are shown on the margins.

```python
fig = px.scatter(
    df_enrichment_scores.apply(lambda x: -1 * np.log(1 - x)).T,
    x="starts_with_space",
    y="starts_with_capital",
    marginal_x="histogram",
    marginal_y="histogram",
    labels={
        "starts_with_space": "Starts with Space",
        "starts_with_capital": "Starts with Capital",
    },
    title="Enrichment Scores for Starts with Space vs Starts with Capital",
    height=800,
    width=800,
)
# reduce point size on the scatter only
fig.update_traces(marker=dict(size=2), selector=dict(mode="markers"))
fig.show()
```

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

### Enrichment Analysis: Starts with Space vs. Starts with Capital

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/logits_lens_with_features.ipynb

Calculates and visualizes enrichment scores for token sets 'starts_with_space' and 'starts_with_capital'. It uses a manhattan plot to display the scores.

```python
# filter our list.
token_sets_index = [
    "starts_with_space",
    "starts_with_capital",
    "all_digits",
    "is_punctuation",
    "all_caps",
]
token_set_selected = {
    k: set(v) for k, v in all_token_sets.items() if k in token_sets_index
}

# calculate the enrichment scores
df_enrichment_scores = get_enrichment_df(
    dec_projection_onto_W_U,  # use the logit weight values as our rankings over tokens.
    features_ordered_by_skew,  # subset by these features
    token_set_selected,  # use token_sets
)

manhattan_plot_enrichment_scores(
    df_enrichment_scores,
    label_threshold=0,
    top_n=3,  # use our enrichment scores
).show()
```

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

### Configure and Run BatchTopK SAE Training

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/synth_sae_bench.ipynb

Sets up and initiates the training of a BatchTopK SAE using the SyntheticSAERunner. Ensure 'device' is defined before running.

```python
from sae_lens.synthetic import SyntheticSAERunner, SyntheticSAERunnerConfig
from sae_lens import BatchTopKTrainingSAEConfig, LoggingConfig

runner_cfg = SyntheticSAERunnerConfig(
    synthetic_model="decoderesearch/synth-sae-bench-16k-v1",
    sae=BatchTopKTrainingSAEConfig(
        d_in=768,
        d_sae=4096,
        k=25,
    ),
    training_samples=200_000_000,
    batch_size=1024,
    lr=3e-4,
    eval_frequency=1000,
    eval_samples=500_000,
    autocast_sae=True,
    autocast_data=True,
    logger=LoggingConfig(log_to_wandb=False),
    device=device,
)

runner = SyntheticSAERunner(runner_cfg)
btk_result = runner.run()
```

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

### Get SAE Input Shape

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/joseph_curt_pairing_gemma_scope_saes.ipynb

Determines the shape of the input tensor prepared for the SAE.

```python
sae_in.shape

```

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

### Configure and Run BatchTopK SAE Training

Source: https://github.com/decoderesearch/saelens/blob/main/docs/training_saes.md

Sets up and runs training for a BatchTopK SAE, a modern architecture that fixes mean L0 across a batch. It uses BatchTopKTrainingSAEConfig and saves as JumpReLU for inference.

```python
from sae_lens import LanguageModelSAERunnerConfig, LanguageModelSAETrainingRunner, BatchTopKTrainingSAEConfig

cfg = LanguageModelSAERunnerConfig( # Full config would be defined here
    # ... other LanguageModelSAERunnerConfig parameters ...
    sae=BatchTopKTrainingSAEConfig(
        k=100, # Set the number of active features
        d_in=1024, # Must match your hook point
        d_sae=16 * 1024,
        # ... other common SAE parameters from SAEConfig if needed ...
    ),
    # ...
)
sparse_autoencoder = LanguageModelSAETrainingRunner(cfg).run()
```

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

### Get Top Features by Skewness

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/logits_lens_with_features.ipynb

Selects the top 5000 features ordered by their skewness from the W_U_stats_df_dec DataFrame.

```python
features_ordered_by_skew = (
    W_U_stats_df_dec["skewness"].sort_values(ascending=False).head(5000).index.to_list()
)
```

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

### Load SynthSAEBench-16k Model

Source: https://github.com/decoderesearch/saelens/blob/main/docs/synth_sae_bench.md

Loads the pretrained SynthSAEBench-16k model from HuggingFace. Ensure you have the 'sae-lens' library installed and specify the desired device (e.g., 'cuda').

```python
from sae_lens.synthetic import SyntheticModel

model = SyntheticModel.from_pretrained(
    "decoderesearch/synth-sae-bench-16k-v1", device="cuda"
)
```

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

### Create Synthetic Model from Scratch

Source: https://github.com/decoderesearch/saelens/blob/main/docs/synth_sae_bench.md

Instantiate a SyntheticModel with a detailed configuration. Customize parameters like hidden dimensions, firing probabilities, hierarchy, and magnitude distributions to create specialized benchmark models.

```python
from sae_lens.synthetic import (
    SyntheticModel,
    SyntheticModelConfig,
    ZipfianFiringProbabilityConfig,
    HierarchyConfig,
    OrthogonalizationConfig,
    LowRankCorrelationConfig,
    LinearMagnitudeConfig,
    FoldedNormalMagnitudeConfig,
)

cfg = SyntheticModelConfig(
    num_features=16_384,
    hidden_dim=768,
    firing_probability=ZipfianFiringProbabilityConfig(
        exponent=0.5,
        max_prob=0.4,
        min_prob=5e-4,
    ),
    hierarchy=HierarchyConfig(
        total_root_nodes=128,
        branching_factor=4,
        max_depth=3,
        mutually_exclusive_portion=1.0,
        mutually_exclusive_min_depth=0,
        compensate_probabilities=True,
        scale_children_by_parent=True,
    ),
    orthogonalization=OrthogonalizationConfig(num_steps=100, lr=3e-4),
    correlation=LowRankCorrelationConfig(rank=25, correlation_scale=0.1),
    mean_firing_magnitudes=LinearMagnitudeConfig(start=5.0, end=4.0),
    std_firing_magnitudes=FoldedNormalMagnitudeConfig(mean=0.5, std=0.5),
    bias=0.5,
    seed=42,
)
model = SyntheticModel(cfg, device="cuda")
```

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

### Load Hugging Face Model with TransformerLens

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/test_loading_llama_8b_distill.ipynb

Loads a Hugging Face causal language model and then wraps it with TransformerLens for further analysis. Ensure you have the necessary libraries installed.

```python
from transformer_lens import HookedTransformer
from transformers import AutoModelForCausalLM

hf_model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
)

model = HookedTransformer.from_pretrained(
    "meta-llama/Llama-3.1-8B", hf_model=hf_model, device="mps"
)
```

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

### Get Neuronpedia Quick List for Features

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/using_an_sae_as_a_steering_vector.ipynb

This snippet uses the `get_neuronpedia_quick_list` function to retrieve information about the top feature activations, allowing for easy lookup on Neuronpedia.

```python
from sae_lens.analysis.neuronpedia_integration import get_neuronpedia_quick_list

get_neuronpedia_quick_list(
    sae=sae,
    features=torch.topk(sv_feature_acts, 3).indices.tolist(),
)
```

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

### Get Development Instance IP Address

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Retrieve the public IP address of the launched development instance using Ansible inventory. This is needed to SSH into the instance.

```bash
ansible-inventory --list --yaml tag_service__dev | grep public_ip_address
```

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

### List Repository Files

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/joseph_curt_pairing_gemma_scope_saes.ipynb

Fetches a list of all files within a specified Hugging Face repository. Requires the 'huggingface_hub' library.

```python
import re
import pandas as pd


from huggingface_hub import HfApi
import os


def list_repo_files(repo_id):
    api = HfApi()
    repo_files = api.list_repo_files(repo_id)
    return repo_files


files = list_repo_files(repo_id)
```

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

### Configure and Run PretokenizeRunner

Source: https://github.com/decoderesearch/saelens/blob/main/docs/training_saes.md

Configure and run the PretokenizeRunner to pre-tokenize a dataset. Adjust parameters like tokenizer_name, dataset_path, context_size, and sequence tokens based on your model and dataset. Uncomment 'hf_repo_id' or 'save_path' to upload to Hugging Face or save locally, respectively.

```python
from sae_lens import PretokenizeRunner, PretokenizeRunnerConfig

cfg = PretokenizeRunnerConfig(
    tokenizer_name="gpt2",
    dataset_path="NeelNanda/c4-10k", # this is just a tiny test dataset
    shuffle=True,
    num_proc=4, # increase this number depending on how many CPUs you have

    # tweak these settings depending on the model
    context_size=128,
    begin_batch_token="bos",
    begin_sequence_token=None,
    sequence_separator_token="eos",

    # uncomment to upload to huggingface
    # hf_repo_id="your-username/c4-10k-tokenized-gpt2"

    # uncomment to save the dataset locally
    # save_path="./c4-10k-tokenized-gpt2"
)

dataset = PretokenizeRunner(cfg).run()
```

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

### Load Pretrained SAE from Huggingface

Source: https://github.com/decoderesearch/saelens/blob/main/docs/index.md

Load a pretrained sparse autoencoder from Huggingface using its release and SAE ID. Ensure you have the necessary libraries installed for Huggingface model loading.

```python
from sae_lens import SAE

sae = SAE.from_pretrained(
    release = "gemma-scope-2b-pt-res-canonical",
    sae_id = "layer_12/width_16k/canonical",
    device = "cuda"
)
```

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

### Using SAETrainer Directly for Custom Training

Source: https://github.com/decoderesearch/saelens/blob/main/docs/extending_saelens.md

This snippet demonstrates how to use the SAETrainer class directly for custom SAE training. It covers creating a TrainingSAE, defining a custom activation data provider, configuring the trainer, and running the training loop.

```python
from sae_lens.training.sae_trainer import SAETrainer
from sae_lens.config import SAETrainerConfig, LoggingConfig
from sae_lens.saes.sae import TrainingSAE, TrainingSAEConfig
from sae_lens import StandardTrainingSAEConfig
from collections.abc import Iterator
import torch

# 1. Create a TrainingSAE with your desired architecture
sae_cfg = StandardTrainingSAEConfig(
    d_in=768,
    d_sae=768 * 8,
    l1_coefficient=5.0,
)
sae = TrainingSAE.from_dict(sae_cfg.to_dict())

# 2. Create a data provider (any iterator that yields activation tensors)
def my_activation_generator() -> Iterator[torch.Tensor]:
    while True:
        # Your logic to generate activations
        yield torch.randn(4096, 768)  # (batch_size, d_in)

data_provider = my_activation_generator()

# 3. Create the trainer config
trainer_cfg = SAETrainerConfig(
    total_training_samples=1_000_000,
    train_batch_size_samples=4096,
    device="cuda",
    lr=3e-4,
    lr_end=3e-5,
    lr_scheduler_name="constant",
    lr_warm_up_steps=1000,
    adam_beta1=0.9,
    adam_beta2=0.999,
    lr_decay_steps=0,
    n_restart_cycles=1,
    autocast=True,
    dead_feature_window=1000,
    feature_sampling_window=2000,
    n_checkpoints=0,
    checkpoint_path=None,
    save_final_checkpoint=False,
    logger=LoggingConfig(log_to_wandb=False),
)

# 4. Create and run the trainer
trainer = SAETrainer(
    cfg=trainer_cfg,
    sae=sae,
    data_provider=data_provider,
    evaluator=None,  # Optional: add custom evaluation
)

trained_sae = trainer.fit()

# Save the trained SAE
trained_sae.save_inference_model("path/to/sae")
```

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

### Configure and Train Sparse Autoencoder

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/training_a_sparse_autoencoder.ipynb

Sets up the configuration for training a Sparse Autoencoder, including model, data, SAE, and training parameters. It then instantiates and runs the training process.

```python
total_training_steps = 30_000
batch_size = 4096
total_training_tokens = total_training_steps * batch_size

lr_warm_up_steps = 0
lr_decay_steps = total_training_steps // 5
l1_warm_up_steps = total_training_steps // 20

cfg = LanguageModelSAERunnerConfig(
    # Data Generating Function (Model + Training Distibuion)
    model_name="tiny-stories-1L-21M",
    hook_name="blocks.0.hook_mlp_out",
    dataset_path="apollo-research/roneneldan-TinyStories-tokenizer-gpt2",
    is_dataset_tokenized=True,
    streaming=True,
    # SAE Parameters
    sae=StandardTrainingSAEConfig(
        d_in=1024,
        d_sae=16384,
        apply_b_dec_to_input=False,
        normalize_activations="expected_average_only_in",
        l1_coefficient=5,
        l1_warm_up_steps=l1_warm_up_steps,
    ),
    # Training Parameters
    lr=5e-5,
    adam_beta1=0.9,
    adam_beta2=0.999,
    lr_scheduler_name="constant",
    lr_warm_up_steps=lr_warm_up_steps,
    lr_decay_steps=lr_decay_steps,
    train_batch_size_tokens=batch_size,
    context_size=512,
    # Activation Store Parameters
    n_batches_in_buffer=64,
    training_tokens=total_training_tokens,
    store_batch_size_prompts=16,
    # Resampling protocol
    feature_sampling_window=1000,
    dead_feature_window=1000,
    dead_feature_threshold=1e-4,
    # WANDB
    logger=LoggingConfig(
        log_to_wandb=True,
        wandb_project="sae_lens_tutorial",
        wandb_log_frequency=30,
        eval_every_n_wandb_logs=20,
    ),
    # Misc
    device=device,
    seed=42,
    n_checkpoints=0,
    checkpoint_path="checkpoints",
    dtype="float32",
)
sparse_autoencoder = SAETrainingRunner(cfg).run()
```

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

### Get Decoder-to-W_U Projection Statistics

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/logits_lens_with_features.ipynb

Calculates statistics for the projection of SAE decoder weights onto the model's W_U matrix. Can optionally use cosine similarity for normalization.

```python
@torch.no_grad()
def get_W_U_W_dec_stats_df(
    W_dec: torch.Tensor, model: HookedTransformer, cosine_sim: bool = False
) -> tuple[pd.DataFrame, torch.Tensor]:
    W_U = model.W_U.detach().cpu()
    if cosine_sim:
        W_U = W_U / W_U.norm(dim=0, keepdim=True)
    dec_projection_onto_W_U = W_dec @ W_U
    W_U_stats_df = get_stats_df(dec_projection_onto_W_U)
    return W_U_stats_df, dec_projection_onto_W_U
```

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

### Get Forward and Backward Activations/Gradients

Source: https://github.com/decoderesearch/saelens/blob/main/tutorials/Hooked_SAE_Transformer_Demo.ipynb

A utility function to capture activations and gradients during forward and backward passes for SAE-related hooks. It requires the model, input tokens, and a metric function.

```python
from transformer_lens import ActivationCache

filter_sae_acts = lambda name: ("hook_sae_acts_post" in name)

def get_cache_fwd_and_bwd(model, tokens, metric):
    model.reset_hooks()
    cache = {}

    def forward_cache_hook(act, hook):
        cache[hook.name] = act.detach()

    model.add_hook(filter_sae_acts, forward_cache_hook, "fwd")

    grad_cache = {}

    def backward_cache_hook(act, hook):
        grad_cache[hook.name] = act.detach()

    model.add_hook(filter_sae_acts, backward_cache_hook, "bwd")

    value = metric(model(tokens))
    print(value)
    value.backward()
    model.reset_hooks()
    return (
        value.item(),
        ActivationCache(cache, model),
        ActivationCache(grad_cache, model),
    )


BASELINE = original_per_prompt_logit_diff


def ioi_metric(logits, answer_tokens=answer_tokens):
    return (
        logits_to_ave_logit_diff(logits, answer_tokens, per_prompt=True) - BASELINE
    ).sum()


clean_tokens = tokens.clone()
clean_value, clean_cache, clean_grad_cache = get_cache_fwd_and_bwd(
    model, clean_tokens, ioi_metric
)
print("Clean Value:", clean_value)
print("Clean Activations Cached:", len(clean_cache))
print("Clean Gradients Cached:", len(clean_grad_cache))
```

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

### Run Development Instance with Ansible

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Launch a development EC2 instance with SAELens pre-configured. Modify `dev.yml` for instance type and ensure configurations are copied.

```bash
cd scripts/ansible
ansible-playbook run-dev.yml
```

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

### SSH into Development Instance

Source: https://github.com/decoderesearch/saelens/blob/main/scripts/ansible/README.md

Connect to the development instance via SSH using the retrieved public IP address and a specified private key. Replace `[PASTE_PUBLIC_IP_ADDRESS]` with the actual IP.

```bash
ssh -i ~/.ssh/saelens_ansible ubuntu@[PASTE_PUBLIC_IP_ADDRESS]
```