### Setup development environment

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/contributing.md

Commands to clone the repository and install dependencies using either uv or pip.

```bash
# Clone your fork
git clone https://github.com/YOUR_USERNAME/mlxterp
cd mlxterp

# Install with all dependencies
uv sync --all-extras

# Run tests
uv run pytest

# Start development server for docs
uv run mkdocs serve
```

```bash
# Clone your fork
git clone https://github.com/YOUR_USERNAME/mlxterp
cd mlxterp

# Create virtual environment
python -m venv .venv
source .venv/bin/activate

# Install in editable mode
pip install -e ".[dev,viz]"

# Run tests
pytest
```

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

### Setup development environment with uv

Source: https://github.com/coairesearch/mlxterp/blob/main/CLAUDE.md

Recommended commands for cloning the repository and installing dependencies using uv.

```bash
# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone and setup
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Install with all dependencies
uv sync --all-extras

# Run tests
uv run pytest

# Serve documentation
uv run mkdocs serve
```

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

### Setup development environment with pip

Source: https://github.com/coairesearch/mlxterp/blob/main/CLAUDE.md

Alternative commands for installing dependencies and running services using pip.

```bash
pip install -e ".[dev,docs,viz]"
pytest
mkdocs serve
```

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

### Development setup with uv

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Commands for contributors to install dependencies, run tests, and serve documentation using uv.

```bash
# Clone your fork
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Install with all dependencies
uv sync --all-extras

# Run tests
uv run pytest

# Build docs
uv run mkdocs serve
```

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

### Development setup with pip

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Commands for contributors to install dependencies, run tests, and serve documentation using pip.

```bash
# Clone your fork
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Install in editable mode with dev dependencies
pip install -e ".[dev,viz]"

# Run tests
pytest

# Build docs
mkdocs serve
```

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

### Initialize mlxterp environment

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/tutorials/01_logit_lens.md

Setup required imports and installation instructions for the mlxterp library.

```python
# Install mlxterp if you haven't already
# pip install mlxterp

from mlxterp import InterpretableModel
import mlx.core as mx
```

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

### SAE Analysis Session Setup

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_evaluation.md

Sets up an example analysis session for a Sparse Autoencoder (SAE) using MLXTERP. Includes loading the SAE, model, tokenizer, and a test dataset.

```python
from mlx_lm import load
from mlxterp import InterpretableModel
from mlxterp.sae import BatchTopKSAE
from examples.sae_feature_analysis import (
    get_top_activating_features,
    get_top_activating_texts
)
from datasets import load_dataset

# Setup
sae = BatchTopKSAE.load("sae_layer12_mlp.mlx")
model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
interp = InterpretableModel(model, tokenizer=tokenizer)

# Load test dataset
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
texts = [item["text"] for item in dataset if len(item["text"]) > 50][:1000]

print("="*80)
print("SAE FEATURE ANALYSIS SESSION")
print("="*80)
```

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

### Install mlxterp using uv

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/index.md

Install uv if you haven't already, clone the mlxterp repository, create a virtual environment, and install dependencies using uv sync.

```bash
# Install uv if you haven't already
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone the repository
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Create environment and install
uv sync

# Activate the environment
source .venv/bin/activate
```

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

### Troubleshoot MLX installation

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Command to install the MLX framework if missing.

```bash
pip install mlx
```

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

### Install mlxterp with uv

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Steps to clone the repository, sync dependencies, and verify the installation using uv.

```bash
# Clone the repository
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Create virtual environment and install dependencies
uv sync

# Activate the environment
source .venv/bin/activate

# Verify installation
python -c "import mlxterp; print(mlxterp.__version__)"
```

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

### Build and view documentation

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/README.md

Use this command to start a local development server for the project documentation.

```bash
mkdocs serve
```

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

### Install mlxterp

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/tutorials/03_causal_tracing.md

Install the mlxterp library using pip. This is a prerequisite for running the tutorial.

```python
# Install mlxterp if you haven't already
# pip install mlxterp
```

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

### Install mlxterp via uv

Source: https://github.com/coairesearch/mlxterp/blob/main/README.md

Recommended installation method using the uv package manager.

```bash
# Install uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone and install
git clone https://github.com/coairesearch/mlxterp
cd mlxterp
uv sync

# Install mlx-lm for real models
uv add mlx-lm
```

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

### Install mlxterp from PyPI

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Commands for installing the package from PyPI once published.

```bash
pip install mlxterp

# With extras
pip install mlxterp[viz]
```

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

### Install mlxterp via pip

Source: https://github.com/coairesearch/mlxterp/blob/main/README.md

