### Run Basic Example Experiment

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Execute a simple two-layer transformer training on multi-query associative recall using the basic example script.

```bash
python -m zoology.launch zoology/experiments/basic_examples/basic.py
```

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

### Install Zoology Project (Lightweight)

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Install the Zoology project with only the essential dependencies, excluding optional ones like mamba_ssm and conv1d.

```bash
pip install -e .
```

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

### Install Zoology Project

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Clone the repository and install the Zoology project with extra and analysis dependencies. Ensure torch and transformers are installed beforehand.

```bash
git clone https://github.com/HazyResearch/zoology.git
cd zoology
pip install -e .[extra,analysis]
```

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

### Run MQAR Example Config Sweep

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Launch a comprehensive sweep of models for Multi-Query Associative Recall (MQAR) using the original MQAR configuration file.

```bash
python -m zoology.launch zoology/experiments/mqar_example_configs/original_mqar_configs.py -p
```

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

### Launch Hyperparameter Sweeps

Source: https://context7.com/hazyresearch/zoology/llms.txt

Provides command-line examples for launching hyperparameter sweeps sequentially or in parallel using Ray, with options to specify GPUs.

```bash
# Run sequentially
python -m zoology.launch sweep_experiment.py

# Run in parallel across GPUs (requires Ray)
python -m zoology.launch sweep_experiment.py -p

# Specify GPUs
python -m zoology.launch sweep_experiment.py -p --gpus 0,1,2,3
```

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

### Run Basic Example Experiment with Learning Rate Sweep

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Train a basic transformer model while sweeping over different learning rates using the basic sweep script.

```bash
python -m zoology.launch zoology/experiments/basic_examples/basic_sweep.py
```

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

### Basic TrainConfig Setup for MQAR

Source: https://context7.com/hazyresearch/zoology/llms.txt

Defines a comprehensive training configuration including data, model, and hyperparameters for the MQAR task. Ensure to update cache_dir and wandb project/entity details.

```python
from zoology.config import TrainConfig, ModelConfig, DataConfig, ModuleConfig, LoggerConfig
from zoology.data.multiquery_ar import MQARConfig

# Basic training configuration
config = TrainConfig(
    # Data configuration
    data=DataConfig(
        train_configs=[
            MQARConfig(
                num_examples=10_000,
                vocab_size=256,
                input_seq_len=64,
                num_kv_pairs=4
            )
        ],
        test_configs=[
            MQARConfig(
                num_examples=1_000,
                vocab_size=256,
                input_seq_len=64,
                num_kv_pairs=4
            )
        ],
        batch_size=32,  # Can also be tuple (train_bs, test_bs)
        cache_dir="/path/to/cache"  # Optional caching
    ),

    # Model configuration
    model=ModelConfig(
        vocab_size=256,
        d_model=128,
        n_layers=2,
        max_position_embeddings=64,
        sequence_mixer=ModuleConfig(
            name="zoology.mixers.attention.MHA",
            kwargs={"dropout": 0.1, "num_heads": 2}
        )
    ),

    # Training parameters
    max_epochs=100,
    learning_rate=1e-3,
    weight_decay=0.1,
    early_stopping_metric="valid/accuracy",
    early_stopping_threshold=0.99,

    # Logging
    logger=LoggerConfig(
        project_name="my_wandb_project",
        entity="my_wandb_entity"
    )
)

configs = [config]  # List of configs for sweeps
```

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

### Define Experiment Configuration for Launching

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Define a configuration object in a Python file and assign it to a global variable named `configs` to launch a single experiment. This setup is used by the `zoology.launch` command.

```python
config = TrainConfig(...)
configs = [config]
```

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

### Import necessary libraries

Source: https://github.com/hazyresearch/zoology/blob/main/zoology/experiments/paper_configs/arxiv24_based_rebuttal_circuits/plot.ipynb

Imports required libraries for data manipulation, plotting, and fetching Weights & Biases runs. Ensure these libraries are installed.

```python
import os

import pandas as pd
from tqdm import tqdm
import seaborn as sns
import matplotlib.pyplot as plt

from zoology.analysis.utils import fetch_wandb_runs
```

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

### Instantiate FunctionConfig

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Use FunctionConfig to configure and instantiate Python functions, useful for dynamic function calls within your experiment setup. The `instantiate()` method handles the import and partial application of arguments.

```python
fn_config = FunctionConfig(name="torch.sort", kwargs={"descending": True})
fn = fn_config.instantiate()
fn(torch.tensor([2,4,3])) # [4, 3, 2]
```

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

### Get Next Token Predictions with Argmax

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Apply a language modeling head to get next token predictions by taking the argmax of the output tensor. This is equivalent to finding the index with the highest value, corresponding to the predicted next token.

```python
# do not modify
z2 = torch.argmax(O2, dim=-1)
```

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

### Configure GatedLinearAttention (GLA) Model

Source: https://context7.com/hazyresearch/zoology/llms.txt

Configure a hardware-efficient Gated Linear Attention (GLA) model with learnable decay gates. This example shows a 'chunk' mode configuration with options for short convolutions and output gates. Adjust parameters like 'num_heads', 'expand_k', 'expand_v', and 'gate_fn' as needed.

```python
from zoology.config import ModelConfig, ModuleConfig

gla_mixer = ModuleConfig(
    name="zoology.mixers.gla.GatedLinearAttention",
    kwargs={
        "mode": "chunk",           # "chunk", "fused_recurrent", "fused_chunk"
        "num_heads": 4,
        "expand_k": 0.5,           # Key dimension expansion
        "expand_v": 1.0,           # Value dimension expansion
        "use_short_conv": True,    # Add short convolution
        "conv_size": 4,
        "use_output_gate": True,
        "gate_fn": "swish"
    }
)

model_config = ModelConfig(
    vocab_size=8192,
    d_model=128,
    n_layers=2,
    max_position_embeddings=0,
    sequence_mixer=gla_mixer
)
```

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

### Configure Based (Linear Attention) Model

Source: https://context7.com/hazyresearch/zoology/llms.txt

Configure a Based model, which implements linear attention with learnable feature maps for efficient O(n) inference. This example shows a hybrid mixer combining a short convolution with the Based module. The 'train_view' can be set to 'linear' or 'quadratic'.

```python
from zoology.config import ModelConfig, ModuleConfig

# Configure Based model with short convolution
conv_mixer = ModuleConfig(
    name="zoology.mixers.base_conv.BaseConv",
    kwargs={"l_max": 1024, "kernel_size": 3, "implicit_long_conv": True}
)

based_mixer = ModuleConfig(
    name="zoology.mixers.based.Based",
    kwargs={
        "l_max": 1024,
        "feature_dim": 16,          # Feature map expansion dimension
        "feature_name": "taylor_exp", # Feature map type
        "num_key_value_heads": 1,
        "num_heads": 1,
        "train_view": "quadratic"   # "linear" or "quadratic" training
    }
)

# Hybrid mixer combining convolution + linear attention
model_config = ModelConfig(
    vocab_size=8192,
    d_model=128,
    n_layers=2,
    max_position_embeddings=0,  # No position embeddings needed
    sequence_mixer=ModuleConfig(
        name="zoology.mixers.hybrid.Hybrid",
        kwargs={"configs": [conv_mixer.dict(), based_mixer.dict()]}
    )
)
```

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

### Define Query Projection Matrix (Layer 1)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Defines the weights for the Query (Q) projection matrix in the first attention layer. This setup shifts positional embeddings by one position, preparing them for attention calculations.

```python
Q[1:position_dim+1, :position_dim] = 100  * # YOUR ANSWER TENSOR
```

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

### Call and Display Incorrect Indices

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Calls the score_solution function and displays the first 10 incorrect indices. This is useful for debugging and identifying specific examples that were not solved correctly.

```python
incorrect_indices = score_solution(repeated, z2, golds)
incorrect_indices[:10]
```

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

### Visualize Intermediate Matrices in Associative Recall

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Visualizes various intermediate matrices (q1, k1, v1, S1, O1, z1, q2, k2, v2, S2, O2, z2) generated during the associative recall task processing. This helps in understanding the transformations applied to the data at each step. The code prints the input example and then displays each matrix using matplotlib.

```python
idx = 0

print("Input")
print(examples[idx])

print(f"Below we are going to print a bunch of N x D matrices for the solution. Note that D is of length (position_dim + embed_dim)\n")

print("q1 -- there are ones on the position encodings side of the matrix, shifted by one position")
plt.imshow(q1[idx].detach().numpy())
plt.show()
print("k1 -- there are ones on the position encodings side of the matrix, corresponding to the position encodings in the original sequence")
plt.imshow(k1[idx].detach().numpy())
plt.show()
print("v1 -- we take the original token embeddings in the sequence")
plt.imshow(v1[idx].detach().numpy())
plt.show()
print("S1 -- there are high values on the indices i, j where j comes before a particular token i in the sequence")
plt.imshow(S1[idx].detach().numpy())
plt.show()
print("O1 -- we have shifted each token embedding forward by one position. the old token 1 is now at position 2, etc.")
plt.imshow(O1[idx].detach().numpy())
plt.show()
print("z1 -- it's hard to visualize because of the relative values, but the original embeddings are now added back")
plt.imshow(z1[idx].detach().numpy())
plt.show()

print("q2 -- remember we said z1 has the original embeddings added back. q2 just takes the original token embeddings from z1 (right hand side)")
plt.imshow(q2[idx].detach().numpy())
plt.show()
print("k2 -- k2 takes the embeddings shifted forward by one position (left hand side of z1)")
plt.imshow(k2[idx].detach().numpy())
plt.show()
print("v2 -- v2 takes the original token embeddings from z1 as well")
plt.imshow(v2[idx].detach().numpy())
plt.show()
print("S2 -- S2 has high values at (i, j) if token i matches token j-1 in the original sequence")
plt.imshow(S2[idx].detach().numpy())
plt.show()
print("O2 -- O2 has the associative recall output tokens")
plt.imshow(O2[idx].detach().numpy())
plt.show()

print("z2 before subtracting position dim -- we take the argmax of O2 to get the token embedding with the 'max' value")
z2 = torch.argmax(O2, dim=-1)
print(z2[0])
print(f"z2 after subtracting position dim ({position_dim})")
z2 = z2 - position_dim
print(z2[0])
```

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

### Launch Single Experiment

Source: https://context7.com/hazyresearch/zoology/llms.txt

Demonstrates how to create an experiment file and launch a single training configuration from the command line.