Standard installation method using pip.

```bash
git clone https://github.com/coairesearch/mlxterp
cd mlxterp
pip install -e .
pip install mlx-lm  # For loading real models
```

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

### Model Initialization Example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/activation_patching.md

Initializes the InterpretableModel for a Qwen-30B model.

```python
from mlxterp import InterpretableModel
from mlx_lm import load

base_model, tokenizer = load('mlx-community/Qwen3-30B-A3B-Thinking-2507-4bit')
model = InterpretableModel(base_model, tokenizer=tokenizer)
```

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

### Install mlxterp

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/QUICKSTART.md

Commands to clone the repository and install the package with optional visualization dependencies.

```bash
# Clone the repository
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Install in development mode
pip install -e .

# Or install with optional dependencies
pip install -e ".[viz]"
```

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

### Execute Complete Workflow

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_feature_analysis.md

Full setup and preparation steps for analyzing SAE learning.

```python
from mlxterp import InterpretableModel
from mlx_lm import load
from datasets import load_dataset

# 1. Setup
print("Loading model and SAE...")
mlx_model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
model = InterpretableModel(mlx_model, tokenizer=tokenizer)
sae = model.load_sae("sae_layer10.mlx")

# 2. Load dataset for searching
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
texts = [item["text"] for item in dataset if len(item["text"].strip()) > 50][:1000]

# 3. Pick an interesting text to analyze
test_text = "The Eiffel Tower is a famous landmark in Paris, France"
```

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

### Install mlxterp with pip

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Standard installation steps using pip, including optional dependency groups.

```bash
# Clone the repository
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Create virtual environment (optional but recommended)
python -m venv .venv
source .venv/bin/activate

# Install in development mode
pip install -e .

# Or install normally
pip install .
```

```bash
# Install with visualization tools
pip install -e ".[viz]"

# Install with development tools
pip install -e ".[dev]"

# Install everything
pip install -e ".[dev,viz]"
```

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

### Install and Login to Weights & Biases

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/dictionary_learning.md

Provides the necessary bash commands to install the Weights & Biases library and log in to your account. This is a prerequisite for using W&B features.

```bash
# Install wandb
pip install wandb

# Login (first time only)
wandb login
```

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

### Install mlx-lm

Source: https://github.com/coairesearch/mlxterp/blob/main/tests/README.md

Installs the mlx-lm package using uv. This is a prerequisite for running the tests.

```bash
uv add mlx-lm
```

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

### Install uv

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Commands to install the uv package manager on macOS/Linux or via pip.

```bash
# macOS/Linux
curl -LsSf https://astral.sh/uv/install.sh | sh

# Or with pip
pip install uv
```

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

### Run analysis scripts via CLI

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/PHASE2_COMPLETE.md

Execute the provided example scripts from the command line.

```bash
# Run basic feature analysis
python examples/phase2_feature_analysis.py

# Run Neuronpedia-style visualization
python examples/neuronpedia_style_viz.py
```

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

### Launch Jupyter Notebook

Source: https://github.com/coairesearch/mlxterp/blob/main/examples/notebooks/README.md

Starts the Jupyter notebook server for interactive exploration.

```bash
jupyter notebook
```

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

### Install mlxterp using pip

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/index.md

Clone the mlxterp repository and install it in development mode using pip.

```bash
# Clone the repository
git clone https://github.com/coairesearch/mlxterp
cd mlxterp

# Install in development mode
pip install -e .
```

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

### Manual activation patching example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/activation_patching.md

Demonstrates the manual process of tracing activations and applying interventions to measure recovery.

```python
import mlx.core as mx
from mlxterp import InterpretableModel, interventions as iv
from mlx_lm import load

# Load model
base_model, tokenizer = load('mlx-community/Llama-3.2-1B-Instruct-4bit')
model = InterpretableModel(base_model, tokenizer=tokenizer)

# Define clean vs corrupted inputs
clean_text = "Paris is the capital of France"
corrupted_text = "London is the capital of France"

# Get baseline outputs
with model.trace(clean_text):
    clean_output = model.output.save()

with model.trace(corrupted_text):
    corrupted_output = model.output.save()

mx.eval(clean_output, corrupted_output)

# Helper function to measure distance
def l2_distance(a, b):
    return float(mx.sqrt(mx.sum((a - b) ** 2)))

baseline = l2_distance(corrupted_output[0, -1], clean_output[0, -1])
print(f"Baseline L2 distance: {baseline:.2f}")

# Patch MLP at layer 10
with model.trace(clean_text) as trace:
    clean_mlp = trace.activations["model.model.layers.10.mlp"]

mx.eval(clean_mlp)

with model.trace(corrupted_text,
                interventions={"layers.10.mlp": iv.replace_with(clean_mlp)}):
    patched_output = model.output.save()

mx.eval(patched_output)

dist = l2_distance(patched_output[0, -1], clean_output[0, -1])
recovery = (baseline - dist) / baseline * 100
print(f"Layer 10 MLP: {recovery:.1f}% recovery")
```

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