```bash
# Create experiment file: my_experiment.py
# Contents:
# from zoology.config import TrainConfig, ModelConfig, DataConfig, ModuleConfig
# from zoology.data.multiquery_ar import MQARConfig
# 
# config = TrainConfig(
#     data=DataConfig(
#         train_configs=[MQARConfig(num_examples=10_000, vocab_size=256, input_seq_len=64)],
#         test_configs=[MQARConfig(num_examples=1_000, vocab_size=256, input_seq_len=64)]
#     ),
#     model=ModelConfig(
#         vocab_size=256,
#         sequence_mixer=ModuleConfig(name="zoology.mixers.attention.MHA", kwargs={"num_heads": 2})
#     )
# )
# configs = [config]

# Launch from command line
python -m zoology.launch my_experiment.py
```

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

### Dynamic Module and Function Instantiation

Source: https://context7.com/hazyresearch/zoology/llms.txt

Demonstrates using ModuleConfig for PyTorch modules and FunctionConfig for partial functions, enabling runtime instantiation with specified arguments.

```python
from zoology.config import ModuleConfig, FunctionConfig
import torch

# ModuleConfig for PyTorch modules
attention_config = ModuleConfig(
    name="zoology.mixers.attention.MHA",
    kwargs={"dropout": 0.1, "num_heads": 2}
)
# Instantiate with additional kwargs
attention = attention_config.instantiate(d_model=128, layer_idx=0)

# FunctionConfig for partial functions
fn_config = FunctionConfig(
    name="torch.sort",
    kwargs={"descending": True}
)
fn = fn_config.instantiate()
result = fn(torch.tensor([2, 4, 3]))  # Returns [4, 3, 2]
```

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

### Run Paper Experiment

Source: https://github.com/hazyresearch/zoology/blob/main/zoology/experiments/paper_configs/arxiv24_based_figure2/README.md

Execute the experiment configuration to reproduce paper results. Ensure WandB is set up for logging.

```bash
python -m zoology.launch zoology/experiments/arxiv24_based_figure2/configs.py -p
```

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

### Configuring Models from Repository

Source: https://context7.com/hazyresearch/zoology/llms.txt

This code demonstrates how to use the model repository to pre-configure various language model architectures. It sets up common configurations and then adds specific models like attention, Based, Mamba2, GLA, and Gated Delta Net. The models can be filtered by name.

```python
from zoology.config import ModelConfig, ModuleConfig
from zoology.experiments.models_repo import (
    add_attention, add_based, add_mamba2, add_gla,
    add_delta_net, add_gated_delta_net, add_rwkv7, add_hyena
)

# Common configuration
input_seq_len = 1024
model_factory_kwargs = {
    "state_mixer": dict(name="torch.nn.Identity", kwargs={}),
    "vocab_size": 8192
}

# Short convolution mixer (used in hybrid architectures)
conv_mixer = dict(
    name="zoology.mixers.base_conv.BaseConv",
    kwargs={"l_max": input_seq_len, "kernel_size": 3, "implicit_long_conv": True}
)

# Build model list
models = []
models = add_attention(models, conv_mixer, input_seq_len, model_factory_kwargs)
models = add_based(models, conv_mixer, input_seq_len, model_factory_kwargs)
models = add_mamba2(models, conv_mixer, input_seq_len, model_factory_kwargs)
models = add_gla(models, conv_mixer, input_seq_len, model_factory_kwargs)
models = add_gated_delta_net(models, conv_mixer, input_seq_len, model_factory_kwargs)

# Filter to specific architectures
included = ["attention", "based", "gated_delta_net"]
models = [m for m in models if any(i in m.name for i in included)]
```

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

### Instantiate TrainConfig with Data and Model Configurations

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Configure an experiment's training parameters, data sources, and model architecture using nested Pydantic models. Ensure all necessary configuration classes are imported.

```python
from zoology.config import TrainConfig, ModelConfig, DataConfig, ModuleConfig, FunctionConfig
from zoology.data.associative_recall import MQARConfig

input_seq_len = 10

config = TrainConfig(
    max_epochs=20,
    data=DataConfig(
        train_configs=[MQARConfig(num_examples=10_000, vocab_size=128, input_seq_len=input_seq_len, )], #**factory_kwargs
        test_configs=[MQARConfig(num_examples=1_000, vocab_size=128, input_seq_len=input_seq_len, )], #**factory_kwargs
    ),
    model=ModelConfig(
        vocab_size=128,
        sequence_mixer=ModuleConfig(name= "zoology.mixers.attention.MHA") ),
)
```

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

### Creating Training Configurations

Source: https://context7.com/hazyresearch/zoology/llms.txt

This code generates training configurations for multiple models and learning rates. It iterates through the selected models and a range of learning rates, creating `TrainConfig` objects with specified data configurations (MQARConfig) and hyperparameters. Each configuration is assigned a unique run ID.

```python
from zoology.config import TrainConfig, DataConfig
from zoology.data.multiquery_ar import MQARConfig
import numpy as np

configs = []
for model in models:
    for lr in np.logspace(-3, -1.5, 4):
        config = TrainConfig(
            model=model,
            data=DataConfig(
                train_configs=[MQARConfig(vocab_size=8192, input_seq_len=64, num_examples=100_000, num_kv_pairs=4)],
                test_configs=[MQARConfig(vocab_size=8192, input_seq_len=64, num_examples=1_000, num_kv_pairs=4)]
            ),
            learning_rate=lr,
            max_epochs=32,
            run_id=f"{model.name}-lr{lr:.1e}"
        )
        configs.append(config)
```

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

### Configure Logger for Weights and Biases

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Set up logging for your experiments by configuring `LoggerConfig` within your `TrainConfig`. Specify your Weights and Biases project and entity names.

```python
from zoology.config import TrainConfig, LoggerConfig

TrainConfig(
    logger=LoggerConfig(
        project="my_wandb_project",
        entity="my_wandb_entity",
    ),
    ...
)
```

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

### Initialize Attention Projection Matrices (Layer 2)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Initializes the projection matrices for the second attention layer. K2 is initialized to zeros, while Q2 and V2 are initialized as identity matrices.

```python
K2 = torch.zeros(d_model, d_model)
Q2 = torch.eye(d_model) 
V2 = torch.eye(d_model)
```

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

### Fetch Weights & Biases Results

Source: https://context7.com/hazyresearch/zoology/llms.txt

Demonstrates fetching experiment results from Weights & Biases using `fetch_wandb_runs` and processing them with pandas for analysis.

```python
from zoology.analysis.utils import fetch_wandb_runs
import pandas as pd

# Fetch runs by launch_id (shared across sweep configs)
df = fetch_wandb_runs(
    launch_id=[
        "default-2024-02-09-05-44-06",
        "default-2024-02-09-14-59-58"
    ],
    project_name="zoology"
)

# DataFrame contains columns like:
# - model.name, model.d_model, model.n_layers
# - learning_rate, valid/accuracy, valid/loss
# - state_size (computed from model architecture)

# Filter best runs per model and state size
idx = df.groupby(["state_size", "model.name"])["valid/accuracy"].idxmax(skipna=True).dropna()
best_runs = df.loc[idx]
print(best_runs[["model.name", "state_size", "valid/accuracy"]])
```

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

### Initialize Attention Projection Matrices (Layer 1)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Initializes the tensors for Query (Q), Key (K), and Value (V) projection matrices used in the first layer of attention. K is initialized as an identity matrix, while Q and V are zero matrices of the same dimension.

```python
K = torch.eye(d_model)
Q = torch.zeros_like(K)
V = torch.zeros_like(K)
```

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

### Fetch and plot single run

Source: https://github.com/hazyresearch/zoology/blob/main/zoology/experiments/paper_configs/arxiv24_based_rebuttal_circuits/plot.ipynb

Fetches a single Weights & Biases run, filters for finished states, identifies the optimal configuration based on validation accuracy, and plots the results.

```python
df = fetch_wandb_runs(
    launch_id=[
        # "redo-vocab-majority-sweep-V654618ae"
        "default-2024-03-29-02-15-15"

        
    ], 
    project_name="zoology"
)

df = df[df["state"] == "finished"]

idxs = df.groupby(
    ["model.name", "data.train_configs.0.input_seq_len"]
){"valid/accuracy"}.idxmax()
plot_df = df.loc[idxs)

sns.set_style("whitegrid")
sns.relplot(
    data=plot_df,
    hue="model.name",
    x="data.train_configs.0.input_seq_len",
    y="valid/accuracy",
    kind="line",
    marker="o"
)

plt.ylim(0.5, 1.0)
```

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

### MQARConfig for Data Generation

Source: https://context7.com/hazyresearch/zoology/llms.txt

Configures and builds the Multi-Query Associative Recall (MQAR) dataset, allowing customization of vocabulary size, sequence length, and key-value pair count.

```python
from zoology.data.multiquery_ar import MQARConfig, multiquery_ar
from zoology.data.utils import DataSegment

# Configure MQAR task
mqar_config = MQARConfig(
    vocab_size=8192,      # Large vocab helps differentiate architectures
    num_examples=10_000,  # Number of examples to generate
    input_seq_len=64,     # Sequence length
    num_kv_pairs=8,       # Number of key-value pairs per sequence
    power_a=0.01,         # Power law parameter for query gap distribution
    random_non_queries=True  # Fill non-query positions with random tokens
)

# Build the dataset
data_segment: DataSegment = mqar_config.build(seed=42)

# Access data
inputs = data_segment.inputs   # Shape: (num_examples, input_seq_len)
labels = data_segment.labels   # Shape: (num_examples, input_seq_len), -100 for ignored positions

# Direct function call for more control
data = multiquery_ar(
    vocab_size=256,
    num_examples=1000,
    input_seq_len=64,
    seed=42,
    num_kv_pairs=4,
    power_a=0.01,
    random_non_queries=True
)
# Example sequence:
# Inputs: [Key1, Val1, Key2, Val2, ..., Query(Key1), ..., Query(Key2), ...]
# Labels: [-100, -100, -100, -100, ..., Val1,        ..., Val2,        ...]
```

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

### Load and Inspect Compositional MQAR Data

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/test_data.ipynb

Imports the `compositional_mqar` dataset and iterates through the first input-label pair, printing values up to the 30th element. Ensure the 'compositional_ar' module is accessible.

```python
# Step 2: Import the data you want to inspect

from compositional_ar import compositional_mqar
data  = compositional_mqar(
    vocab_size=256,
    num_examples=10,
    input_seq_len=24,
    num_kv_pairs=4,
    seed=0,
)
x= data.inputs[0]
y= data.labels[0]

for i, (_x, _y) in enumerate(zip(x, y)):
    if i == 30:
        break
    print(f"{_x.item()}, {_y.item()}")
```

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

### Configure Mamba/Mamba2 (Selective State Space Models)

Source: https://context7.com/hazyresearch/zoology/llms.txt