### Install mlx-lm

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/JUPYTER_GUIDE.md

Install the required mlx-lm package in a Jupyter notebook environment.

```python
!uv add mlx-lm
```

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

### Implement Context Manager Pattern

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/architecture.md

Provides a clean setup and teardown lifecycle for model tracing and activation saving.

```python
class Trace:
    def __enter__(self):
        # Setup: Create and push context
        self.context = TraceContext()
        TraceContext.push(self.context)

        # Execute: Run model forward pass immediately
        # This allows users to access activations inside the with block
        self.output = self.model_forward(self.inputs)
        return self

    def __exit__(self, *args):
        # Copy saved_values from context (for values saved inside the block)
        self.saved_values = self.context.saved_values.copy()

        # Cleanup: Pop context
        TraceContext.pop()
```

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

### Initialize InterpretableModel with SAE

Source: https://github.com/coairesearch/mlxterp/blob/main/SAE_ROADMAP.md

Setup an interpretable model using a pre-trained mlx-lm model and SAE configuration.

```python
from mlx_lm import load
from mlxterp import InterpretableModel, SAEConfig

# Load model
model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
interp = InterpretableModel(model, tokenizer=tokenizer)
```

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

### Launch Jupyter Lab

Source: https://github.com/coairesearch/mlxterp/blob/main/examples/notebooks/README.md

Starts the Jupyter Lab interface for interactive exploration.

```bash
jupyter lab
```

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

### L2 Distance Usage Example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/activation_patching.md

Demonstrates activation patching using the default L2 metric.

```python
results = model.activation_patching(
    clean_text="Paris is the capital of France",
    corrupted_text="London is the capital of France",
    metric="l2"  # Default
)
```

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

### Trace model activations

Source: https://github.com/coairesearch/mlxterp/blob/main/CLAUDE.md

Example of loading a model and using the trace context manager to capture layer outputs.

```python
from mlxterp import InterpretableModel
import mlx.core as mx

# Load any MLX model with automatic wrapping
model = InterpretableModel("mlx-community/Llama-3.2-1B-Instruct")

# Simple tracing with context manager
with model.trace("The capital of France is"):
    # Direct attribute access to layer outputs
    attn_3 = model.layers[3].attn.output.save()
    mlp_8 = model.layers[8].mlp.output.save()
    logits = model.output.save()
```

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

### Intervention Examples

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/API.md

Examples of using built-in interventions like clamp and noise within a model trace.

```APIDOC
## Clamp to [-1, 1]

```python
with model.trace(input, interventions={"layers.3": iv.clamp(-1.0, 1.0)}):
    output = model.output.save()
```

## Only maximum

```python
with model.trace(input, interventions={"layers.3": iv.clamp(max_val=10.0)}):
    output = model.output.save()
```
```

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

### Perform basic feature analysis

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/PHASE2_COMPLETE.md

A complete workflow example for loading a model and SAE to analyze feature activations.

```python
from mlxterp import InterpretableModel
from mlx_lm import load

# Setup
mlx_model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
model = InterpretableModel(mlx_model, tokenizer=tokenizer)
sae = model.load_sae("sae_layer10.mlx")

# What features activate for this text?
features = model.get_top_features_for_text(
    "Machine learning processes data",
    sae=sae, layer=10, component="mlp"
)

for feat_id, act in features:
    print(f"Feature {feat_id}: {act:.3f}")
```

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

### Troubleshooting Model Loading

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/QUICKSTART.md

Example of an incorrect model loading attempt missing the quantization suffix.

```python
# WRONG - base model name doesn't exist
model = InterpretableModel("mlx-community/Llama-3.2-1B-Instruct")
```

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

### Launch Interactive Feature Dashboard

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/dictionary_learning.md

Start an interactive dashboard for exploring model features. Requires the model and dataset.

```python
# Interactive feature dashboard (Phase 2)
sae.launch_dashboard(model=model, dataset=texts)
```

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

### Define Python functions with type hints

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/contributing.md

Example of a function following Google-style docstrings and PEP 8 standards.