Configure Mamba and Mamba2 models, which are state space models featuring input-dependent selection mechanisms. Mamba uses 'MambaBlock' and requires specific mixer arguments like 'd_state', 'd_conv', and 'expand'. Mamba2 uses 'Mamba2Block' with a simpler configuration.

```python
from zoology.config import ModelConfig, ModuleConfig

# Mamba configuration (uses special MambaBlock)
mamba_config = ModelConfig(
    vocab_size=8192,
    d_model=128,
    n_layers=2,
    max_position_embeddings=0,
    block_type="MambaBlock",  # Special block type for Mamba
    sequence_mixer=ModuleConfig(
        name="zoology.mixers.mamba.Mamba",
        kwargs={
            "d_state": 16,    # SSM state dimension
            "d_conv": 4,      # Local convolution width
            "expand": 2       # Inner dimension expansion factor
        }
    )
)

# Mamba2 configuration
mamba2_config = ModelConfig(
    vocab_size=8192,
    d_model=128,
    n_layers=2,
    max_position_embeddings=0,
    block_type="Mamba2Block",
    sequence_mixer=ModuleConfig(
        name="zoology.mixers.mamba2.Mamba2",
        kwargs={"d_state": 16}
    )
)
```

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

### Create Recall-Memory Tradeoff Plots

Source: https://context7.com/hazyresearch/zoology/llms.txt

Illustrates how to generate plots comparing model architectures by fetching data from Weights & Biases and using seaborn and matplotlib.

```python
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# Assuming df is fetched from WandB
df = fetch_wandb_runs(launch_id=["your-sweep-id"], project_name="your-project")

```

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

### Train Model Programmatically

Source: https://context7.com/hazyresearch/zoology/llms.txt

Shows how to use the `train` function directly within a Python script to control model training programmatically.

```python
from zoology.train import train
from zoology.config import TrainConfig, ModelConfig, DataConfig, ModuleConfig
from zoology.data.multiquery_ar import MQARConfig

config = TrainConfig(
    data=DataConfig(
        train_configs=[MQARConfig(num_examples=10_000, vocab_size=256, input_seq_len=64, num_kv_pairs=4)],
        test_configs=[MQARConfig(num_examples=1_000, vocab_size=256, input_seq_len=64, num_kv_pairs=4)]
    ),
    model=ModelConfig(
        vocab_size=256,
        d_model=128,
        n_layers=2,
        sequence_mixer=ModuleConfig(name="zoology.mixers.attention.MHA", kwargs={"num_heads": 2})
    ),
    max_epochs=50,
    learning_rate=1e-3
)

# Run training
train(config)
```

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

### Load and Inspect Parity Dataset

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/test_data.ipynb

Loads the `parity` dataset with a specified input sequence length and prints the lengths of the input and label sequences. It then iterates through the sequences, printing each item and counting the number of ones in the input.

```python
# Parity dataset

from circuits import parity

import random
random_number = random.randint(0, 1000)

input_seq_len=4
data  = parity(
    vocab_size=256,
    num_examples=10,
    input_seq_len=input_seq_len,
    seed=random_number,
)
x= data.inputs[0]
y= data.labels[0]

print(f"Lengths: {len(x)}, {len(y)}\n")
num_one = 0
for i, (_x, _y) in enumerate(zip(x, y)):
    if _x.item() == 1:
        num_one += 1
    print(_x.item(), _y.item())

print(num_one)
```

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

### Define Sweep Configurations

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Create a list of `TrainConfig` objects by iterating through parameter variations, such as learning rates, to define a sweep. Each configuration in the list will be launched as a separate job.

```python
import numpy as np
from zoology.config import TrainConfig

configs = []
for lr in np.logspace(-4, -2, 10):
   configs.append(TrainConfig(learning_rate=lr))
```

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

### Launch Paper Experiment Sweep

Source: https://github.com/hazyresearch/zoology/blob/main/zoology/experiments/paper_configs/iclr24_zoology_figure2/README.md

Use this command to launch the main synthetic data experiment sweep for reproducing Figure 2 results. Ensure WandB is set up for logging. The -p flag enables parallel execution.

```bash
python -m zoology.launch zoology/experiments/paper/figure2.py -p
```

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

### Define Sweep Experiment Configuration

Source: https://context7.com/hazyresearch/zoology/llms.txt

Sets up a Python script to define multiple training configurations for hyperparameter sweeps, iterating over learning rates and model dimensions.

```python
# sweep_experiment.py
import numpy as np
from zoology.config import TrainConfig, ModelConfig, DataConfig, ModuleConfig
from zoology.data.multiquery_ar import MQARConfig

configs = []

# Sweep over learning rates
for lr in np.logspace(-4, -2, 10):
    # Sweep over model dimensions
    for d_model in [64, 128, 256]:
        config = TrainConfig(
            data=DataConfig(
                train_configs=[MQARConfig(num_examples=10_000, vocab_size=8192, input_seq_len=64, num_kv_pairs=4)],
                test_configs=[MQARConfig(num_examples=1_000, vocab_size=8192, input_seq_len=64, num_kv_pairs=4)]
            ),
            model=ModelConfig(
                vocab_size=8192,
                d_model=d_model,
                n_layers=2,
                sequence_mixer=ModuleConfig(
                    name="zoology.mixers.attention.MHA",
                    kwargs={"num_heads": 2, "dropout": 0.1}
                )
            ),
            learning_rate=lr,
            max_epochs=100,
            run_id=f"d{d_model}-lr{lr:.1e}"
        )
        configs.append(config)
```

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

### Fetch and filter Weights & Biases runs

Source: https://github.com/hazyresearch/zoology/blob/main/zoology/experiments/paper_configs/arxiv24_based_rebuttal_circuits/plot.ipynb

Fetches experiment runs from Weights & Biases for specified launch IDs and project name. Filters the results to include only runs that have finished.

```python
df = fetch_wandb_runs(
    launch_id=[
        # "redo-vocab-majority-sweep-V654618ae"
        "default-2024-03-29-00-48-10",
        "default-2024-03-29-01-21-59"

        
    ], 
    project_name="zoology"
)

df = df[df["state"] == "finished"]
```

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

### Configure Training Data with MyDataSegmentConfig

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Add a custom data segment configuration to the `TrainConfig`. Specify parameters like `num_examples`, `vocab_size`, and `input_seq_len` for training and testing splits.

```python
from zoology.config import TrainConfig, DataConfig, FunctionConfig
config = TrainConfig(
    DataConfig(
        train_configs=[MyDataSegmentConfig(num_examples=10_000, vocab_size=128, input_seq_len=input_seq_len, **other_kwargs)],
        test_configs=[MyDataSegmentConfig(num_examples=1_000, vocab_size=128, input_seq_len=input_seq_len, **other_kwargs)],
    ),
)
```

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

### Identify optimal configurations

Source: https://github.com/hazyresearch/zoology/blob/main/zoology/experiments/paper_configs/arxiv24_based_rebuttal_circuits/plot.ipynb

Groups the DataFrame by model name and input sequence length, then selects the row with the maximum validation accuracy for each group. This is useful for finding the best performing configuration.

```python
idxs = df.groupby(
    ["model.name", "data.train_configs.0.input_seq_len"]
){"valid/accuracy"}.idxmax()
plot_df = df.loc[idxs]
```

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

### Load and Inspect Cumulative Parity Dataset

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/test_data.ipynb

Loads the `cumulative_parity` dataset with a specified input sequence length and prints the lengths of the input and label sequences. It then iterates through the sequences, printing each item and counting the number of ones in the input.

```python
# Parity dataset

from circuits import cumulative_parity

import random
random_number = random.randint(0, 1000)

input_seq_len=8
data  = cumulative_parity(
    vocab_size=256,
    num_examples=10,
    input_seq_len=input_seq_len,
    seed=random_number,
)
x= data.inputs[0]
y= data.labels[0]

print(f"Lengths: {len(x)}, {len(y)}\n")
num_one = 0
for i, (_x, _y) in enumerate(zip(x, y)):
    if _x.item() == 1:
        num_one += 1
    print(_x.item(), _y.item())

print(num_one)
```

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

### Define Query Projection Matrix (Layer 2)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Sets the weights for the Query (Q2) projection matrix in the second attention layer. This configuration ensures that Q2 yields only the token embeddings from the input, preparing them for matching with keys.

```python
Q2[:position_dim] = # YOUR ANSWER SCALAR
```

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

### Load and Inspect Forgetting MQAR Data

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/test_data.ipynb

Imports the `forgetting_mqar` dataset and iterates through the first input-label pair, printing values up to the 30th element. Ensure the 'forgetting_ar' module is accessible.

```python
import sys
sys.path.append("/home/simarora/code/zoology/zoology/data/")
```

```python
# Step 2: Import the data you want to inspect

from forgetting_ar import forgetting_mqar
data  = forgetting_mqar(
    vocab_size=256,
    num_examples=10,
    input_seq_len=24,
    num_kv_pairs=4,
    seed=0,
)
x= data.inputs[0]
y= data.labels[0]

for i, (_x, _y) in enumerate(zip(x, y)):
    if i == 30:
        break
    print(f"{_x.item()}, {_y.item()}")
```

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

### Initialize Positional Embeddings

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Creates positional encodings for the input sequence. Each position in the sequence is represented by a unique vector.

```python
# Define the positional embeddings
batch_size, sequence_length = examples.shape[0], examples.shape[1]
seq_len = sequence_length 
positional_encodings = torch.zeros((seq_len, seq_len))
for pos in range(seq_len):
    positional_encodings[pos][pos] = 1
plt.imshow(positional_encodings.detach().numpy())
print(seq_len)
```

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

### Create Custom Data Task with DataSegmentConfig

Source: https://context7.com/hazyresearch/zoology/llms.txt

Subclass DataSegmentConfig and implement the build method to create custom synthetic data tasks. Define custom parameters and generation logic within the build method. Use the custom config by instantiating it with desired parameters.

```python
from zoology.config import DataSegmentConfig
from zoology.data.utils import DataSegment
import torch

class MyCustomTaskConfig(DataSegmentConfig):
    vocab_size: int = 1024
    num_examples: int = 1000
    input_seq_len: int = 128
    custom_param: float = 0.5  # Add custom parameters

    def build(self, seed: int) -> DataSegment:
        import numpy as np
        np.random.seed(seed)

        # Generate custom input/label sequences
        inputs = torch.randint(0, self.vocab_size, (self.num_examples, self.input_seq_len))
        labels = torch.full((self.num_examples, self.input_seq_len), -100)

        # Your custom logic here to set labels at specific positions
        # labels at position i should be the target token (not -100)

        return DataSegment(
            inputs=inputs,
            labels=labels,
            slices={"custom_param": self.custom_param}
        )

# Use in config
custom_config = MyCustomTaskConfig(
    num_examples=5000,
    custom_param=0.8
)
```

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