```python
def my_function(x: mx.array, scale: float = 1.0) -> mx.array:
    """
    Brief description of function.

    Args:
        x: Input array
        scale: Scaling factor

    Returns:
        Scaled array

    Example:
        >>> result = my_function(mx.array([1, 2, 3]), scale=2.0)
    """
    return x * scale
```

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

### MSE Distance Usage Example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/activation_patching.md

Demonstrates activation patching using the MSE metric, recommended for large vocabularies.

```python
# Recommended for Qwen (151k vocab), GPT-4 scale models
results = model.activation_patching(
    clean_text="Paris is the capital of France",
    corrupted_text="London is the capital of France",
    metric="mse"  # Most stable
)
```

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

### Load Datasets and Train SAE

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/dictionary_learning.md

Example of preparing a HuggingFace dataset for SAE training by filtering and truncating text samples.

```python
from datasets import load_dataset
from mlxterp import InterpretableModel, SAEConfig

# Load dataset
dataset = load_dataset("wikitext", "wikitext-103-raw-v1", split="train")

# Prepare texts
texts = []
for item in dataset:
    text = item['text'].strip()
    if len(text) > 50:  # Filter short texts
        texts.append(text[:512])  # Truncate long texts
    if len(texts) >= 20000:
        break

# Train SAE
model = InterpretableModel("mlx-community/Llama-3.2-1B-Instruct")
sae = model.train_sae(
    layer=10,
    dataset=texts,
    save_path="sae_wikitext_20k.mlx"
)
```

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

### Feature Naming Convention and Registry

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_evaluation.md

Demonstrates how to create and manage a feature registry for interpretable SAE features. Includes naming, description, examples, and confidence levels.

```python
# Create a feature registry
feature_registry = {
    1234: {
        "name": "mathematical_equations",
        "description": "Activates on mathematical formulas and equations",
        "examples": ["2+2=4", "E=mc²", "ax²+bx+c=0"],
        "confidence": "high"
    },
    5678: {
        "name": "geographic_locations",
        "description": "Activates on city and country names",
        "examples": ["Paris, France", "Tokyo, Japan"],
        "confidence": "high"
    },
    9999: {
        "name": "unknown_pattern",
        "description": "No clear pattern identified yet",
        "examples": [...],
        "confidence": "low"
    }
}

# Save registry
import json
with open("sae_features.json", "w") as f:
    json.dump(feature_registry, f, indent=2)
```

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

### Compute and Apply Steering Vectors

Source: https://github.com/coairesearch/mlxterp/blob/main/README.md

Derive steering vectors from contrastive examples to guide model behavior.

```python
# Compute steering vector from contrastive examples
with model.trace("I love this") as pos:
    pos_h = pos.activations['model.model.layers.10']

with model.trace("I hate this") as neg:
    neg_h = neg.activations['model.model.layers.10']

steering_vector = pos_h - neg_h

# Apply steering to guide model behavior
with model.trace("This movie is",
                 interventions={'model.model.layers.10': iv.add_vector(steering_vector)}) as steered:
    steered_output = steered.activations['__model_output__']
```

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

### Build and serve documentation

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/contributing.md

Commands to manage the documentation site locally.

```bash
# Serve locally
mkdocs serve

# Build static site
mkdocs build
```

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

### Configure Sparse Autoencoder Training

Source: https://github.com/coairesearch/mlxterp/blob/main/TROUBLESHOOTING.md

Configuration settings for training Sparse Autoencoders, ranging from rapid development testing to validated production setups.

```python
config = SAEConfig(
    sae_type="batchtopk",  # Modern BatchTopK (recommended)
    expansion_factor=4,
    k=64,
    learning_rate=3e-4,    # SAELens-validated
    num_epochs=1,          # Just 1 epoch for testing
    batch_size=64,
    text_batch_size=32,
    warmup_steps=100,
    lr_scheduler="cosine", # Cosine decay after warmup
    validation_split=0.1,
    normalize_input=True,
    use_ghost_grads=True,  # Reduce dead features
    gradient_clip=1.0,
    use_wandb=False,       # Disable W&B for testing
)
```

```python
config = SAEConfig(
    # Architecture (SAELens defaults)
    sae_type="batchtopk",      # Modern BatchTopK (recommended over "topk")
    expansion_factor=32,        # SAELens-validated (increased from 16)
    k=128,

    # Optimization
    learning_rate=3e-4,         # SAELens-validated (increased from 1e-4)
    num_epochs=3,
    batch_size=256,
    text_batch_size=32,

    # Learning rate schedule
    warmup_steps=1000,
    lr_scheduler="cosine",      # Cosine decay after warmup
    lr_decay_steps=None,        # Auto: uses total training steps

    # Sparsity warmup (reduces dead features early)
    sparsity_warm_up_steps=None,  # Auto: set to total_steps

    # Ghost gradients (critical for reducing dead features)
    use_ghost_grads=True,       # Apply ghost grads to dead features
    feature_sampling_window=1000,
    dead_feature_window=5000,

    # Other settings
    validation_split=0.05,
    normalize_input=True,
    gradient_clip=1.0,
    use_wandb=True,
    wandb_project="your-project",
    checkpoint_every=5000,
)
```

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