### Load and Inspect Multiquery AR Data

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/test_data.ipynb

Imports the `multiquery_ar` dataset and iterates through the first input-label pair, printing tensor values up to the 30th element. Ensure the 'associative_recall' module is accessible.

```python
# Step 2: Import the data you want to inspect

from associative_recall import multiquery_ar
data  = multiquery_ar(
    vocab_size=256,
    num_examples=10,
    input_seq_len=128,
    seed=0,
)
x= data.inputs[0]
y= data.labels[0]

for i, (_x, _y) in enumerate(zip(x, y)):
    if i == 30:
        break
    print(_x, _y)
```

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

### Generate Associative Recall Dataset

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Generates a dataset for the associative recall task. Use this function to create sample data for training and evaluation.

```python
def generate_dataset(
    vocab_size: int = 50,
    num_examples: int = 1000,
    seed: int = 0,
) -> List[str]:
    """
    Generate a dataset for the associative recall task.
    """
    np.random.seed(seed)
    torch.manual_seed(seed)
    tokens = torch.arange(vocab_size) + 1
    keys   = tokens[: (vocab_size // 2) - 1]
    values = tokens[((vocab_size // 2) - 1):-1]
    examples = []
    repeated = []
    golds = []
    for _ in range(num_examples):
        shuffled_keys = keys[torch.randperm(len(keys))]
        shuffled_values = values[torch.randperm(len(values))]
        kv_pairs = list(zip(shuffled_keys, shuffled_values))  
        example = kv_pairs[:4*len(kv_pairs) // 8]
        example += kv_pairs[:4*len(kv_pairs) // 8]
        example += kv_pairs[:4*len(kv_pairs) // 8]
        # shuffle example and repeat_pos
        perm = torch.randperm(len(example))
        example = [example[i] for i in perm]
        repeat_pos = []
        repeats = {}
        for i, (k, v) in enumerate(example):
            if int(k) in repeats:
                repeat_pos.append(i)
            repeats[int(k)] = int(v)
        example_tensor = torch.tensor(example).flatten()
        repeat_tensor = [i*2 for i in repeat_pos]
        gold = []
        for i in repeat_tensor:
            gold.append(int(example_tensor[i+1]))
        repeated.append(repeat_tensor)
        examples.append(example_tensor)
        golds.append(gold)
    examples = torch.stack(examples)
    return examples, repeated, golds
```

```python
vocab_size=10
examples, repeated, golds = generate_dataset(vocab_size=vocab_size)
examples[0], repeated[0], examples[0][repeated[0]], golds[0]
```

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

### Define Key Projection Matrix (Layer 1)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Defines the weights for the Key (K) projection matrix in the first attention layer. This configuration aims to zero out token embeddings while preserving positional embeddings, influencing how keys are generated.

```python
K[position_dim:] = 0
```

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

### Define Key Projection Matrix (Layer 2)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Defines the weights for the Key (K2) projection matrix in the second attention layer. This configuration is designed to capture shifted token embeddings from the input.

```python
K2[:embed_dim, position_dim:] =  100  * #YOUR ANSWER TENSOR
```

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

### Initialize Token Embeddings

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Defines one-hot token embeddings. Each token is represented by a unique vector where only the corresponding dimension is set to 1.

```python
# Define the one-hot token embeddings 
emb_pos_2_tok_id = {}
embeddings_init = torch.zeros(vocab_size, vocab_size)
for i in range(vocab_size):
    embeddings_init[i][i] = 1
    emb_pos_2_tok_id[i] = embeddings_init[i]
plt.imshow(embeddings_init.detach().numpy())
```

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

### Tensor Manipulation: Argmax and Subtraction

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Demonstrates taking the argmax of O2 to find the token embedding with the maximum value, followed by subtracting the position dimension. This is useful for analyzing token importance and positional encoding.

```python
z2 before subtracting position dim -- we take the argmax of O2 to get the token embedding with the 'max' value
tensor([14, 14, 21, 14, 21, 14, 14, 14, 17, 13, 17, 13])
```

```python
z2 after subtracting position dim (12)
tensor([2, 2, 9, 2, 9, 2, 2, 2, 5, 1, 5, 1])
```

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

### Load and Auto-reload Python Modules

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Imports necessary libraries for data manipulation and visualization. Ensures that changes in imported modules are automatically reloaded.

```python
%load_ext autoreload
%autoreload 2

import matplotlib.pyplot as plt
import math
from typing import List
import numpy as np
import torch
```

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

### Perform Attention Calculation (Layer 1)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Executes the self-attention mechanism for the first layer using the defined Q, K, and V matrices. It computes attention scores, applies softmax, masks future positions, and generates the output embeddings.

```python
q1 = embs @ Q 
k1 = embs @ K 
v1 = embs @ V 
S1 = torch.matmul(q1, k1.transpose(1, 2)) / math.sqrt(Q.shape[0])
P1 = torch.nn.functional.softmax(S1, dim=1)
P1 = torch.tril(P1)
O1 = torch.matmul(P1, v1)
z1 = O1 + embs
```

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

### Configure Hyena Long Convolution Operator

Source: https://context7.com/hazyresearch/zoology/llms.txt

Defines the configuration for the Hyena sequence mixer, specifying parameters like maximum sequence length and filter order.

```python
from zoology.config import ModelConfig, ModuleConfig

hyena_mixer = ModuleConfig(
    name="zoology.mixers.hyena.Hyena",
    kwargs={
        "l_max": 1024,           # Maximum sequence length
        "filter_order": 64,       # Width of filter MLP
        "num_heads": 1,
        "short_filter_order": 3,  # Short convolution kernel size
        "dropout": 0.0
    }
)

model_config = ModelConfig(
    vocab_size=8192,
    d_model=128,
    n_layers=2,
    max_position_embeddings=0,
    sequence_mixer=hyena_mixer
)
```

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

### Configure MHA (Multi-Head Attention) Model

Source: https://context7.com/hazyresearch/zoology/llms.txt

Configure a standard causal multi-head self-attention model using ModelConfig. Specify vocabulary size, model dimensions, number of layers, and attention module details. The MHA module can also be instantiated and used directly for custom applications.

```python
from zoology.config import ModelConfig, ModuleConfig

# Configure attention model
model_config = ModelConfig(
    vocab_size=8192,
    d_model=128,
    n_layers=2,
    max_position_embeddings=1024,
    sequence_mixer=ModuleConfig(
        name="zoology.mixers.attention.MHA",
        kwargs={
            "num_heads": 4,
            "dropout": 0.1
        }
    )
)

# The MHA module can also be used directly
from zoology.mixers.attention import MHA
import torch

mha = MHA(d_model=128, num_heads=4, dropout=0.1)
x = torch.randn(2, 64, 128)  # (batch, seq_len, d_model)
output = mha(x)  # (batch, seq_len, d_model)

# Get state size for memory analysis
state_size = mha.state_size(sequence_length=2048)
```

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

### Embed Input Sequence

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Initializes the embedding tensor for the input sequence by concatenating token and positional embeddings.

```python
# embed the input
position_dim, embed_dim = positional_encodings.shape[1], embeddings_init.shape[1] 
d_model = position_dim + embed_dim 
embs = torch.zeros((batch_size, seq_len, d_model))
```

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

### Define a New Data Task with DataSegmentConfig

Source: https://github.com/hazyresearch/zoology/blob/main/README.md

Subclass `DataSegmentConfig` to define a new synthetic task. Implement the `build` method to return a `DataSegment` object containing integer tensors for inputs and labels within the range [0, vocab_size).

```python
class DataSegmentConfig(BaseConfig):
    """
    This class should be subclassed to define per task. For example, MQARConfig
    """
    vocab_size: int = 8_192
    num_examples: int = 1_000
    input_seq_len: int = 64

    def build(self, **kwargs):
        raise NotImplementedError()
```

```python
@dataclass
class DataSegment:
    inputs: torch.Tensor
    labels: torch.Tensor
    slices: Dict[str, any] = None
```

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

### Define Value Projection Matrix (Layer 2)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Sets the weights for the Value (V2) projection matrix in the second attention layer. This configuration allows V2 to extract original token embeddings, facilitating associative recall based on the attention scores.

```python
V2[:position_dim] = # YOUR ANSWER SCALAR
```

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

### Perform Attention Calculation (Layer 2)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Executes the self-attention mechanism for the second layer using the defined Q2, K2, and V2 matrices. This layer rearranges shifted token embeddings to the correct output positions for subsequent predictions.

```python
q2 = z1 @ Q2
k2 = z1 @ K2
v2 = z1 @ V2
S2 = torch.matmul(q2, k2.transpose(1, 2)) / math.sqrt(Q2.shape[0])
P2 = torch.nn.functional.softmax(S2, dim=1)
P2 = torch.tril(P2)
O2 = torch.matmul(P2, v2)
```

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

### O2 Output Interpretation

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

O2 represents associative recall output tokens. It is used to understand the model's recall capabilities.

```text
O2 -- O2 has the associative recall output tokens
```

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

### Define Value Projection Matrix (Layer 1)

Source: https://github.com/hazyresearch/zoology/blob/main/notebooks/program_transformer_scaffolding.ipynb

Defines the weights for the Value (V) projection matrix in the first attention layer. This matrix isolates token embeddings, ensuring that the output of the attention mechanism contains the value of the token embedding that follows a key token.

```python
V[position_dim:, :embed_dim] = 100 * #YOUR ANSWER TENSOR
```

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

### Plotting Model Accuracy vs. State Size

Source: https://context7.com/hazyresearch/zoology/llms.txt

This code preprocesses data, cleans model names, and generates a scatter plot to visualize the recall accuracy of different models against their state size. It uses Seaborn for plotting and Matplotlib for saving the figure.

```python
idx = df.groupby(["state_size", "model.name"])["valid/accuracy"].idxmax(skipna=True).dropna()
plot_df = df.loc[idx]

# Clean up model names
plot_df["Model"] = plot_df["model.name"].str.replace("_", " ").str.title()

# Color palette
model2color = {
    "Attention": "black",
    "Based": "#59A14F",
    "Mamba2": "#4E79A7",
    "Gated Delta Net": "#9C755F",
    "Gla": "#F28E2B"
}

# Create scatter plot
g = sns.relplot(
    data=plot_df,
    y="valid/accuracy",
    x="state_size",
    hue="Model",
    kind="scatter",
    height=5,
    aspect=1,
    palette=model2color,
    s=60,
    edgecolor="black",
    linewidth=0.5
)
g.set(xscale="log", ylabel="Recall Accuracy", xlabel="State Size (log scale)")
plt.title("MQAR Recall-Memory Tradeoff")
plt.savefig("results.png", dpi=300, bbox_inches="tight")
```