### Perform Clean Run and Get Baseline Output

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/tutorials/03_causal_tracing.md

Execute a clean run of the model with the factual prompt and save the output logits for baseline comparison.

```python
# Clean run - should predict "Paris"
with model.trace(clean_text) as trace:
    clean_output = model.output.save()

# Check prediction (note: model.output is logits, not hidden states)
clean_pred = get_top_from_logits(model, clean_output, top_k=3)
print("Clean prediction:", [(model.token_to_str(t), f"{s:.2f}") for t, s in clean_pred])
```

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

### Apply Steering Vectors

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/JUPYTER_GUIDE.md

Guides model behavior by adding a directional vector derived from contrasting examples to specific layers.

```python
# Compute steering vector (difference between contrasting examples)
with model.trace("I love this") as pos:
    pos_h = pos.activations['model.model.layers.10']

with model.trace("I hate this") as neg:
    neg_h = neg.activations['model.model.layers.10']

steering_vector = pos_h - neg_h

# Apply steering
with model.trace("This movie is",
                 interventions={'model.model.layers.10': iv.add_vector(steering_vector)}) as steered:
    steered_output = steered.activations['__model_output__']
```

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

### Get Activations for Batch Prompts and Single Layer

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/API.md

Collect activations from a specified layer for multiple prompts. The output shape for a single position is `(batch_size, hidden_dim)`. This example processes three prompts for layer 5.

```python
from mlxterp import get_activations

# Batch prompts
acts = get_activations(
    model,
    ["Hello", "World", "Test"],
    layers=[5],
    positions=-1
)
print(acts["layer_5"].shape)  # (3, hidden_dim)
```

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

### Get Activations for Single Prompt and Multiple Layers

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/API.md

Collect activations from specified layers for a single prompt. The output shape for a single position is `(batch_size, hidden_dim)`. This example retrieves activations for layers 3, 8, and 12.

```python
from mlxterp import get_activations

# Single prompt, multiple layers
acts = get_activations(model, "Hello world", layers=[3, 8, 12])
print(acts["layer_3"].shape)  # (1, hidden_dim)
```

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

### Improve Feature Interpretability

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_evaluation.md

Enhance the interpretability of SAE features when top activating texts seem random or lack clear patterns. This involves getting more examples, potentially with a higher top_k value, and analyzing features across different model layers.

```python
# Get more examples (patterns may emerge)
examples = get_top_activating_texts(
    sae, interp, feature_id,
    texts=dataset,
    top_k=100,  # Was 20
)

# Try different layers
for layer in [6, 12, 18, 24]:
    sae_layer = BatchTopKSAE.load(f"sae_layer{layer}.mlx")
    # Analyze features...
```

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

### Get Activations for Single Prompt and Multiple Positions

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/API.md

Collect activations from specified layers at multiple token positions for a single prompt. The output shape for multiple positions is `(batch_size, num_positions, hidden_dim)`. This example retrieves activations for the first and last tokens of layer 3.

```python
from mlxterp import get_activations

# Multiple positions
acts = get_activations(
    model,
    "Test prompt",
    layers=[3],
    positions=[0, -1]  # First and last token
)
print(acts["layer_3"].shape)  # (1, 2, hidden_dim)
```

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

### Batch Get Activations for Large Datasets

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/API.md

Process a large number of prompts efficiently using `batch_get_activations`. This function handles batching internally to manage memory. The output shape for a single position is `(total_prompts, hidden_dim)`. This example processes 1000 prompts with a batch size of 32.

```python
from mlxterp import batch_get_activations

# Process 1000 prompts efficiently
large_dataset = [f"Prompt {i}" for i in range(1000)]

acts = batch_get_activations(
    model,
    prompts=large_dataset,
    layers=[3, 8, 12],
    batch_size=32
)

print(acts["layer_3"].shape)  # (1000, hidden_dim)
```

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

### Install mlxterp Package

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/QUICKSTART.md

Resolve ModuleNotFoundError by installing the package in editable mode.

```bash
cd mlxterp
pip install -e .
```

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

### Apply Best Practices for Activation Saving

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/architecture.md

Demonstrates efficient activation saving by selecting specific layers rather than saving all activations.

```python
# ✅ Good: Only save needed activations
with model.trace(input):
    important = [3, 8, 12]
    acts = {i: model.layers[i].output.save() for i in important}

# ❌ Avoid: Saving everything
with model.trace(input):
    all_acts = [model.layers[i].output.save() for i in range(100)]
```

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

### Tutorial file organization structure

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/TUTORIAL_PAPERS_PLAN.md

Directory structure for organizing tutorial markdown files, Python scripts, and figures.

```text
examples/tutorials/
├── 01_logit_lens/
│   ├── tutorial.md
│   ├── logit_lens_tutorial.py
│   └── figures/
├── 01b_tuned_lens/
│   ├── tutorial.md
│   ├── tuned_lens_tutorial.py
│   └── figures/
├── 02_causal_tracing/
│   ├── tutorial.md
│   ├── causal_tracing_tutorial.py
│   └── figures/
├── 03_steering_vectors/
│   ├── tutorial.md
│   ├── steering_tutorial.py
│   └── figures/
├── 04_induction_heads/
│   ├── tutorial.md
│   ├── induction_heads_tutorial.py
│   └── figures/
└── 05_sparse_autoencoders/
    ├── tutorial.md
    ├── sae_tutorial.py
    └── figures/
```

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

### Configure Fast Testing with Pre-collected Activations

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/PERFORMANCE_NOTES.md

Use this configuration for smaller datasets where activations can be stored in memory to achieve high training speeds.

```python
config = SAEConfig(
    expansion_factor=16,
    k=64,
    num_epochs=1,
    batch_size=256,
    dead_feature_window=500,  # Start ghost grads early
)

# Pre-collect activations
activations = collect_activations(model, texts)  # Fast!
sae = trainer.train(activations)  # 80-200 it/s
```

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

### Large Vocabulary Model Example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/API.md

Example of using the activation patching method with a large vocabulary model.

```python
# Without correct metric - gets NaN
results = model.activation_patching(..., metric="l2")  # ❌ May overflow
```

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

### SAELens Configuration for BatchTopK

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/SAELENS_COMPARISON.md

This configuration block outlines parameters for SAELens when using the BatchTopK method. It includes settings for expansion factor, learning rate, and normalization strategies.

```python
# For BatchTopK (recommended):
expansion_factor: 32          # Larger than ours
k: 100                        # Similar to ours
learning_rate: 3e-4           # Higher than ours
normalize_activations: "expected_average_only_in"
use_ghost_grads: True
feature_sampling_window: 1000
dead_feature_window: 5000
l1_warm_up_steps: total_steps
autocast: True
dtype: "float32"
```

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

### Analyzing Texts and Features

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_evaluation.md

Performs analysis on specific texts to find top activating features and their corresponding examples. Demonstrates retrieving feature IDs, activations, and example texts.

```python
# Analysis 1: Pick interesting texts
test_texts = [
    "Python is a programming language used for data science",
    "The Eiffel Tower is located in Paris, France",
    "Photosynthesis converts sunlight into chemical energy",
]

for text in test_texts:
    print(f"\nText: '{text}'")

    # Find top features
    features = get_top_activating_features(
        sae, interp, text,
        layer=12, component="mlp",
        top_k=3
    )

    print("Top 3 features:")
    for feat_id, activation in features:
        print(f"  Feature {feat_id}: {activation:.4f}")

        # Get examples for this feature
        examples = get_top_activating_texts(
            sae, interp, feat_id,
            texts=texts[:200],  # Search subset for speed
            layer=12, component="mlp",
            top_k=3
        )

        print(f"    Top examples:")
        for ex_text, ex_act, pos in examples:
            print(f"      [{ex_act:.4f}] {ex_text[:60]}...")

print("\n" + "="*80)
print("ANALYSIS COMPLETE")
print("="*80)
```

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

### Run Causal Tracing Tutorial

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/tutorials/index.md

Execute the Python script for the Causal Tracing tutorial. This script demonstrates how to localize information within a transformer model.

```bash
python examples/tutorials/02_causal_tracing/causal_tracing_tutorial.py
```

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

### Analyze Features for Text

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_feature_analysis.md

Example usage of get_top_features_for_text to inspect model activations.

```python
# Analyze what the model "thinks about" when processing this text
features = model.get_top_features_for_text(
    "The Eiffel Tower is in Paris",
    sae=sae,
    layer=10,
    component="mlp",
    top_k=5
)

# Output:
# Feature 42: 0.856  (might represent "landmarks")
# Feature 128: 0.743  (might represent "France/French")
# Feature 91: 0.621  (might represent "locations")
```

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

### Run Logit Lens Tutorial

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/tutorials/index.md

Execute the Python script for the Logit Lens tutorial. This script demonstrates the core concept of examining intermediate predictions in transformers.

```bash
python examples/tutorials/01_logit_lens/logit_lens_tutorial.py
```

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

### Upgrade Python version

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Command to install a newer Python version via Homebrew.

```bash
# Using Homebrew
brew install python@3.11

# Or download from python.org
```

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

### Discover Available Model Components

Source: https://github.com/coairesearch/mlxterp/blob/main/TROUBLESHOOTING.md

Use this script to inspect and print the available activation keys for a specific model layer.

```python
from mlx_lm import load
from mlxterp import InterpretableModel

model, tokenizer = load("your-model-name")
interp = InterpretableModel(model, tokenizer=tokenizer)

with interp.trace("Test") as trace:
    pass

# Find components for layer 23:
for key in sorted(trace.activations.keys()):
    if "layers.23" in key:
        print(f"{key}: {trace.activations[key].shape}")
```

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

### Immediate Actions for SAE Analysis

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_evaluation.md

Lists immediate actions to take after setting up the SAE analysis, including running evaluation, analyzing features, and documenting findings.

```bash
# Run evaluation on your trained SAE:
python examples/evaluate_sae.py

# Analyze features if quality is good:
python examples/sae_feature_analysis.py

# Document findings in a feature registry
```

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

### Analyze Texts for Feature

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_feature_analysis.md

Example usage of get_top_texts_for_feature to interpret what a specific feature represents.

```python
# What does feature 1234 represent?
examples = model.get_top_texts_for_feature(
    feature_id=1234,
    sae=sae,
    texts=dataset_texts,
    layer=10,
    component="mlp",
    top_k=20
)

# Examine examples to understand the feature
for text, activation, pos in examples[:5]:
    print(f"[{activation:.3f}] {text}")

# If they all relate to "geography", then feature 1234
# likely represents geographic concepts
```

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

### Running Integration and Unit Tests

Source: https://github.com/coairesearch/mlxterp/blob/main/README.md

Execute the provided test scripts using uv to verify integration, activation validity, and intervention logic.

```bash
# Run main integration test
uv run python tests/test_comprehensive.py

# Test activation validity
uv run python tests/test_activation_validity.py

# Test interventions
uv run python tests/test_interventions.py
```

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

### Interpret Feature Interpretability

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/sae_evaluation.md

Examples of what constitutes good versus poor feature interpretability.

```python
# Good interpretability:
Feature 1234: Activates on mathematical equations
  - "2 + 2 = 4"
  - "E = mc²"
  - "ax² + bx + c = 0"

Feature 5678: Activates on geographic locations
  - "Paris, France"
  - "Tokyo, Japan"
  - "New York City"

# Poor interpretability:
Feature 9999: No clear pattern
  - Random mix of unrelated texts
  - No semantic coherence
```

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

### Cosine Distance Usage Example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/guides/activation_patching.md

Demonstrates activation patching using the cosine metric.

```python
results = model.activation_patching(
    clean_text="Paris is the capital of France",
    corrupted_text="London is the capital of France",
    metric="cosine"
)
```

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

### Resolve permission errors

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Commands to set up a virtual environment to avoid permission issues.

```bash
python -m venv .venv
source .venv/bin/activate
pip install -e .
```

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

### Verify mlxterp installation

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/installation.md

Python script to check the version and test basic model functionality.

```python
import mlxterp
import mlx.core as mx
import mlx.nn as nn

# Check version
print(f"mlxterp version: {mlxterp.__version__}")

# Test basic functionality
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(32, 32) for _ in range(3)]

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model = mlxterp.InterpretableModel(SimpleModel())
print(f"Model has {len(model.layers)} layers")
print("✓ Installation successful!")
```

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

### Initialize InterpretableModel and Define Prompts

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/tutorials/03_causal_tracing.md

Load the InterpretableModel and define clean and corrupted text prompts for causal tracing experiments.

```python
from mlxterp import InterpretableModel
from mlxterp import interventions as iv

model = InterpretableModel("mlx-community/Llama-3.2-1B-Instruct-4bit")

# Factual prompt: The model should complete with "Paris"
clean_text = "The Eiffel Tower is located in the city of"

# Corrupted prompt: Different subject, should change prediction
corrupted_text = "The Louvre Museum is located in the city of"
```

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

### Construct Dynamic Computation Graphs

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/architecture.md

Shows how computation graphs are built on-the-fly during model execution, similar to PyTorch eager mode.

```python
# Graph constructed on-the-fly
for layer in model.layers:
    x = layer(x)  # Each call adds to graph
```

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

### SAELens Sparsity Warmup Configuration

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/SAELENS_COMPARISON.md

SAELens includes configuration for gradual sparsity increase (L0/L1 warmup) over training steps. This prevents dead features early in training by allowing the network to adapt before sparsity penalties are fully enforced.

```python
l1_warm_up_steps: int = total_training_steps  # Gradual sparsity increase
l0_warm_up_steps: int = total_training_steps  # For JumpReLU
```

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

### Mechanistic interpretability example

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/JUPYTER_GUIDE.md

Perform activation patching by comparing clean and corrupted model passes.

```python
from mlxterp import InterpretableModel, interventions as iv
from mlx_lm import load
import mlx.core as mx

# 1. Load model
base_model, tokenizer = load('mlx-community/Llama-3.2-1B-Instruct-4bit')
model = InterpretableModel(base_model, tokenizer=tokenizer)

# 2. Run clean forward pass
prompt = "The Eiffel Tower is located in"
with model.trace(prompt) as clean:
    clean_attn_8 = clean.activations['model.model.layers.8.self_attn']
    clean_output = clean.activations['__model_output__']

# 3. Run with corrupted input
corrupted_prompt = "The Big Ben is located in"
with model.trace(corrupted_prompt) as corrupted:
    corrupted_attn_8 = corrupted.activations['model.model.layers.8.self_attn']

# 4. Activation patching: restore clean attention in layer 8
from mlxterp import interventions as iv
```

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

### Initialize InterpretableModel with Default Support

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/examples.md

Initialize InterpretableModel for standard models like Llama and Mistral. Auto-detection handles common architectures without explicit configuration.

```python
from mlxterp import InterpretableModel
from mlx_lm import load

# Auto-detection works for standard models
base_model, tokenizer = load("mlx-community/Llama-3.2-1B-Instruct-4bit")
model = InterpretableModel(base_model, tokenizer=tokenizer)

# logit_lens and get_token_predictions work automatically
results = model.logit_lens("Hello world")
```

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

### Logit Lens Prediction Output

Source: https://github.com/coairesearch/mlxterp/blob/main/tests/README.md

Example output showing logit lens predictions at different layers.

```text
Layer  0: ' the'
Layer  5: ' a'
Layer 10: ' Paris'
Layer 15: ' Paris'
```

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

### Analyze top texts for a feature

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/PHASE2_COMPLETE.md

Retrieve text examples where a specific feature exhibits strong activation.

```python
examples = model.get_top_texts_for_feature(
    feature_id=1234,
    sae=sae,
    texts=dataset,
    layer=10, component="mlp", top_k=20
)
# Returns: [(text, activation, position), ...]
```

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

### Access model outputs

Source: https://github.com/coairesearch/mlxterp/blob/main/examples/notebooks/Test Library.ipynb

Demonstrates how to retrieve outputs from the wrapped model versus the base model directly.

```python
output = model.model(input_data)  # Access the underlying base_model
print(f"Direct output: {output}")

# Option 3: Use base_model directly
output = base_model(input_data)
print(f"Base model output: {output.shape}")
```

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

### Trace Model Execution

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/QUICKSTART.md

Example of wrapping a custom model with InterpretableModel and tracing specific layer outputs.

```python
from mlxterp import InterpretableModel
import mlx.core as mx
import mlx.nn as nn

# 1. Create a simple model
class SimpleTransformer(nn.Module):
    def __init__(self, hidden_dim=64):
        super().__init__()
        self.layers = [
            nn.Linear(hidden_dim, hidden_dim)
            for _ in range(4)
        ]

    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
            x = nn.relu(x)
        return x

# 2. Wrap with InterpretableModel
base_model = SimpleTransformer()
model = InterpretableModel(base_model)

# 3. Create input
input_data = mx.random.normal((1, 10, 64))  # (batch, seq, hidden)

# 4. Trace execution
with model.trace(input_data):
    layer_0 = model.layers[0].output.save()
    layer_2 = model.layers[2].output.save()

print(f"Layer 0 shape: {layer_0.shape}")
print(f"Layer 2 shape: {layer_2.shape}")
```

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

### Run Neuronpedia Style Visualization Demo

Source: https://github.com/coairesearch/mlxterp/blob/main/docs/dev/PHASE2_COMPLETE.md

Execute the visualization demo script to observe Neuronpedia-style visualizations. This script requires a trained SAE.

```python
python examples/neuronpedia_style_viz.py
```