diff --git a/sae_lens/config.py b/sae_lens/config.py
index 2175df33..085ee36c 100644
--- a/sae_lens/config.py
+++ b/sae_lens/config.py
@@ -124,6 +124,7 @@ class LanguageModelSAERunnerConfig:
     )
 
     # SAE Parameters
+    architecture: Literal["standard", "gated"] = "standard"
     d_in: int = 512
     d_sae: Optional[int] = None
     b_dec_init_method: str = "geometric_median"
@@ -349,6 +350,8 @@ def total_training_steps(self) -> int:
 
     def get_base_sae_cfg_dict(self) -> dict[str, Any]:
         return {
+            # TEMP
+            "architecture": self.architecture,
             "d_in": self.d_in,
             "d_sae": self.d_sae,
             "dtype": self.dtype,
@@ -474,6 +477,9 @@ def __post_init__(self):
 
 @dataclass
 class ToyModelSAERunnerConfig:
+
+    architecture: Literal["standard", "gated"] = "standard"
+
     # ReLu Model Parameters
     n_features: int = 5
     n_hidden: int = 2
@@ -527,6 +533,7 @@ def __post_init__(self):
     def get_base_sae_cfg_dict(self) -> dict[str, Any]:
         # TO DO: Have the same hyperparameters as in the main sae runner.
         return {
+            "architecture": self.architecture,
             "d_in": self.d_in,
             "d_sae": self.d_sae,
             "dtype": self.dtype,
diff --git a/sae_lens/sae.py b/sae_lens/sae.py
index dd51d7e1..bd0f91b2 100644
--- a/sae_lens/sae.py
+++ b/sae_lens/sae.py
@@ -5,7 +5,7 @@
 import json
 import os
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Literal, Optional, Tuple
 
 import einops
 import torch
@@ -28,6 +28,8 @@
 
 @dataclass
 class SAEConfig:
+    # architecture details
+    architecture: Literal["standard", "gated"]
 
     # forward pass details.
     d_in: int
@@ -76,6 +78,7 @@ def from_dict(cls, config_dict: dict[str, Any]) -> "SAEConfig":
 
     def to_dict(self) -> dict[str, Any]:
         return {
+            "architecture": self.architecture,
             "d_in": self.d_in,
             "d_sae": self.d_sae,
             "dtype": self.dtype,
@@ -121,7 +124,12 @@ def __init__(
         self.device = torch.device(cfg.device)
         self.use_error_term = use_error_term
 
-        self.initialize_weights_basic()
+        if self.cfg.architecture == "standard":
+            self.initialize_weights_basic()
+            self.encode_fn = self.encode
+        elif self.cfg.architecture == "gated":
+            self.initialize_weights_gated()
+            self.encode_fn = self.encode_gated
 
         # handle presence / absence of scaling factor.
         if self.cfg.finetuning_scaling_factor:
@@ -209,16 +217,50 @@ def initialize_weights_basic(self):
                 torch.ones(self.cfg.d_sae, dtype=self.dtype, device=self.device)
             )
 
+    def initialize_weights_gated(self):
+        # Initialize the weights and biases for the gated encoder
+        self.W_enc = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(
+                    self.cfg.d_in, self.cfg.d_sae, dtype=self.dtype, device=self.device
+                )
+            )
+        )
+
+        self.b_gate = nn.Parameter(
+            torch.zeros(self.cfg.d_sae, dtype=self.dtype, device=self.device)
+        )
+
+        self.r_mag = nn.Parameter(
+            torch.zeros(self.cfg.d_sae, dtype=self.dtype, device=self.device)
+        )
+
+        self.b_mag = nn.Parameter(
+            torch.zeros(self.cfg.d_sae, dtype=self.dtype, device=self.device)
+        )
+
+        self.W_dec = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(
+                    self.cfg.d_sae, self.cfg.d_in, dtype=self.dtype, device=self.device
+                )
+            )
+        )
+
+        self.b_dec = nn.Parameter(
+            torch.zeros(self.cfg.d_in, dtype=self.dtype, device=self.device)
+        )
+
     # Basic Forward Pass Functionality.
     def forward(
         self,
         x: torch.Tensor,
     ) -> torch.Tensor:
-
-        feature_acts = self.encode(x)
+        feature_acts = self.encode_fn(x)
         sae_out = self.decode(feature_acts)
 
-        if self.use_error_term:
+        # TEMP
+        if self.use_error_term and self.cfg.architecture != "gated":
             with torch.no_grad():
                 # Recompute everything without hooks to get true error term
                 # Otherwise, the output with error term will always equal input, even for causal interventions that affect x_reconstruct
@@ -248,16 +290,62 @@ def forward(
 
                 sae_out = self.run_time_activation_norm_fn_out(sae_out)
                 sae_error = self.hook_sae_error(x - x_reconstruct_clean)
+            return self.hook_sae_output(sae_out + sae_error)
+
+        # TODO: Add tests
+        elif self.use_error_term and self.cfg.architecture == "gated":
+            with torch.no_grad():
+                x = x.to(self.dtype)
+                sae_in = self.reshape_fn_in(x)  # type: ignore
+                gating_pre_activation = sae_in @ self.W_enc + self.b_gate
+                active_features = (gating_pre_activation > 0).float()
+
+                # Magnitude path with weight sharing
+                magnitude_pre_activation = self.hook_sae_acts_pre(
+                    sae_in @ (self.W_enc * self.r_mag.exp()) + self.b_mag
+                )
+                feature_magnitudes = self.hook_sae_acts_post(
+                    self.activation_fn(magnitude_pre_activation)
+                )
+                feature_acts_clean = active_features * feature_magnitudes
+                x_reconstruct_clean = self.reshape_fn_out(
+                    self.apply_finetuning_scaling_factor(feature_acts_clean)
+                    @ self.W_dec
+                    + self.b_dec,
+                    d_head=self.d_head,
+                )
 
+                sae_error = self.hook_sae_error(x - x_reconstruct_clean)
             return self.hook_sae_output(sae_out + sae_error)
 
         return self.hook_sae_output(sae_out)
 
+    def encode_gated(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        x = x.to(self.dtype)
+        x = self.reshape_fn_in(x)
+        sae_in = self.hook_sae_input(x - self.b_dec * self.cfg.apply_b_dec_to_input)
+
+        # Gating path
+        gating_pre_activation = sae_in @ self.W_enc + self.b_gate
+        active_features = (gating_pre_activation > 0).float()
+
+        # Magnitude path with weight sharing
+        magnitude_pre_activation = self.hook_sae_acts_pre(
+            sae_in @ (self.W_enc * self.r_mag.exp()) + self.b_mag
+        )
+        feature_magnitudes = self.hook_sae_acts_post(
+            self.activation_fn(magnitude_pre_activation)
+        )
+
+        return active_features * feature_magnitudes
+
     def encode(
         self, x: Float[torch.Tensor, "... d_in"]
     ) -> Float[torch.Tensor, "... d_sae"]:
         """
-        Calcuate SAE features from inputs
+        Calculate SAE features from inputs
         """
 
         # move x to correct dtype
@@ -301,7 +389,12 @@ def fold_W_dec_norm(self):
         W_dec_norms = self.W_dec.norm(dim=-1).unsqueeze(1)
         self.W_dec.data = self.W_dec.data / W_dec_norms
         self.W_enc.data = self.W_enc.data * W_dec_norms.T
-        self.b_enc.data = self.b_enc.data * W_dec_norms.squeeze()
+        if self.cfg.architecture == "gated":
+            self.r_mag.data = self.r_mag.data * W_dec_norms.squeeze()
+            self.b_gate.data = self.b_gate.data * W_dec_norms.squeeze()
+            self.b_mag.data = self.b_mag.data * W_dec_norms.squeeze()
+        else:
+            self.b_enc.data = self.b_enc.data * W_dec_norms.squeeze()
 
     @torch.no_grad()
     def fold_activation_norm_scaling_factor(
diff --git a/sae_lens/toolkit/pretrained_sae_loaders.py b/sae_lens/toolkit/pretrained_sae_loaders.py
index e5e9eb1d..3bc72a4d 100644
--- a/sae_lens/toolkit/pretrained_sae_loaders.py
+++ b/sae_lens/toolkit/pretrained_sae_loaders.py
@@ -124,6 +124,7 @@ def connor_rob_hook_z_loader(
     # }
 
     cfg_dict = {
+        "architecture": "standard",
         "d_in": old_cfg_dict["act_size"],
         "d_sae": old_cfg_dict["dict_size"],
         "dtype": "float32",
@@ -169,6 +170,10 @@ def load_pretrained_sae_lens_sae_components(
     with open(cfg_path, "r") as f:
         cfg_dict = json.load(f)
 
+    cfg_dict["architecture"] = (
+        "standard"  # TODO: modify this when we add support for loading more architectures
+    )
+
     # filter config for varnames
     cfg_dict["device"] = device
     cfg_dict["dtype"] = dtype
diff --git a/sae_lens/training/sae_trainer.py b/sae_lens/training/sae_trainer.py
index 87ec8c27..34c7c0d5 100644
--- a/sae_lens/training/sae_trainer.py
+++ b/sae_lens/training/sae_trainer.py
@@ -287,6 +287,7 @@ def _build_train_step_log_dict(
             "losses/mse_loss": mse_loss,
             "losses/l1_loss": l1_loss
             / self.current_l1_coefficient,  # normalize by l1 coefficient
+            "losses/auxiliary_reconstruction_loss": output.auxiliary_reconstruction_loss,
             "losses/ghost_grad_loss": ghost_grad_loss,
             "losses/overall_loss": loss,
             # variance explained
@@ -317,19 +318,17 @@ def _run_and_log_evals(self):
                 model_kwargs=self.cfg.model_kwargs,
             )
 
-            W_dec_norm_dist = self.sae.W_dec.norm(dim=1).detach().float().cpu().numpy()
-            b_e_dist = self.sae.b_enc.detach().float().cpu().numpy()
-
-            # More detail on loss.
-
-            # add weight histograms
-            eval_metrics = {
-                **eval_metrics,
-                **{
-                    "weights/W_dec_norms": wandb.Histogram(W_dec_norm_dist),
-                    "weights/b_e": wandb.Histogram(b_e_dist),
-                },
-            }
+            W_dec_norm_dist = self.sae.W_dec.norm(dim=1).detach().cpu().numpy()
+            eval_metrics["weights/W_dec_norms"] = wandb.Histogram(W_dec_norm_dist)  # type: ignore
+
+            if self.sae.cfg.architecture == "standard":
+                b_e_dist = self.sae.b_enc.detach().cpu().numpy()
+                eval_metrics["weights/b_e"] = wandb.Histogram(b_e_dist)  # type: ignore
+            elif self.sae.cfg.architecture == "gated":
+                b_gate_dist = self.sae.b_gate.detach().cpu().numpy()
+                eval_metrics["weights/b_gate"] = wandb.Histogram(b_gate_dist)  # type: ignore
+                b_mag_dist = self.sae.b_mag.detach().cpu().numpy()
+                eval_metrics["weights/b_mag"] = wandb.Histogram(b_mag_dist)  # type: ignore
 
             wandb.log(
                 eval_metrics,
diff --git a/sae_lens/training/training_sae.py b/sae_lens/training/training_sae.py
index 7f99d1d2..e21acff3 100644
--- a/sae_lens/training/training_sae.py
+++ b/sae_lens/training/training_sae.py
@@ -27,6 +27,7 @@ class TrainStepOutput:
     mse_loss: float
     l1_loss: float
     ghost_grad_loss: float
+    auxiliary_reconstruction_loss: float = 0.0
 
 
 @dataclass
@@ -50,7 +51,8 @@ def from_sae_runner_config(
     ) -> "TrainingSAEConfig":
 
         return cls(
-            # base confg
+            # base config
+            architecture=cfg.architecture,
             d_in=cfg.d_in,
             d_sae=cfg.d_sae,  # type: ignore
             dtype=cfg.dtype,
@@ -104,6 +106,7 @@ def to_dict(self) -> dict[str, Any]:
     # parameters. Maybe there's a cleaner way to do this
     def get_base_sae_cfg_dict(self) -> dict[str, Any]:
         return {
+            "architecture": self.architecture,
             "d_in": self.d_in,
             "d_sae": self.d_sae,
             "activation_fn_str": self.activation_fn_str,
@@ -140,6 +143,15 @@ def __init__(self, cfg: TrainingSAEConfig, use_error_term: bool = False):
         base_sae_cfg = SAEConfig.from_dict(cfg.get_base_sae_cfg_dict())
         super().__init__(base_sae_cfg)
         self.cfg = cfg  # type: ignore
+
+        self.encode_with_hidden_pre_fn = (
+            self.encode_with_hidden_pre
+            if cfg.architecture != "gated"
+            else self.encode_with_hidden_pre_gated
+        )
+
+        self.check_cfg_compatibility()
+
         self.use_error_term = use_error_term
 
         self.initialize_weights_complex()
@@ -154,13 +166,20 @@ def __init__(self, cfg: TrainingSAEConfig, use_error_term: bool = False):
     def from_dict(cls, config_dict: dict[str, Any]) -> "TrainingSAE":
         return cls(TrainingSAEConfig.from_dict(config_dict))
 
+    def check_cfg_compatibility(self):
+        if self.cfg.architecture == "gated":
+            assert (
+                self.cfg.use_ghost_grads is False
+            ), "Gated SAEs do not support ghost grads"
+            assert self.use_error_term is False, "Gated SAEs do not support error terms"
+
     def encode(
         self, x: Float[torch.Tensor, "... d_in"]
     ) -> Float[torch.Tensor, "... d_sae"]:
         """
         Calcuate SAE features from inputs
         """
-        feature_acts, _ = self.encode_with_hidden_pre(x)
+        feature_acts, _ = self.encode_with_hidden_pre_fn(x)
         return feature_acts
 
     def encode_with_hidden_pre(
@@ -188,12 +207,44 @@ def encode_with_hidden_pre(
 
         return feature_acts, hidden_pre_noised
 
+    def encode_with_hidden_pre_gated(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+
+        # move x to correct dtype
+        x = x.to(self.dtype)
+
+        # handle hook z reshaping if needed.
+        x = self.reshape_fn_in(x)  # type: ignore
+
+        # apply b_dec_to_input if using that method.
+        sae_in = self.hook_sae_input(x - (self.b_dec * self.cfg.apply_b_dec_to_input))
+
+        # Gating path with Heaviside step function
+        gating_pre_activation = sae_in @ self.W_enc + self.b_gate
+        active_features = (gating_pre_activation > 0).float()
+
+        # Magnitude path with weight sharing
+        magnitude_pre_activation = sae_in @ (self.W_enc * self.r_mag.exp()) + self.b_mag
+        # magnitude_pre_activation_noised = magnitude_pre_activation + (
+        #     torch.randn_like(magnitude_pre_activation) * self.cfg.noise_scale * self.training
+        # )
+        feature_magnitudes = self.activation_fn(
+            magnitude_pre_activation
+        )  # magnitude_pre_activation_noised)
+
+        # Return both the gated feature activations and the magnitude pre-activations
+        return (
+            active_features * feature_magnitudes,
+            magnitude_pre_activation,
+        )  # magnitude_pre_activation_noised
+
     def forward(
         self,
         x: Float[torch.Tensor, "... d_in"],
     ) -> Float[torch.Tensor, "... d_in"]:
 
-        feature_acts, _ = self.encode_with_hidden_pre(x)
+        feature_acts, _ = self.encode_with_hidden_pre_fn(x)
         sae_out = self.decode(feature_acts)
 
         return sae_out
@@ -207,7 +258,7 @@ def training_forward_pass(
 
         # do a forward pass to get SAE out, but we also need the
         # hidden pre.
-        feature_acts, _ = self.encode_with_hidden_pre(sae_in)
+        feature_acts, _ = self.encode_with_hidden_pre_fn(sae_in)
         sae_out = self.decode(feature_acts)
 
         # MSE LOSS
@@ -218,7 +269,7 @@ def training_forward_pass(
         if self.cfg.use_ghost_grads and self.training and dead_neuron_mask is not None:
 
             # first half of second forward pass
-            _, hidden_pre = self.encode_with_hidden_pre(sae_in)
+            _, hidden_pre = self.encode_with_hidden_pre_fn(sae_in)
             ghost_grad_loss = self.calculate_ghost_grad_loss(
                 x=sae_in,
                 sae_out=sae_out,
@@ -229,17 +280,40 @@ def training_forward_pass(
         else:
             ghost_grad_loss = 0.0
 
-        # SPARSITY LOSS
-        # either the W_dec norms are 1 and this won't do anything or they are not 1
-        # and we're using their norm in the loss function.
-        weighted_feature_acts = feature_acts * self.W_dec.norm(dim=1)
-        sparsity = weighted_feature_acts.norm(
-            p=self.cfg.lp_norm, dim=-1
-        )  # sum over the feature dimension
+        if self.cfg.architecture == "gated":
+            # Gated SAE Loss Calculation
+
+            # Shared variables
+            sae_in_centered = (
+                self.reshape_fn_in(sae_in) - self.b_dec * self.cfg.apply_b_dec_to_input
+            )
+            pi_gate = sae_in_centered @ self.W_enc + self.b_gate
+            pi_gate_act = torch.relu(pi_gate)
+
+            # SFN sparsity loss - summed over the feature dimension and averaged over the batch
+            l1_loss = (
+                current_l1_coefficient
+                * torch.sum(pi_gate_act * self.W_dec.norm(dim=1), dim=-1).mean()
+            )
+
+            # Auxiliary reconstruction loss - summed over the feature dimension and averaged over the batch
+            via_gate_reconstruction = pi_gate_act @ self.W_dec + self.b_dec
+            aux_reconstruction_loss = torch.sum(
+                (via_gate_reconstruction - sae_in) ** 2, dim=-1
+            ).mean()
+
+            loss = mse_loss + l1_loss + aux_reconstruction_loss
+        else:
+            # default SAE sparsity loss
+            weighted_feature_acts = feature_acts * self.W_dec.norm(dim=1)
+            sparsity = weighted_feature_acts.norm(
+                p=self.cfg.lp_norm, dim=-1
+            )  # sum over the feature dimension
 
-        l1_loss = (current_l1_coefficient * sparsity).mean()
+            l1_loss = (current_l1_coefficient * sparsity).mean()
+            loss = mse_loss + l1_loss + ghost_grad_loss
 
-        loss = mse_loss + l1_loss + ghost_grad_loss
+            aux_reconstruction_loss = torch.tensor(0.0)
 
         return TrainStepOutput(
             sae_in=sae_in,
@@ -253,6 +327,7 @@ def training_forward_pass(
                 if isinstance(ghost_grad_loss, torch.Tensor)
                 else ghost_grad_loss
             ),
+            auxiliary_reconstruction_loss=aux_reconstruction_loss.item(),
         )
 
     def calculate_ghost_grad_loss(
@@ -354,7 +429,7 @@ def initialize_weights_complex(self):
         elif self.cfg.normalize_sae_decoder:
             self.set_decoder_norm_to_unit_norm()
 
-        # Then we intialize the encoder weights (either as the transpose of decoder or not)
+        # Then we initialize the encoder weights (either as the transpose of decoder or not)
         if self.cfg.init_encoder_as_decoder_transpose:
             self.W_enc.data = self.W_dec.data.T.clone().contiguous()
         else:
diff --git a/tests/benchmark/test_language_model_sae_runner.py b/tests/benchmark/test_language_model_sae_runner.py
index d5eb3e29..55293b04 100644
--- a/tests/benchmark/test_language_model_sae_runner.py
+++ b/tests/benchmark/test_language_model_sae_runner.py
@@ -103,3 +103,101 @@ def test_language_model_sae_runner():
 
     assert sae is not None
     # know whether or not this works by looking at the dashboard!
+
+
+def test_language_model_sae_runner_gated():
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    # total_training_steps = 20_000
+    total_training_steps = 500
+    batch_size = 4096
+    total_training_tokens = total_training_steps * batch_size
+    print(f"Total Training Tokens: {total_training_tokens}")
+
+    lr_warm_up_steps = 0
+    lr_decay_steps = 40_000
+    print(f"lr_decay_steps: {lr_decay_steps}")
+    l1_warmup_steps = 10_000
+    print(f"l1_warmup_steps: {l1_warmup_steps}")
+
+    cfg = LanguageModelSAERunnerConfig(
+        # Pick a tiny model to make this easier.
+        model_name="gelu-1l",
+        architecture="gated",
+        ## MLP Layer 0 ##
+        hook_name="blocks.0.hook_mlp_out",
+        hook_layer=0,
+        d_in=512,
+        dataset_path="NeelNanda/c4-tokenized-2b",
+        context_size=256,
+        is_dataset_tokenized=True,
+        prepend_bos=True,  # I used to train GPT2 SAEs with a prepended-bos but no longer think we should do this.
+        # How big do we want our SAE to be?
+        expansion_factor=16,
+        # Dataset / Activation Store
+        # When we do a proper test
+        # training_tokens= 820_000_000, # 200k steps * 4096 batch size ~ 820M tokens (doable overnight on an A100)
+        # For now.
+        training_tokens=total_training_tokens,  # For initial testing I think this is a good number.
+        train_batch_size_tokens=4096,
+        # Loss Function
+        ## Reconstruction Coefficient.
+        mse_loss_normalization=None,  # MSE Loss Normalization is not mentioned (so we use stanrd MSE Loss). But not we take an average over the batch.
+        ## Anthropic does not mention using an Lp norm other than L1.
+        l1_coefficient=5,
+        lp_norm=1.0,
+        # Instead, they multiply the L1 loss contribution
+        # from each feature of the activations by the decoder norm of the corresponding feature.
+        scale_sparsity_penalty_by_decoder_norm=True,
+        # Learning Rate
+        lr_scheduler_name="constant",  # we set this independently of warmup and decay steps.
+        l1_warm_up_steps=l1_warmup_steps,
+        lr_warm_up_steps=lr_warm_up_steps,
+        lr_decay_steps=lr_warm_up_steps,
+        ## No ghost grad term.
+        use_ghost_grads=False,
+        # Initialization / Architecture
+        apply_b_dec_to_input=False,
+        # encoder bias zero's. (I'm not sure what it is by default now)
+        # decoder bias zero's.
+        b_dec_init_method="zeros",
+        normalize_sae_decoder=False,
+        decoder_heuristic_init=True,
+        init_encoder_as_decoder_transpose=True,
+        # Optimizer
+        lr=4e-5,
+        ## adam optimizer has no weight decay by default so worry about this.
+        adam_beta1=0.9,
+        adam_beta2=0.999,
+        # Buffer details won't matter in we cache / shuffle our activations ahead of time.
+        n_batches_in_buffer=64,
+        store_batch_size_prompts=16,
+        normalize_activations="none",
+        # Feature Store
+        feature_sampling_window=1000,
+        dead_feature_window=1000,
+        dead_feature_threshold=1e-4,
+        # performance enhancement:
+        compile_sae=False,
+        # WANDB
+        log_to_wandb=True,  # always use wandb unless you are just testing code.
+        wandb_project="benchmark",
+        wandb_log_frequency=100,
+        # Misc
+        device=device,
+        seed=42,
+        n_checkpoints=0,
+        checkpoint_path="checkpoints",
+        dtype="float32",
+    )
+
+    # look at the next cell to see some instruction for what to do while this is running.
+    sae = SAETrainingRunner(cfg).run()
+
+    assert sae is not None
+    # know whether or not this works by looking at the dashboard!
diff --git a/tests/unit/analysis/test_hooked_sae.py b/tests/unit/analysis/test_hooked_sae.py
index 6d37fac2..61f09a82 100644
--- a/tests/unit/analysis/test_hooked_sae.py
+++ b/tests/unit/analysis/test_hooked_sae.py
@@ -42,6 +42,7 @@ def get_hooked_sae(model: HookedTransformer, act_name: str) -> SAE:
     d_in = site_to_size[site]
 
     sae_cfg = SAEConfig(
+        architecture="standard",
         d_in=d_in,
         d_sae=d_in * 2,
         dtype="float32",
diff --git a/tests/unit/analysis/test_hooked_sae_transformer.py b/tests/unit/analysis/test_hooked_sae_transformer.py
index 100a3c38..2cc32e97 100644
--- a/tests/unit/analysis/test_hooked_sae_transformer.py
+++ b/tests/unit/analysis/test_hooked_sae_transformer.py
@@ -42,6 +42,7 @@ def get_hooked_sae(model: HookedTransformer, act_name: str) -> SAE:
     d_in = site_to_size[site]
 
     sae_cfg = SAEConfig(
+        architecture="standard",
         d_in=d_in,
         d_sae=d_in * 2,
         dtype="float32",
diff --git a/tests/unit/training/test_config.py b/tests/unit/training/test_config.py
index 2de701f4..81899610 100644
--- a/tests/unit/training/test_config.py
+++ b/tests/unit/training/test_config.py
@@ -45,6 +45,7 @@ def test_sae_training_runner_config_get_sae_base_parameters():
     cfg = LanguageModelSAERunnerConfig()
 
     expected_config = {
+        "architecture": "standard",
         "d_in": 512,
         "d_sae": 2048,
         "activation_fn_str": "relu",
diff --git a/tests/unit/training/test_gated_sae.py b/tests/unit/training/test_gated_sae.py
new file mode 100644
index 00000000..1afa894d
--- /dev/null
+++ b/tests/unit/training/test_gated_sae.py
@@ -0,0 +1,137 @@
+import pytest
+import torch
+
+from sae_lens.training.training_sae import TrainingSAE
+from tests.unit.helpers import build_sae_cfg
+
+
+def test_gated_sae_initialization():
+    cfg = build_sae_cfg()
+    setattr(cfg, "architecture", "gated")
+    sae = TrainingSAE.from_dict(cfg.get_training_sae_cfg_dict())
+
+    assert sae.W_enc.shape == (cfg.d_in, cfg.d_sae)
+    assert sae.W_dec.shape == (cfg.d_sae, cfg.d_in)
+    # assert sae.b_enc.shape == (cfg.d_sae,)
+    assert sae.b_mag.shape == (cfg.d_sae,)
+    assert sae.b_gate.shape == (cfg.d_sae,)
+    assert sae.r_mag.shape == (cfg.d_sae,)
+    assert sae.b_dec.shape == (cfg.d_in,)
+    assert isinstance(sae.activation_fn, torch.nn.ReLU)
+    assert sae.device == torch.device("cpu")
+    assert sae.dtype == torch.float32
+
+    # biases
+    assert torch.allclose(sae.b_dec, torch.zeros_like(sae.b_dec), atol=1e-6)
+    assert torch.allclose(sae.b_mag, torch.zeros_like(sae.b_mag), atol=1e-6)
+    assert torch.allclose(sae.b_gate, torch.zeros_like(sae.b_gate), atol=1e-6)
+
+    # check if the decoder weight norm is 1 by default
+    assert torch.allclose(
+        sae.W_dec.norm(dim=1), torch.ones_like(sae.W_dec.norm(dim=1)), atol=1e-6
+    )
+
+
+def test_gated_sae_encoding():
+    cfg = build_sae_cfg()
+    setattr(cfg, "architecture", "gated")
+    sae = TrainingSAE.from_dict(cfg.get_training_sae_cfg_dict())
+
+    batch_size = 32
+    d_in = sae.cfg.d_in
+    d_sae = sae.cfg.d_sae
+
+    x = torch.randn(batch_size, d_in)
+    feature_acts, hidden_pre = sae.encode_with_hidden_pre_gated(x)
+
+    assert feature_acts.shape == (batch_size, d_sae)
+    assert hidden_pre.shape == (batch_size, d_sae)
+
+    # Check the gating mechanism
+    gating_pre_activation = x @ sae.W_enc + sae.b_gate
+    active_features = (gating_pre_activation > 0).float()
+    magnitude_pre_activation = x @ (sae.W_enc * sae.r_mag.exp()) + sae.b_mag
+    feature_magnitudes = sae.activation_fn(magnitude_pre_activation)
+
+    expected_feature_acts = active_features * feature_magnitudes
+    assert torch.allclose(feature_acts, expected_feature_acts, atol=1e-6)
+
+
+def test_gated_sae_loss():
+    cfg = build_sae_cfg()
+    setattr(cfg, "architecture", "gated")
+    sae = TrainingSAE.from_dict(cfg.get_training_sae_cfg_dict())
+
+    batch_size = 32
+    d_in = sae.cfg.d_in
+    x = torch.randn(batch_size, d_in)
+
+    train_step_output = sae.training_forward_pass(
+        sae_in=x,
+        current_l1_coefficient=sae.cfg.l1_coefficient,
+    )
+
+    assert train_step_output.sae_out.shape == (batch_size, d_in)
+    assert train_step_output.feature_acts.shape == (batch_size, sae.cfg.d_sae)
+
+    sae_in_centered = x - sae.b_dec
+    via_gate_feature_magnitudes = torch.relu(sae_in_centered @ sae.W_enc + sae.b_gate)
+    preactivation_l1_loss = (
+        sae.cfg.l1_coefficient * torch.sum(via_gate_feature_magnitudes, dim=-1).mean()
+    )
+
+    via_gate_reconstruction = (
+        via_gate_feature_magnitudes @ sae.W_dec.detach() + sae.b_dec.detach()
+    )
+    aux_reconstruction_loss = torch.sum(
+        (via_gate_reconstruction - x) ** 2, dim=-1
+    ).mean()
+
+    expected_loss = (
+        train_step_output.mse_loss + preactivation_l1_loss + aux_reconstruction_loss
+    )
+    assert (
+        pytest.approx(train_step_output.loss.item(), rel=1e-3) == expected_loss.item()
+    )
+
+
+def test_gated_sae_forward_pass():
+    cfg = build_sae_cfg()
+    setattr(cfg, "architecture", "gated")
+    sae = TrainingSAE.from_dict(cfg.get_training_sae_cfg_dict())
+
+    batch_size = 32
+    d_in = sae.cfg.d_in
+
+    x = torch.randn(batch_size, d_in)
+    sae_out = sae(x)
+
+    assert sae_out.shape == (batch_size, d_in)
+
+
+def test_gated_sae_training_forward_pass():
+    cfg = build_sae_cfg()
+    setattr(cfg, "architecture", "gated")
+    sae = TrainingSAE.from_dict(cfg.get_training_sae_cfg_dict())
+
+    batch_size = 32
+    d_in = sae.cfg.d_in
+
+    x = torch.randn(batch_size, d_in)
+    train_step_output = sae.training_forward_pass(
+        sae_in=x,
+        current_l1_coefficient=sae.cfg.l1_coefficient,
+    )
+
+    assert train_step_output.sae_out.shape == (batch_size, d_in)
+    assert train_step_output.feature_acts.shape == (batch_size, sae.cfg.d_sae)
+
+    # Detach the loss tensor and convert to numpy for comparison
+    detached_loss = train_step_output.loss.detach().cpu().numpy()
+    expected_loss = (
+        train_step_output.mse_loss
+        + train_step_output.l1_loss
+        + train_step_output.auxiliary_reconstruction_loss
+    )
+
+    assert pytest.approx(detached_loss, rel=1e-3) == expected_loss
diff --git a/tests/unit/training/test_sae_initialization.py b/tests/unit/training/test_sae_initialization.py
index 9c8eaade..fe73ee23 100644
--- a/tests/unit/training/test_sae_initialization.py
+++ b/tests/unit/training/test_sae_initialization.py
@@ -33,6 +33,37 @@ def test_SparseAutoencoder_initialization_standard():
     assert not torch.allclose(unit_normed_W_enc, unit_normed_W_dec, atol=1e-6)
 
 
+def test_SparseAutoencoder_initialization_gated():
+    cfg = build_sae_cfg()
+    setattr(cfg, "architecture", "gated")
+    sae = TrainingSAE.from_dict(cfg.get_training_sae_cfg_dict())
+
+    assert sae.W_enc.shape == (cfg.d_in, cfg.d_sae)
+    assert sae.W_dec.shape == (cfg.d_sae, cfg.d_in)
+    assert sae.b_mag.shape == (cfg.d_sae,)
+    assert sae.b_gate.shape == (cfg.d_sae,)
+    assert sae.r_mag.shape == (cfg.d_sae,)
+    assert sae.b_dec.shape == (cfg.d_in,)
+    assert isinstance(sae.activation_fn, torch.nn.ReLU)
+    assert sae.device == torch.device("cpu")
+    assert sae.dtype == torch.float32
+
+    # biases
+    assert torch.allclose(sae.b_dec, torch.zeros_like(sae.b_dec), atol=1e-6)
+    assert torch.allclose(sae.b_mag, torch.zeros_like(sae.b_mag), atol=1e-6)
+    assert torch.allclose(sae.b_gate, torch.zeros_like(sae.b_gate), atol=1e-6)
+
+    # check if the decoder weight norm is 1 by default
+    assert torch.allclose(
+        sae.W_dec.norm(dim=1), torch.ones_like(sae.W_dec.norm(dim=1)), atol=1e-6
+    )
+
+    #  Default currently shouldn't be tranpose initialization
+    unit_normed_W_enc = sae.W_enc / torch.norm(sae.W_enc, dim=0)
+    unit_normed_W_dec = sae.W_dec.T
+    assert not torch.allclose(unit_normed_W_enc, unit_normed_W_dec, atol=1e-6)
+
+
 def test_SparseAutoencoder_initialization_orthogonal_enc_dec():
     cfg = build_sae_cfg(decoder_orthogonal_init=True)
 
diff --git a/tests/unit/training/test_sae_trainer.py b/tests/unit/training/test_sae_trainer.py
index adb3a8fe..2573140e 100644
--- a/tests/unit/training/test_sae_trainer.py
+++ b/tests/unit/training/test_sae_trainer.py
@@ -185,6 +185,7 @@ def test_build_train_step_log_dict(trainer: SAETrainer) -> None:
         # l1 loss is scaled by l1_coefficient
         "losses/l1_loss": train_output.l1_loss / trainer.cfg.l1_coefficient,
         "losses/ghost_grad_loss": pytest.approx(0.15),
+        "losses/auxiliary_reconstruction_loss": 0.0,
         "losses/overall_loss": 0.5,
         "metrics/explained_variance": 0.75,
         "metrics/explained_variance_std": 0.25,
diff --git a/tutorials/training_a_gated_sae.ipynb b/tutorials/training_a_gated_sae.ipynb
new file mode 100644
index 00000000..ff319f9f
--- /dev/null
+++ b/tutorials/training_a_gated_sae.ipynb
@@ -0,0 +1,700 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5O8tQblzOVHu"
+      },
+      "source": [
+        "# A Very Basic Gated SAE Training Run"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "shAFb9-lOVHu"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "LeRi_tw2dhae"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Requirement already satisfied: sae-lens in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (3.3.0)\n",
+            "Requirement already satisfied: transformer-lens in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (1.19.0)\n",
+            "Requirement already satisfied: circuitsvis in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (1.43.2)\n",
+            "Requirement already satisfied: automated-interpretability<0.0.4,>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.0.3)\n",
+            "Requirement already satisfied: babe<0.0.8,>=0.0.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.0.7)\n",
+            "Requirement already satisfied: datasets<3.0.0,>=2.17.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (2.19.2)\n",
+            "Requirement already satisfied: matplotlib<4.0.0,>=3.8.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (3.9.0)\n",
+            "Requirement already satisfied: matplotlib-inline<0.2.0,>=0.1.6 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.1.7)\n",
+            "Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (3.8.1)\n",
+            "Requirement already satisfied: plotly<6.0.0,>=5.19.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (5.22.0)\n",
+            "Requirement already satisfied: plotly-express<0.5.0,>=0.4.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.4.1)\n",
+            "Requirement already satisfied: pytest-profiling<2.0.0,>=1.7.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (1.7.0)\n",
+            "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (1.0.1)\n",
+            "Requirement already satisfied: pyyaml<7.0.0,>=6.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (6.0.1)\n",
+            "Requirement already satisfied: pyzmq==26.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (26.0.0)\n",
+            "Requirement already satisfied: sae-vis<0.3.0,>=0.2.18 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.2.18)\n",
+            "Requirement already satisfied: safetensors<0.5.0,>=0.4.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.4.3)\n",
+            "Requirement already satisfied: transformers<5.0.0,>=4.38.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (4.41.2)\n",
+            "Requirement already satisfied: typer<0.13.0,>=0.12.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.12.3)\n",
+            "Requirement already satisfied: accelerate>=0.23.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.31.0)\n",
+            "Requirement already satisfied: beartype<0.15.0,>=0.14.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.14.1)\n",
+            "Requirement already satisfied: better-abc<0.0.4,>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.0.3)\n",
+            "Requirement already satisfied: einops>=0.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.7.0)\n",
+            "Requirement already satisfied: fancy-einsum>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.0.3)\n",
+            "Requirement already satisfied: jaxtyping>=0.2.11 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.2.29)\n",
+            "Requirement already satisfied: numpy>=1.24 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (1.26.4)\n",
+            "Requirement already satisfied: pandas>=1.1.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (2.2.2)\n",
+            "Requirement already satisfied: rich>=12.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (13.7.1)\n",
+            "Requirement already satisfied: sentencepiece in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.2.0)\n",
+            "Requirement already satisfied: torch>=1.10 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (2.1.2)\n",
+            "Requirement already satisfied: tqdm>=4.64.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (4.66.4)\n",
+            "Requirement already satisfied: typing-extensions in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (4.12.2)\n",
+            "Requirement already satisfied: wandb>=0.13.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.17.1)\n",
+            "Requirement already satisfied: importlib-metadata>=5.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (7.1.0)\n",
+            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.3.1)\n",
+            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (8.9.2.26)\n",
+            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (11.0.2.54)\n",
+            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (10.3.2.106)\n",
+            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (11.4.5.107)\n",
+            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.0.106)\n",
+            "Requirement already satisfied: nvidia-nccl-cu12==2.18.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (2.18.1)\n",
+            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: triton==2.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (2.1.0)\n",
+            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->circuitsvis) (12.5.40)\n",
+            "Requirement already satisfied: filelock in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from triton==2.1.0->circuitsvis) (3.14.0)\n",
+            "Requirement already satisfied: packaging>=20.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (24.0)\n",
+            "Requirement already satisfied: psutil in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (5.9.0)\n",
+            "Requirement already satisfied: huggingface-hub in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (0.23.3)\n",
+            "Requirement already satisfied: blobfile<3.0.0,>=2.1.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.1.1)\n",
+            "Requirement already satisfied: boostedblob<0.16.0,>=0.15.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.15.3)\n",
+            "Requirement already satisfied: httpx<0.28.0,>=0.27.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.27.0)\n",
+            "Requirement already satisfied: orjson<4.0.0,>=3.10.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.10.4)\n",
+            "Requirement already satisfied: pytest<9.0.0,>=8.1.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (8.2.2)\n",
+            "Requirement already satisfied: scikit-learn<2.0.0,>=1.4.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.5.0)\n",
+            "Requirement already satisfied: tiktoken<0.7.0,>=0.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.6.0)\n",
+            "Requirement already satisfied: py2store in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from babe<0.0.8,>=0.0.7->sae-lens) (0.1.20)\n",
+            "Requirement already satisfied: graze in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from babe<0.0.8,>=0.0.7->sae-lens) (0.1.17)\n",
+            "Requirement already satisfied: pyarrow>=12.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (16.1.0)\n",
+            "Requirement already satisfied: pyarrow-hotfix in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.6)\n",
+            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.3.8)\n",
+            "Requirement already satisfied: requests>=2.32.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (2.32.3)\n",
+            "Requirement already satisfied: xxhash in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (3.4.1)\n",
+            "Requirement already satisfied: multiprocess in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.70.16)\n",
+            "Requirement already satisfied: fsspec<=2024.3.1,>=2023.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from fsspec[http]<=2024.3.1,>=2023.1.0->datasets<3.0.0,>=2.17.1->sae-lens) (2024.3.1)\n",
+            "Requirement already satisfied: aiohttp in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (3.9.5)\n",
+            "Requirement already satisfied: zipp>=0.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from importlib-metadata>=5.1.0->circuitsvis) (3.19.2)\n",
+            "Requirement already satisfied: typeguard==2.13.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from jaxtyping>=0.2.11->transformer-lens) (2.13.3)\n",
+            "Requirement already satisfied: contourpy>=1.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (1.2.1)\n",
+            "Requirement already satisfied: cycler>=0.10 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (0.12.1)\n",
+            "Requirement already satisfied: fonttools>=4.22.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (4.53.0)\n",
+            "Requirement already satisfied: kiwisolver>=1.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (1.4.5)\n",
+            "Requirement already satisfied: pillow>=8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (10.3.0)\n",
+            "Requirement already satisfied: pyparsing>=2.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (3.1.2)\n",
+            "Requirement already satisfied: python-dateutil>=2.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (2.9.0)\n",
+            "Requirement already satisfied: traitlets in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib-inline<0.2.0,>=0.1.6->sae-lens) (5.14.3)\n",
+            "Requirement already satisfied: click in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (8.1.7)\n",
+            "Requirement already satisfied: joblib in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (1.4.2)\n",
+            "Requirement already satisfied: regex>=2021.8.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (2024.5.15)\n",
+            "Requirement already satisfied: pytz>=2020.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pandas>=1.1.5->transformer-lens) (2024.1)\n",
+            "Requirement already satisfied: tzdata>=2022.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pandas>=1.1.5->transformer-lens) (2024.1)\n",
+            "Requirement already satisfied: tenacity>=6.2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly<6.0.0,>=5.19.0->sae-lens) (8.3.0)\n",
+            "Requirement already satisfied: statsmodels>=0.9.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (0.14.2)\n",
+            "Requirement already satisfied: scipy>=0.18 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (1.13.1)\n",
+            "Requirement already satisfied: patsy>=0.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (0.5.6)\n",
+            "Requirement already satisfied: six in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens) (1.16.0)\n",
+            "Requirement already satisfied: gprof2dot in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens) (2024.6.6)\n",
+            "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from rich>=12.6.0->transformer-lens) (3.0.0)\n",
+            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from rich>=12.6.0->transformer-lens) (2.18.0)\n",
+            "Requirement already satisfied: dataclasses-json<0.7.0,>=0.6.4 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens) (0.6.7)\n",
+            "Requirement already satisfied: eindex-callum<0.2.0,>=0.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens) (0.1.1)\n",
+            "Requirement already satisfied: sympy in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (1.12.1)\n",
+            "Requirement already satisfied: networkx in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (3.3)\n",
+            "Requirement already satisfied: jinja2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (3.1.4)\n",
+            "Requirement already satisfied: tokenizers<0.20,>=0.19 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformers<5.0.0,>=4.38.1->sae-lens) (0.19.1)\n",
+            "Requirement already satisfied: shellingham>=1.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from typer<0.13.0,>=0.12.3->sae-lens) (1.5.4)\n",
+            "Requirement already satisfied: docker-pycreds>=0.4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (0.4.0)\n",
+            "Requirement already satisfied: gitpython!=3.1.29,>=1.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (3.1.43)\n",
+            "Requirement already satisfied: platformdirs in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (4.2.2)\n",
+            "Requirement already satisfied: protobuf!=4.21.0,<6,>=3.19.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (5.27.1)\n",
+            "Requirement already satisfied: sentry-sdk>=1.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (2.5.1)\n",
+            "Requirement already satisfied: setproctitle in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (1.3.3)\n",
+            "Requirement already satisfied: setuptools in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (70.0.0)\n",
+            "Requirement already satisfied: pycryptodomex~=3.8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.20.0)\n",
+            "Requirement already satisfied: urllib3<3,>=1.25.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.2.1)\n",
+            "Requirement already satisfied: lxml~=4.9 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (4.9.4)\n",
+            "Requirement already satisfied: uvloop>=0.16.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from boostedblob<0.16.0,>=0.15.3->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.19.0)\n",
+            "Requirement already satisfied: aiosignal>=1.1.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.3.1)\n",
+            "Requirement already satisfied: attrs>=17.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (23.2.0)\n",
+            "Requirement already satisfied: frozenlist>=1.1.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.4.1)\n",
+            "Requirement already satisfied: multidict<7.0,>=4.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (6.0.5)\n",
+            "Requirement already satisfied: yarl<2.0,>=1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.9.4)\n",
+            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (4.0.3)\n",
+            "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (3.21.3)\n",
+            "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (0.9.0)\n",
+            "Requirement already satisfied: gitdb<5,>=4.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer-lens) (4.0.11)\n",
+            "Requirement already satisfied: anyio in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (4.4.0)\n",
+            "Requirement already satisfied: certifi in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2024.6.2)\n",
+            "Requirement already satisfied: httpcore==1.* in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.0.5)\n",
+            "Requirement already satisfied: idna in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.7)\n",
+            "Requirement already satisfied: sniffio in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.3.1)\n",
+            "Requirement already satisfied: h11<0.15,>=0.13 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpcore==1.*->httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.14.0)\n",
+            "Requirement already satisfied: mdurl~=0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich>=12.6.0->transformer-lens) (0.1.2)\n",
+            "Requirement already satisfied: iniconfig in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.0.0)\n",
+            "Requirement already satisfied: pluggy<2.0,>=1.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.5.0)\n",
+            "Requirement already satisfied: exceptiongroup>=1.0.0rc8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.2.0)\n",
+            "Requirement already satisfied: tomli>=1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.0.1)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from requests>=2.32.1->datasets<3.0.0,>=2.17.1->sae-lens) (3.3.2)\n",
+            "Requirement already satisfied: threadpoolctl>=3.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from scikit-learn<2.0.0,>=1.4.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.5.0)\n",
+            "Requirement already satisfied: dol in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from graze->babe<0.0.8,>=0.0.7->sae-lens) (0.2.47)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from jinja2->torch>=1.10->transformer-lens) (2.1.5)\n",
+            "Requirement already satisfied: config2py in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens) (0.1.33)\n",
+            "Requirement already satisfied: importlib-resources in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens) (6.4.0)\n",
+            "Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sympy->torch>=1.10->transformer-lens) (1.3.0)\n",
+            "Requirement already satisfied: smmap<6,>=3.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer-lens) (5.0.1)\n",
+            "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (1.0.0)\n",
+            "Requirement already satisfied: i2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from config2py->py2store->babe<0.0.8,>=0.0.7->sae-lens) (0.1.17)\n",
+            "Note: you may need to restart the kernel to use updated packages.\n"
+          ]
+        }
+      ],
+      "source": [
+        "try:\n",
+        "    #import google.colab # type: ignore\n",
+        "    #from google.colab import output\n",
+        "    %pip install sae-lens transformer-lens circuitsvis\n",
+        "except:\n",
+        "    from IPython import get_ipython # type: ignore\n",
+        "    ipython = get_ipython(); assert ipython is not None\n",
+        "    ipython.run_line_magic(\"load_ext\", \"autoreload\")\n",
+        "    ipython.run_line_magic(\"autoreload\", \"2\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "uy-b3CcSOVHu",
+        "outputId": "58ce28d0-f91f-436d-cf87-76bb26e2ecaf"
+      },
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "  from .autonotebook import tqdm as notebook_tqdm\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Using device: cuda\n"
+          ]
+        }
+      ],
+      "source": [
+        "import torch\n",
+        "import os\n",
+        "\n",
+        "from sae_lens import LanguageModelSAERunnerConfig, SAETrainingRunner\n",
+        "\n",
+        "if torch.cuda.is_available():\n",
+        "    device = \"cuda\"\n",
+        "elif torch.backends.mps.is_available():\n",
+        "    device = \"mps\"\n",
+        "else:\n",
+        "    device = \"cpu\"\n",
+        "\n",
+        "print(\"Using device:\", device)\n",
+        "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jCHtPycOOVHw"
+      },
+      "source": [
+        "## Training on MLP Out"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "oAsZCAdJOVHw"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Run name: 16384-L1-5-LR-5e-05-Tokens-1.229e+08\n",
+            "n_tokens_per_buffer (millions): 0.262144\n",
+            "Lower bound: n_contexts_per_buffer (millions): 0.001024\n",
+            "Total training steps: 30000\n",
+            "Total wandb updates: 1000\n",
+            "n_tokens_per_feature_sampling_window (millions): 1048.576\n",
+            "n_tokens_per_dead_feature_window (millions): 1048.576\n",
+            "We will reset the sparsity calculation 30 times.\n",
+            "Number tokens in sparsity calculation window: 4.10e+06\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
+            "  warnings.warn(\n",
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
+            "  return self.fget.__get__(instance, owner)()\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Loaded pretrained model tiny-stories-1L-21M into HookedTransformer\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcurt-tigges\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "Tracking run with wandb version 0.17.1"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Run data is saved locally in <code>/home/curttigges/projects/SAELens/tutorials/wandb/run-20240611_143204-n7cy5v24</code>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Syncing run <strong><a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/n7cy5v24' target=\"_blank\">16384-L1-5-LR-5e-05-Tokens-1.229e+08</a></strong> to <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View project at <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial</a>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View run at <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/n7cy5v24' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/n7cy5v24</a>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Estimating norm scaling factor: 100%|██████████| 1000/1000 [00:33<00:00, 30.13it/s]\n",
+            "5500| MSE Loss 208.944 | L1 167.607:   0%|          | 225280/122880000 [08:05<71:26:53, 476.86it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "interrupted, saving progress\n",
+            "done saving\n"
+          ]
+        },
+        {
+          "ename": "InterruptedException",
+          "evalue": "",
+          "output_type": "error",
+          "traceback": [
+            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+            "\u001b[0;31mInterruptedException\u001b[0m                      Traceback (most recent call last)",
+            "Cell \u001b[0;32mIn[3], line 63\u001b[0m\n\u001b[1;32m      9\u001b[0m cfg \u001b[38;5;241m=\u001b[39m LanguageModelSAERunnerConfig(\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;66;03m# Data Generating Function (Model + Training Distribution)\u001b[39;00m\n\u001b[1;32m     11\u001b[0m     variant\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbaseline\u001b[39m\u001b[38;5;124m\"\u001b[39m,  \u001b[38;5;66;03m# we'll use the gated variant.\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     60\u001b[0m     dtype\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     61\u001b[0m )\n\u001b[1;32m     62\u001b[0m \u001b[38;5;66;03m# look at the next cell to see some instruction for what to do while this is running.\u001b[39;00m\n\u001b[0;32m---> 63\u001b[0m sparse_autoencoder \u001b[38;5;241m=\u001b[39m \u001b[43mSAETrainingRunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcfg\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+            "File \u001b[0;32m~/projects/SAELens/sae_lens/sae_training_runner.py:87\u001b[0m, in \u001b[0;36mSAETrainingRunner.run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     78\u001b[0m trainer \u001b[38;5;241m=\u001b[39m SAETrainer(\n\u001b[1;32m     79\u001b[0m     model\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodel,\n\u001b[1;32m     80\u001b[0m     sae\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msae,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     83\u001b[0m     cfg\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcfg,\n\u001b[1;32m     84\u001b[0m )\n\u001b[1;32m     86\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compile_if_needed()\n\u001b[0;32m---> 87\u001b[0m sae \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_trainer_with_interruption_handling\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrainer\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     89\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcfg\u001b[38;5;241m.\u001b[39mlog_to_wandb:\n\u001b[1;32m     90\u001b[0m     wandb\u001b[38;5;241m.\u001b[39mfinish()\n",
+            "File \u001b[0;32m~/projects/SAELens/sae_lens/sae_training_runner.py:130\u001b[0m, in \u001b[0;36mSAETrainingRunner.run_trainer_with_interruption_handling\u001b[0;34m(self, trainer)\u001b[0m\n\u001b[1;32m    127\u001b[0m     signal\u001b[38;5;241m.\u001b[39msignal(signal\u001b[38;5;241m.\u001b[39mSIGTERM, interrupt_callback)\n\u001b[1;32m    129\u001b[0m     \u001b[38;5;66;03m# train SAE\u001b[39;00m\n\u001b[0;32m--> 130\u001b[0m     sae \u001b[38;5;241m=\u001b[39m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfit\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    132\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m, InterruptedException):\n\u001b[1;32m    133\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minterrupted, saving progress\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+            "File \u001b[0;32m~/projects/SAELens/sae_lens/training/sae_trainer.py:162\u001b[0m, in \u001b[0;36mSAETrainer.fit\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    159\u001b[0m layer_acts \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactivation_store\u001b[38;5;241m.\u001b[39mnext_batch()[:, \u001b[38;5;241m0\u001b[39m, :]\n\u001b[1;32m    160\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mn_training_tokens \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcfg\u001b[38;5;241m.\u001b[39mtrain_batch_size_tokens\n\u001b[0;32m--> 162\u001b[0m step_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_train_step\u001b[49m\u001b[43m(\u001b[49m\u001b[43msae\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msae\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msae_in\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlayer_acts\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    164\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcfg\u001b[38;5;241m.\u001b[39mlog_to_wandb:\n\u001b[1;32m    165\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_log_train_step(step_output)\n",
+            "File \u001b[0;32m~/projects/SAELens/sae_lens/training/sae_trainer.py:216\u001b[0m, in \u001b[0;36mSAETrainer._train_step\u001b[0;34m(self, sae, sae_in)\u001b[0m\n\u001b[1;32m    212\u001b[0m \u001b[38;5;66;03m# for documentation on autocasting see:\u001b[39;00m\n\u001b[1;32m    213\u001b[0m \u001b[38;5;66;03m# https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html\u001b[39;00m\n\u001b[1;32m    214\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mautocast_if_enabled:\n\u001b[0;32m--> 216\u001b[0m     train_step_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msae\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtraining_forward_pass\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    217\u001b[0m \u001b[43m        \u001b[49m\u001b[43msae_in\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msae_in\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    218\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdead_neuron_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdead_neurons\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    219\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcurrent_l1_coefficient\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcurrent_l1_coefficient\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    220\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    222\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[1;32m    223\u001b[0m         did_fire \u001b[38;5;241m=\u001b[39m (train_step_output\u001b[38;5;241m.\u001b[39mfeature_acts \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m.\u001b[39mfloat()\u001b[38;5;241m.\u001b[39msum(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m2\u001b[39m) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m\n",
+            "File \u001b[0;32m~/projects/SAELens/sae_lens/training/training_sae.py:303\u001b[0m, in \u001b[0;36mTrainingSAE.training_forward_pass\u001b[0;34m(self, sae_in, current_l1_coefficient, dead_neuron_mask)\u001b[0m\n\u001b[1;32m    295\u001b[0m     l1_loss \u001b[38;5;241m=\u001b[39m (current_l1_coefficient \u001b[38;5;241m*\u001b[39m sparsity)\u001b[38;5;241m.\u001b[39mmean()\n\u001b[1;32m    296\u001b[0m     loss \u001b[38;5;241m=\u001b[39m mse_loss \u001b[38;5;241m+\u001b[39m l1_loss \u001b[38;5;241m+\u001b[39m ghost_grad_loss\n\u001b[1;32m    298\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m TrainStepOutput(\n\u001b[1;32m    299\u001b[0m     sae_in\u001b[38;5;241m=\u001b[39msae_in,\n\u001b[1;32m    300\u001b[0m     sae_out\u001b[38;5;241m=\u001b[39msae_out,\n\u001b[1;32m    301\u001b[0m     feature_acts\u001b[38;5;241m=\u001b[39mfeature_acts,\n\u001b[1;32m    302\u001b[0m     loss\u001b[38;5;241m=\u001b[39mloss,\n\u001b[0;32m--> 303\u001b[0m     mse_loss\u001b[38;5;241m=\u001b[39m\u001b[43mmse_loss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m    304\u001b[0m     l1_loss\u001b[38;5;241m=\u001b[39ml1_loss\u001b[38;5;241m.\u001b[39mitem(),\n\u001b[1;32m    305\u001b[0m     ghost_grad_loss\u001b[38;5;241m=\u001b[39m(\n\u001b[1;32m    306\u001b[0m         ghost_grad_loss\u001b[38;5;241m.\u001b[39mitem()\n\u001b[1;32m    307\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(ghost_grad_loss, torch\u001b[38;5;241m.\u001b[39mTensor)\n\u001b[1;32m    308\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m ghost_grad_loss\n\u001b[1;32m    309\u001b[0m     ),\n\u001b[1;32m    310\u001b[0m )\n",
+            "File \u001b[0;32m~/projects/SAELens/sae_lens/sae_training_runner.py:25\u001b[0m, in \u001b[0;36minterrupt_callback\u001b[0;34m(sig_num, stack_frame)\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21minterrupt_callback\u001b[39m(sig_num: Any, stack_frame: Any):\n\u001b[0;32m---> 25\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m InterruptedException()\n",
+            "\u001b[0;31mInterruptedException\u001b[0m: "
+          ]
+        }
+      ],
+      "source": [
+        "total_training_steps = 30_000  # probably we should do more\n",
+        "batch_size = 4096\n",
+        "total_training_tokens = total_training_steps * batch_size\n",
+        "\n",
+        "lr_warm_up_steps = 0\n",
+        "lr_decay_steps = total_training_steps // 5  # 20% of training\n",
+        "l1_warm_up_steps = total_training_steps // 20  # 5% of training\n",
+        "\n",
+        "cfg = LanguageModelSAERunnerConfig(\n",
+        "    # Data Generating Function (Model + Training Distribution)\n",
+        "    architecture=\"baseline\",  # we'll use the gated variant.\n",
+        "    model_name=\"tiny-stories-1L-21M\",  # our model (more options here: https://neelnanda-io.github.io/TransformerLens/generated/model_properties_table.html)\n",
+        "    hook_name=\"blocks.0.hook_mlp_out\",  # A valid hook point (see more details here: https://neelnanda-io.github.io/TransformerLens/generated/demos/Main_Demo.html#Hook-Points)\n",
+        "    hook_layer=0,  # Only one layer in the model.\n",
+        "    d_in=1024,  # the width of the mlp output.\n",
+        "    dataset_path=\"apollo-research/roneneldan-TinyStories-tokenizer-gpt2\",  # this is a tokenized language dataset on Huggingface for the Tiny Stories corpus.\n",
+        "    is_dataset_tokenized=True,\n",
+        "    streaming=True,  # we could pre-download the token dataset if it was small.\n",
+        "    # SAE Parameters\n",
+        "    mse_loss_normalization=None,  # We won't normalize the mse loss,\n",
+        "    expansion_factor=16,  # the width of the SAE. Larger will result in better stats but slower training.\n",
+        "    b_dec_init_method=\"zeros\",  # The geometric median can be used to initialize the decoder weights.\n",
+        "    apply_b_dec_to_input=False,  # We won't apply the decoder weights to the input.\n",
+        "    normalize_sae_decoder=False,\n",
+        "    scale_sparsity_penalty_by_decoder_norm=True,\n",
+        "    decoder_heuristic_init=True,\n",
+        "    init_encoder_as_decoder_transpose=True,\n",
+        "    normalize_activations=True,\n",
+        "    # Training Parameters\n",
+        "    lr=5e-5,  # lower the better, we'll go fairly high to speed up the tutorial.\n",
+        "    adam_beta1=0.9,  # adam params (default, but once upon a time we experimented with these.)\n",
+        "    adam_beta2=0.999,\n",
+        "    lr_scheduler_name=\"constant\",  # constant learning rate with warmup. Could be better schedules out there.\n",
+        "    lr_warm_up_steps=lr_warm_up_steps,  # this can help avoid too many dead features initially.\n",
+        "    lr_decay_steps=lr_decay_steps,  # this will help us avoid overfitting.\n",
+        "    l1_coefficient=5,  # will control how sparse the feature activations are\n",
+        "    l1_warm_up_steps=l1_warm_up_steps,  # this can help avoid too many dead features initially.\n",
+        "    lp_norm=1.0,  # the L1 penalty (and not a Lp for p < 1)\n",
+        "    train_batch_size_tokens=batch_size,\n",
+        "    context_size=256,  # will control the lenght of the prompts we feed to the model. Larger is better but slower. so for the tutorial we'll use a short one.\n",
+        "    # Activation Store Parameters\n",
+        "    n_batches_in_buffer=64,  # controls how many activations we store / shuffle.\n",
+        "    training_tokens=total_training_tokens,  # 100 million tokens is quite a few, but we want to see good stats. Get a coffee, come back.\n",
+        "    store_batch_size_prompts=16,\n",
+        "    # Resampling protocol\n",
+        "    use_ghost_grads=False,  # we don't use ghost grads anymore.\n",
+        "    feature_sampling_window=1000,  # this controls our reporting of feature sparsity stats\n",
+        "    dead_feature_window=1000,  # would effect resampling or ghost grads if we were using it.\n",
+        "    dead_feature_threshold=1e-4,  # would effect resampling or ghost grads if we were using it.\n",
+        "    # WANDB\n",
+        "    log_to_wandb=True,  # always use wandb unless you are just testing code.\n",
+        "    wandb_project=\"sae_lens_tutorial\",\n",
+        "    wandb_log_frequency=30,\n",
+        "    eval_every_n_wandb_logs=20,\n",
+        "    # Misc\n",
+        "    device=device,\n",
+        "    seed=42,\n",
+        "    n_checkpoints=0,\n",
+        "    checkpoint_path=\"checkpoints\",\n",
+        "    dtype=\"float32\"\n",
+        ")\n",
+        "# look at the next cell to see some instruction for what to do while this is running.\n",
+        "sparse_autoencoder = SAETrainingRunner(cfg).run()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {},
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Run name: 16384-L1-20-LR-5e-05-Tokens-1.229e+08\n",
+            "n_tokens_per_buffer (millions): 0.262144\n",
+            "Lower bound: n_contexts_per_buffer (millions): 0.001024\n",
+            "Total training steps: 30000\n",
+            "Total wandb updates: 1000\n",
+            "n_tokens_per_feature_sampling_window (millions): 1048.576\n",
+            "n_tokens_per_dead_feature_window (millions): 1048.576\n",
+            "We will reset the sparsity calculation 30 times.\n",
+            "Number tokens in sparsity calculation window: 4.10e+06\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
+            "  warnings.warn(\n",
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
+            "  return self.fget.__get__(instance, owner)()\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Loaded pretrained model tiny-stories-1L-21M into HookedTransformer\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: W&B API key is configured. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "Tracking run with wandb version 0.17.1"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Run data is saved locally in <code>/home/curttigges/projects/SAELens/tutorials/wandb/run-20240616_143959-ch6e0a5s</code>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Syncing run <strong><a href='https://wandb.ai/curt-tigges/gated_sae_testing/runs/ch6e0a5s' target=\"_blank\">16384-L1-20-LR-5e-05-Tokens-1.229e+08</a></strong> to <a href='https://wandb.ai/curt-tigges/gated_sae_testing' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View project at <a href='https://wandb.ai/curt-tigges/gated_sae_testing' target=\"_blank\">https://wandb.ai/curt-tigges/gated_sae_testing</a>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View run at <a href='https://wandb.ai/curt-tigges/gated_sae_testing/runs/ch6e0a5s' target=\"_blank\">https://wandb.ai/curt-tigges/gated_sae_testing/runs/ch6e0a5s</a>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "30000| MSE Loss 143.062 | L1 0.000:   1%|          | 1228800/122880000 [1:04:38<106:39:53, 316.81it/s]\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "<style>\n",
+              "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
+              "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
+              "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
+              "    </style>\n",
+              "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>details/current_l1_coefficient</td><td>▁▂▂▃▃▄▄▅▅▆▆▇████████████████████████████</td></tr><tr><td>details/current_learning_rate</td><td>████████████████████████████████▇▇▅▅▄▃▂▁</td></tr><tr><td>details/n_training_tokens</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>losses/auxiliary_reconstruction_loss</td><td>▁▃▃▄▄▅▅▅▆▆▆▇▇███████████████████████████</td></tr><tr><td>losses/ghost_grad_loss</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/l1_loss</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/mse_loss</td><td>▁▂▃▄▄▅▅▅▆▆▆▆▆▆▇▆▆▆▆▆▆▇▆▆▇▆▆▆▇▆▇▆▆▆▆▇▆▇▆█</td></tr><tr><td>losses/overall_loss</td><td>▁▃▄▄▅▆▆▆▇▇▇▇██████████████▇████▇▇█▇█▇█▇█</td></tr><tr><td>losses/sfn_sparsity_loss</td><td>▂▃▅▆▆▇████▇▆▅▄▃▃▃▃▃▃▃▃▄▃▃▃▂▃▃▃▃▃▃▂▂▂▁▁▁▁</td></tr><tr><td>metrics/CE_loss_score</td><td>██▇▇▆▆▆▆▆▆▆▆▆▅▆▆▅▆▅▃▅▁▆▅▆▆▅▅▃▆▅▅▆▅▆▆▅▅▃▄</td></tr><tr><td>metrics/ce_loss_with_ablation</td><td>▅▃▁▆▄▃▅▆▅▄▄▁▅▃▄▃▃▄▁▃▄▆▄▄▄▃▆▄▃█▄▄▁▄▄▅█▄▄▃</td></tr><tr><td>metrics/ce_loss_with_sae</td><td>▁▁▂▂▃▂▃▃▃▃▃▃▃▃▃▃▄▃▄▆▄█▃▄▃▃▄▄▆▃▃▄▃▄▃▃▄▄▆▅</td></tr><tr><td>metrics/ce_loss_without_sae</td><td>▄▂▃█▃▁▃▃▄▅▃▄▂▃▄▃▄▃▃▃▃▃▄▁▂▂▆▃▄▅▃▄▃▅▇▅▃▃▃▂</td></tr><tr><td>metrics/explained_variance</td><td>█▇▆▅▅▄▄▄▃▃▃▃▃▃▂▃▃▃▃▃▃▂▃▃▃▃▃▃▂▃▃▃▃▃▃▂▃▂▃▁</td></tr><tr><td>metrics/explained_variance_std</td><td>▁▂▃▃▃▃▃▃▃▃▃▃▃▃▄▃▃▄▄▄▃▆▄▄▄▄▄▄▅▄▄▄▄▄▄▅▄▅▄█</td></tr><tr><td>metrics/l0</td><td>█▅▅▂▃▅▆▄▅█▃▅▃▅▇▄▆▄▆▄▄▁▄▅▅▄▁▃▇▄▄▅▃▆▃▄▄▄▃▁</td></tr><tr><td>metrics/l2_norm</td><td>▇▆▇▁▄▃▆▃▃▂▂▄▂▂▃▂▄▃▂▃▄▅▅▃▃▃▄▄▃▃▂▅▃▃▃▃▃▆▆█</td></tr><tr><td>metrics/l2_norm_in</td><td>▃▃▇▂▄▂▅▅▅▅▄▆▃▃▁▃▄▅▃▂▅▂▅▄▃▄▃▅▂▅▄▃▃▄▁▄▂█▃▅</td></tr><tr><td>metrics/l2_ratio</td><td>█▆▆▁▃▃▅▂▃▁▂▃▂▂▃▂▄▂▂▄▄▅▅▃▃▂▅▃▃▃▂▅▃▂▃▃▃▅▆█</td></tr><tr><td>metrics/mean_log10_feature_sparsity</td><td>█▇▇▇▇▇▇▇▇▇▇▇▇▇▇▆▆▆▆▆▆▆▆▆▅▄▄▂▁▁</td></tr><tr><td>sparsity/below_1e-5</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▁▁▁▁▁▁▁▁▂▃▆██</td></tr><tr><td>sparsity/below_1e-6</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▁▁▁▁▁▁▁▁▂▃▆██</td></tr><tr><td>sparsity/dead_features</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▂▁▁▁▁▁▁▁▁▁▁▃▃▄▇██</td></tr><tr><td>sparsity/mean_passes_since_fired</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▂▁▁▁▁▁▁▁▁▁▂▂▂▄▅▇█</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>details/current_l1_coefficient</td><td>20</td></tr><tr><td>details/current_learning_rate</td><td>0.0</td></tr><tr><td>details/n_training_tokens</td><td>122880000</td></tr><tr><td>losses/auxiliary_reconstruction_loss</td><td>227.78122</td></tr><tr><td>losses/ghost_grad_loss</td><td>0.0</td></tr><tr><td>losses/l1_loss</td><td>0.0</td></tr><tr><td>losses/mse_loss</td><td>143.06226</td></tr><tr><td>losses/overall_loss</td><td>434.46942</td></tr><tr><td>losses/sfn_sparsity_loss</td><td>63.62593</td></tr><tr><td>metrics/CE_loss_score</td><td>0.59248</td></tr><tr><td>metrics/ce_loss_with_ablation</td><td>8.29373</td></tr><tr><td>metrics/ce_loss_with_sae</td><td>4.50411</td></tr><tr><td>metrics/ce_loss_without_sae</td><td>1.8969</td></tr><tr><td>metrics/explained_variance</td><td>0.15973</td></tr><tr><td>metrics/explained_variance_std</td><td>0.24142</td></tr><tr><td>metrics/l0</td><td>7705.52734</td></tr><tr><td>metrics/l2_norm</td><td>14.99578</td></tr><tr><td>metrics/l2_norm_in</td><td>17.58649</td></tr><tr><td>metrics/l2_ratio</td><td>0.8463</td></tr><tr><td>metrics/mean_log10_feature_sparsity</td><td>-0.74933</td></tr><tr><td>sparsity/below_1e-5</td><td>681</td></tr><tr><td>sparsity/below_1e-6</td><td>681</td></tr><tr><td>sparsity/dead_features</td><td>681</td></tr><tr><td>sparsity/mean_passes_since_fired</td><td>138.74988</td></tr></table><br/></div></div>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View run <strong style=\"color:#cdcd00\">16384-L1-20-LR-5e-05-Tokens-1.229e+08</strong> at: <a href='https://wandb.ai/curt-tigges/gated_sae_testing/runs/ch6e0a5s' target=\"_blank\">https://wandb.ai/curt-tigges/gated_sae_testing/runs/ch6e0a5s</a><br/> View project at: <a href='https://wandb.ai/curt-tigges/gated_sae_testing' target=\"_blank\">https://wandb.ai/curt-tigges/gated_sae_testing</a><br/>Synced 6 W&B file(s), 0 media file(s), 3 artifact file(s) and 0 other file(s)"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Find logs at: <code>./wandb/run-20240616_143959-ch6e0a5s/logs</code>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        }
+      ],
+      "source": [
+        "total_training_steps = 30_000  # probably we should do more\n",
+        "batch_size = 4096\n",
+        "total_training_tokens = total_training_steps * batch_size\n",
+        "\n",
+        "lr_warm_up_steps = 0\n",
+        "lr_decay_steps = total_training_steps // 5  # 20% of training\n",
+        "l1_warm_up_steps = 10_000 #total_training_steps // 20  # 5% of training\n",
+        "\n",
+        "cfg = LanguageModelSAERunnerConfig(\n",
+        "    # Data Generating Function (Model + Training Distribution)\n",
+        "    architecture=\"gated\",  # we'll use the gated variant.\n",
+        "    model_name=\"tiny-stories-1L-21M\",  # our model (more options here: https://neelnanda-io.github.io/TransformerLens/generated/model_properties_table.html)\n",
+        "    hook_name=\"blocks.0.hook_mlp_out\",  # A valid hook point (see more details here: https://neelnanda-io.github.io/TransformerLens/generated/demos/Main_Demo.html#Hook-Points)\n",
+        "    hook_layer=0,  # Only one layer in the model.\n",
+        "    d_in=1024,  # the width of the mlp output.\n",
+        "    dataset_path=\"apollo-research/roneneldan-TinyStories-tokenizer-gpt2\",  # this is a tokenized language dataset on Huggingface for the Tiny Stories corpus.\n",
+        "    is_dataset_tokenized=True,\n",
+        "    streaming=True,  # we could pre-download the token dataset if it was small.\n",
+        "    # SAE Parameters\n",
+        "    mse_loss_normalization=None,  # We won't normalize the mse loss,\n",
+        "    expansion_factor=16,  # the width of the SAE. Larger will result in better stats but slower training.\n",
+        "    b_dec_init_method=\"zeros\",  # The geometric median can be used to initialize the decoder weights.\n",
+        "    apply_b_dec_to_input=True,  # We won't apply the decoder weights to the input.\n",
+        "    normalize_sae_decoder=False,\n",
+        "    scale_sparsity_penalty_by_decoder_norm=False,\n",
+        "    decoder_heuristic_init=True,\n",
+        "    init_encoder_as_decoder_transpose=True,\n",
+        "    normalize_activations=False,\n",
+        "    # Training Parameters\n",
+        "    lr=5e-5,  # lower the better, we'll go fairly high to speed up the tutorial.\n",
+        "    adam_beta1=0.9,  # adam params (default, but once upon a time we experimented with these.)\n",
+        "    adam_beta2=0.999,\n",
+        "    lr_scheduler_name=\"constant\",  # constant learning rate with warmup. Could be better schedules out there.\n",
+        "    lr_warm_up_steps=lr_warm_up_steps,  # this can help avoid too many dead features initially.\n",
+        "    lr_decay_steps=lr_decay_steps,  # this will help us avoid overfitting.\n",
+        "    l1_coefficient=20,  # will control how sparse the feature activations are\n",
+        "    l1_warm_up_steps=l1_warm_up_steps,  # this can help avoid too many dead features initially.\n",
+        "    lp_norm=1.0,  # the L1 penalty (and not a Lp for p < 1)\n",
+        "    train_batch_size_tokens=batch_size,\n",
+        "    context_size=256,  # will control the lenght of the prompts we feed to the model. Larger is better but slower. so for the tutorial we'll use a short one.\n",
+        "    # Activation Store Parameters\n",
+        "    n_batches_in_buffer=64,  # controls how many activations we store / shuffle.\n",
+        "    training_tokens=total_training_tokens,  # 100 million tokens is quite a few, but we want to see good stats. Get a coffee, come back.\n",
+        "    store_batch_size_prompts=16,\n",
+        "    # Resampling protocol\n",
+        "    use_ghost_grads=False,  # we don't use ghost grads anymore.\n",
+        "    feature_sampling_window=1000,  # this controls our reporting of feature sparsity stats\n",
+        "    dead_feature_window=1000,  # would effect resampling or ghost grads if we were using it.\n",
+        "    dead_feature_threshold=1e-4,  # would effect resampling or ghost grads if we were using it.\n",
+        "    # WANDB\n",
+        "    log_to_wandb=True,  # always use wandb unless you are just testing code.\n",
+        "    wandb_project=\"gated_sae_testing\",\n",
+        "    wandb_log_frequency=30,\n",
+        "    eval_every_n_wandb_logs=20,\n",
+        "    # Misc\n",
+        "    device=device,\n",
+        "    seed=42,\n",
+        "    n_checkpoints=0,\n",
+        "    checkpoint_path=\"checkpoints\",\n",
+        "    dtype=\"float32\"\n",
+        ")\n",
+        "# look at the next cell to see some instruction for what to do while this is running.\n",
+        "sparse_autoencoder = SAETrainingRunner(cfg).run()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.10.14"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
diff --git a/tutorials/training_a_sparse_autoencoder.ipynb b/tutorials/training_a_sparse_autoencoder.ipynb
index 006b880c..31a25850 100644
--- a/tutorials/training_a_sparse_autoencoder.ipynb
+++ b/tutorials/training_a_sparse_autoencoder.ipynb
@@ -22,15 +22,162 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 2,
       "metadata": {
         "id": "LeRi_tw2dhae"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Requirement already satisfied: sae-lens in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (3.3.0)\n",
+            "Requirement already satisfied: transformer-lens in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (1.19.0)\n",
+            "Requirement already satisfied: circuitsvis in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (1.43.2)\n",
+            "Requirement already satisfied: automated-interpretability<0.0.4,>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.0.3)\n",
+            "Requirement already satisfied: babe<0.0.8,>=0.0.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.0.7)\n",
+            "Requirement already satisfied: datasets<3.0.0,>=2.17.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (2.19.2)\n",
+            "Requirement already satisfied: matplotlib<4.0.0,>=3.8.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (3.9.0)\n",
+            "Requirement already satisfied: matplotlib-inline<0.2.0,>=0.1.6 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.1.7)\n",
+            "Requirement already satisfied: nltk<4.0.0,>=3.8.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (3.8.1)\n",
+            "Requirement already satisfied: plotly<6.0.0,>=5.19.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (5.22.0)\n",
+            "Requirement already satisfied: plotly-express<0.5.0,>=0.4.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.4.1)\n",
+            "Requirement already satisfied: pytest-profiling<2.0.0,>=1.7.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (1.7.0)\n",
+            "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (1.0.1)\n",
+            "Requirement already satisfied: pyyaml<7.0.0,>=6.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (6.0.1)\n",
+            "Requirement already satisfied: pyzmq==26.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (26.0.0)\n",
+            "Requirement already satisfied: sae-vis<0.3.0,>=0.2.18 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.2.18)\n",
+            "Requirement already satisfied: safetensors<0.5.0,>=0.4.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.4.3)\n",
+            "Requirement already satisfied: transformers<5.0.0,>=4.38.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (4.41.2)\n",
+            "Requirement already satisfied: typer<0.13.0,>=0.12.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-lens) (0.12.3)\n",
+            "Requirement already satisfied: accelerate>=0.23.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.31.0)\n",
+            "Requirement already satisfied: beartype<0.15.0,>=0.14.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.14.1)\n",
+            "Requirement already satisfied: better-abc<0.0.4,>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.0.3)\n",
+            "Requirement already satisfied: einops>=0.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.7.0)\n",
+            "Requirement already satisfied: fancy-einsum>=0.0.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.0.3)\n",
+            "Requirement already satisfied: jaxtyping>=0.2.11 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.2.29)\n",
+            "Requirement already satisfied: numpy>=1.24 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (1.26.4)\n",
+            "Requirement already satisfied: pandas>=1.1.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (2.2.2)\n",
+            "Requirement already satisfied: rich>=12.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (13.7.1)\n",
+            "Requirement already satisfied: sentencepiece in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.2.0)\n",
+            "Requirement already satisfied: torch>=1.10 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (2.1.2)\n",
+            "Requirement already satisfied: tqdm>=4.64.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (4.66.4)\n",
+            "Requirement already satisfied: typing-extensions in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (4.12.2)\n",
+            "Requirement already satisfied: wandb>=0.13.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformer-lens) (0.17.1)\n",
+            "Requirement already satisfied: importlib-metadata>=5.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (7.1.0)\n",
+            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.3.1)\n",
+            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (8.9.2.26)\n",
+            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (11.0.2.54)\n",
+            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (10.3.2.106)\n",
+            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (11.4.5.107)\n",
+            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.0.106)\n",
+            "Requirement already satisfied: nvidia-nccl-cu12==2.18.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (2.18.1)\n",
+            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (12.1.105)\n",
+            "Requirement already satisfied: triton==2.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from circuitsvis) (2.1.0)\n",
+            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->circuitsvis) (12.5.40)\n",
+            "Requirement already satisfied: filelock in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from triton==2.1.0->circuitsvis) (3.14.0)\n",
+            "Requirement already satisfied: packaging>=20.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (24.0)\n",
+            "Requirement already satisfied: psutil in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (5.9.0)\n",
+            "Requirement already satisfied: huggingface-hub in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from accelerate>=0.23.0->transformer-lens) (0.23.3)\n",
+            "Requirement already satisfied: blobfile<3.0.0,>=2.1.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.1.1)\n",
+            "Requirement already satisfied: boostedblob<0.16.0,>=0.15.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.15.3)\n",
+            "Requirement already satisfied: httpx<0.28.0,>=0.27.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.27.0)\n",
+            "Requirement already satisfied: orjson<4.0.0,>=3.10.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.10.4)\n",
+            "Requirement already satisfied: pytest<9.0.0,>=8.1.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (8.2.2)\n",
+            "Requirement already satisfied: scikit-learn<2.0.0,>=1.4.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.5.0)\n",
+            "Requirement already satisfied: tiktoken<0.7.0,>=0.6.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.6.0)\n",
+            "Requirement already satisfied: py2store in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from babe<0.0.8,>=0.0.7->sae-lens) (0.1.20)\n",
+            "Requirement already satisfied: graze in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from babe<0.0.8,>=0.0.7->sae-lens) (0.1.17)\n",
+            "Requirement already satisfied: pyarrow>=12.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (16.1.0)\n",
+            "Requirement already satisfied: pyarrow-hotfix in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.6)\n",
+            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.3.8)\n",
+            "Requirement already satisfied: requests>=2.32.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (2.32.3)\n",
+            "Requirement already satisfied: xxhash in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (3.4.1)\n",
+            "Requirement already satisfied: multiprocess in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (0.70.16)\n",
+            "Requirement already satisfied: fsspec<=2024.3.1,>=2023.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from fsspec[http]<=2024.3.1,>=2023.1.0->datasets<3.0.0,>=2.17.1->sae-lens) (2024.3.1)\n",
+            "Requirement already satisfied: aiohttp in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from datasets<3.0.0,>=2.17.1->sae-lens) (3.9.5)\n",
+            "Requirement already satisfied: zipp>=0.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from importlib-metadata>=5.1.0->circuitsvis) (3.19.2)\n",
+            "Requirement already satisfied: typeguard==2.13.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from jaxtyping>=0.2.11->transformer-lens) (2.13.3)\n",
+            "Requirement already satisfied: contourpy>=1.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (1.2.1)\n",
+            "Requirement already satisfied: cycler>=0.10 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (0.12.1)\n",
+            "Requirement already satisfied: fonttools>=4.22.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (4.53.0)\n",
+            "Requirement already satisfied: kiwisolver>=1.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (1.4.5)\n",
+            "Requirement already satisfied: pillow>=8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (10.3.0)\n",
+            "Requirement already satisfied: pyparsing>=2.3.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (3.1.2)\n",
+            "Requirement already satisfied: python-dateutil>=2.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib<4.0.0,>=3.8.3->sae-lens) (2.9.0)\n",
+            "Requirement already satisfied: traitlets in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from matplotlib-inline<0.2.0,>=0.1.6->sae-lens) (5.14.3)\n",
+            "Requirement already satisfied: click in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (8.1.7)\n",
+            "Requirement already satisfied: joblib in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (1.4.2)\n",
+            "Requirement already satisfied: regex>=2021.8.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from nltk<4.0.0,>=3.8.1->sae-lens) (2024.5.15)\n",
+            "Requirement already satisfied: pytz>=2020.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pandas>=1.1.5->transformer-lens) (2024.1)\n",
+            "Requirement already satisfied: tzdata>=2022.7 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pandas>=1.1.5->transformer-lens) (2024.1)\n",
+            "Requirement already satisfied: tenacity>=6.2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly<6.0.0,>=5.19.0->sae-lens) (8.3.0)\n",
+            "Requirement already satisfied: statsmodels>=0.9.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (0.14.2)\n",
+            "Requirement already satisfied: scipy>=0.18 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (1.13.1)\n",
+            "Requirement already satisfied: patsy>=0.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from plotly-express<0.5.0,>=0.4.1->sae-lens) (0.5.6)\n",
+            "Requirement already satisfied: six in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens) (1.16.0)\n",
+            "Requirement already satisfied: gprof2dot in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest-profiling<2.0.0,>=1.7.0->sae-lens) (2024.6.6)\n",
+            "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from rich>=12.6.0->transformer-lens) (3.0.0)\n",
+            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from rich>=12.6.0->transformer-lens) (2.18.0)\n",
+            "Requirement already satisfied: dataclasses-json<0.7.0,>=0.6.4 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens) (0.6.7)\n",
+            "Requirement already satisfied: eindex-callum<0.2.0,>=0.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sae-vis<0.3.0,>=0.2.18->sae-lens) (0.1.1)\n",
+            "Requirement already satisfied: sympy in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (1.12.1)\n",
+            "Requirement already satisfied: networkx in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (3.3)\n",
+            "Requirement already satisfied: jinja2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from torch>=1.10->transformer-lens) (3.1.4)\n",
+            "Requirement already satisfied: tokenizers<0.20,>=0.19 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from transformers<5.0.0,>=4.38.1->sae-lens) (0.19.1)\n",
+            "Requirement already satisfied: shellingham>=1.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from typer<0.13.0,>=0.12.3->sae-lens) (1.5.4)\n",
+            "Requirement already satisfied: docker-pycreds>=0.4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (0.4.0)\n",
+            "Requirement already satisfied: gitpython!=3.1.29,>=1.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (3.1.43)\n",
+            "Requirement already satisfied: platformdirs in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (4.2.2)\n",
+            "Requirement already satisfied: protobuf!=4.21.0,<6,>=3.19.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (5.27.1)\n",
+            "Requirement already satisfied: sentry-sdk>=1.0.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (2.5.1)\n",
+            "Requirement already satisfied: setproctitle in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (1.3.3)\n",
+            "Requirement already satisfied: setuptools in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from wandb>=0.13.5->transformer-lens) (70.0.0)\n",
+            "Requirement already satisfied: pycryptodomex~=3.8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.20.0)\n",
+            "Requirement already satisfied: urllib3<3,>=1.25.3 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.2.1)\n",
+            "Requirement already satisfied: lxml~=4.9 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from blobfile<3.0.0,>=2.1.1->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (4.9.4)\n",
+            "Requirement already satisfied: uvloop>=0.16.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from boostedblob<0.16.0,>=0.15.3->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.19.0)\n",
+            "Requirement already satisfied: aiosignal>=1.1.2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.3.1)\n",
+            "Requirement already satisfied: attrs>=17.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (23.2.0)\n",
+            "Requirement already satisfied: frozenlist>=1.1.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.4.1)\n",
+            "Requirement already satisfied: multidict<7.0,>=4.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (6.0.5)\n",
+            "Requirement already satisfied: yarl<2.0,>=1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (1.9.4)\n",
+            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from aiohttp->datasets<3.0.0,>=2.17.1->sae-lens) (4.0.3)\n",
+            "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (3.21.3)\n",
+            "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (0.9.0)\n",
+            "Requirement already satisfied: gitdb<5,>=4.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer-lens) (4.0.11)\n",
+            "Requirement already satisfied: anyio in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (4.4.0)\n",
+            "Requirement already satisfied: certifi in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2024.6.2)\n",
+            "Requirement already satisfied: httpcore==1.* in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.0.5)\n",
+            "Requirement already satisfied: idna in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.7)\n",
+            "Requirement already satisfied: sniffio in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.3.1)\n",
+            "Requirement already satisfied: h11<0.15,>=0.13 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from httpcore==1.*->httpx<0.28.0,>=0.27.0->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (0.14.0)\n",
+            "Requirement already satisfied: mdurl~=0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich>=12.6.0->transformer-lens) (0.1.2)\n",
+            "Requirement already satisfied: iniconfig in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.0.0)\n",
+            "Requirement already satisfied: pluggy<2.0,>=1.5 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.5.0)\n",
+            "Requirement already satisfied: exceptiongroup>=1.0.0rc8 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (1.2.0)\n",
+            "Requirement already satisfied: tomli>=1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from pytest<9.0.0,>=8.1.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (2.0.1)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from requests>=2.32.1->datasets<3.0.0,>=2.17.1->sae-lens) (3.3.2)\n",
+            "Requirement already satisfied: threadpoolctl>=3.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from scikit-learn<2.0.0,>=1.4.2->automated-interpretability<0.0.4,>=0.0.3->sae-lens) (3.5.0)\n",
+            "Requirement already satisfied: dol in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from graze->babe<0.0.8,>=0.0.7->sae-lens) (0.2.47)\n",
+            "Requirement already satisfied: MarkupSafe>=2.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from jinja2->torch>=1.10->transformer-lens) (2.1.5)\n",
+            "Requirement already satisfied: config2py in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens) (0.1.33)\n",
+            "Requirement already satisfied: importlib-resources in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from py2store->babe<0.0.8,>=0.0.7->sae-lens) (6.4.0)\n",
+            "Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from sympy->torch>=1.10->transformer-lens) (1.3.0)\n",
+            "Requirement already satisfied: smmap<6,>=3.0.1 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from gitdb<5,>=4.0.1->gitpython!=3.1.29,>=1.0.0->wandb>=0.13.5->transformer-lens) (5.0.1)\n",
+            "Requirement already satisfied: mypy-extensions>=0.3.0 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7.0,>=0.6.4->sae-vis<0.3.0,>=0.2.18->sae-lens) (1.0.0)\n",
+            "Requirement already satisfied: i2 in /home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages (from config2py->py2store->babe<0.0.8,>=0.0.7->sae-lens) (0.1.17)\n",
+            "Note: you may need to restart the kernel to use updated packages.\n"
+          ]
+        }
+      ],
       "source": [
         "try:\n",
-        "    import google.colab # type: ignore\n",
-        "    from google.colab import output\n",
+        "    #import google.colab # type: ignore\n",
+        "    #from google.colab import output\n",
         "    %pip install sae-lens transformer-lens circuitsvis\n",
         "except:\n",
         "    from IPython import get_ipython # type: ignore\n",
@@ -41,7 +188,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 3,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -49,7 +196,23 @@
         "id": "uy-b3CcSOVHu",
         "outputId": "58ce28d0-f91f-436d-cf87-76bb26e2ecaf"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "  from .autonotebook import tqdm as notebook_tqdm\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Using device: cuda\n"
+          ]
+        }
+      ],
       "source": [
         "import torch\n",
         "import os\n",
@@ -85,11 +248,29 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 4,
       "metadata": {
         "id": "hFz6JUMuOVHv"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
+            "  warnings.warn(\n",
+            "/home/curttigges/miniconda3/envs/saelens/lib/python3.10/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
+            "  return self.fget.__get__(instance, owner)()\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Loaded pretrained model tiny-stories-1L-21M into HookedTransformer\n"
+          ]
+        }
+      ],
       "source": [
         "from transformer_lens import HookedTransformer\n",
         "\n",
@@ -118,11 +299,57 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 5,
       "metadata": {
         "id": "G4ad4Zz1OVHv"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "'Once upon a time, Bobby was hungry and needed something to do. He went to the subway but was far away.\\n\\nThe man wanted to get the hat, so the people wanted it. He found it hard to be a big, powerful bird. It wanted'"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/plain": [
+              "'Once upon a time, there was a trunk. The trunk was very rich, and it was a very special trunk. All the animals came across the trunk, and was very colorful. They took turns to fill it up when they bumped into a dragon. \\n\\n'"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/plain": [
+              "'Once upon a time, there was a young man. He was three years old. He wanted to learn how to keep the match safe. So he kept checking it every day.\\n\\nOne day a 3 year old girl wanted to learn about the fun trunk of the'"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/plain": [
+              "'Once upon a time, there was a little girl named Sally. She liked to play with her toys. One sunny day, Sally found a butterfly. She was so happy! She wanted to play with something new. So, she called her friend Tom.\\n\\nTom'"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/plain": [
+              "'Once upon a time, there was a little girl named Lola. She really loved playing with her pet cat, Tom and show them his appreciation.\\n\\nOne day, Lola licked the couch closer and soon found herself in a magical land! As soon as'"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "# here we use generate to get 10 completeions with temperature 1. Feel free to play with the prompt to make it more interesting.\n",
         "for i in range(5):\n",
@@ -157,11 +384,64 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 6,
       "metadata": {
         "id": "TpmPoj7uOVHv"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Tokenized prompt: ['<|endoftext|>', 'Once', ' upon', ' a', ' time', ',', ' there', ' was', ' a', ' little', ' girl', ' named', ' Lily', '.', ' She', ' lived', ' in', ' a', ' big', ',', ' happy', ' little', ' girl', '.', ' On', ' her', ' big', ' adventure', ',']\n",
+            "Tokenized answer: [' Lily']\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Performance on answer token:\n",
+              "<span style=\"font-weight: bold\">Rank: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">        Logit: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">18.81</span><span style=\"font-weight: bold\"> Prob: </span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">13.46</span><span style=\"font-weight: bold\">% Token: | Lily|</span>\n",
+              "</pre>\n"
+            ],
+            "text/plain": [
+              "Performance on answer token:\n",
+              "\u001b[1mRank: \u001b[0m\u001b[1;36m1\u001b[0m\u001b[1m        Logit: \u001b[0m\u001b[1;36m18.81\u001b[0m\u001b[1m Prob: \u001b[0m\u001b[1;36m13.46\u001b[0m\u001b[1m% Token: | Lily|\u001b[0m\n"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Top 0th token. Logit: 20.48 Prob: 71.06% Token: | she|\n",
+            "Top 1th token. Logit: 18.81 Prob: 13.46% Token: | Lily|\n",
+            "Top 2th token. Logit: 17.35 Prob:  3.11% Token: | the|\n",
+            "Top 3th token. Logit: 17.26 Prob:  2.86% Token: | her|\n",
+            "Top 4th token. Logit: 16.74 Prob:  1.70% Token: | there|\n",
+            "Top 5th token. Logit: 16.43 Prob:  1.25% Token: | they|\n",
+            "Top 6th token. Logit: 15.80 Prob:  0.66% Token: | all|\n",
+            "Top 7th token. Logit: 15.64 Prob:  0.56% Token: | things|\n",
+            "Top 8th token. Logit: 15.28 Prob:  0.39% Token: | one|\n",
+            "Top 9th token. Logit: 15.24 Prob:  0.38% Token: | lived|\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Ranks of the answer tokens:</span> <span style=\"font-weight: bold\">[(</span><span style=\"color: #008000; text-decoration-color: #008000\">' Lily'</span>, <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">)]</span>\n",
+              "</pre>\n"
+            ],
+            "text/plain": [
+              "\u001b[1mRanks of the answer tokens:\u001b[0m \u001b[1m[\u001b[0m\u001b[1m(\u001b[0m\u001b[32m' Lily'\u001b[0m, \u001b[1;36m1\u001b[0m\u001b[1m)\u001b[0m\u001b[1m]\u001b[0m\n"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "from transformer_lens.utils import test_prompt\n",
         "\n",
@@ -203,11 +483,33 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 7,
       "metadata": {
         "id": "Tic0RCUpOVHw"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/html": [
+              "<div id=\"circuits-vis-118354ce-b30f\" style=\"margin: 15px 0;\"/>\n",
+              "    <script crossorigin type=\"module\">\n",
+              "    import { render, TokenLogProbs } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
+              "    render(\n",
+              "      \"circuits-vis-118354ce-b30f\",\n",
+              "      TokenLogProbs,\n",
+              "      {\"prompt\": [\"<|endoftext|>\", \"Hi\", \",\", \" how\", \" are\", \" you\", \" doing\", \" this\", \"?\", \" I\", \"'m\", \" really\", \" enjoying\", \" your\", \" posts\"], \"topKLogProbs\": [[-0.020237978547811508, -6.104498863220215, -7.070311546325684, -7.073490142822266, -7.178879737854004, -7.389684677124023, -7.525315284729004, -7.588573455810547, -7.763055801391602, -8.224309921264648], [-1.1883893013000488, -3.4683852195739746, -3.4808287620544434, -3.4896111488342285, -3.500030994415283, -3.554358959197998, -4.017881870269775, -4.030544757843018, -4.298159122467041, -4.418168544769287], [-2.9201290607452393, -3.379486322402954, -3.4266788959503174, -3.531764268875122, -3.956162691116333, -4.256065368652344, -4.443942070007324, -4.4448394775390625, -4.559991836547852, -4.583307266235352], [-0.9917522668838501, -2.928676128387451, -3.2575316429138184, -3.317981243133545, -3.4740567207336426, -3.6419081687927246, -3.8674941062927246, -3.868630886077881, -3.8917040824890137, -4.005284786224365], [-0.037123147398233414, -4.683361530303955, -4.9722676277160645, -5.854288578033447, -5.9421305656433105, -6.388034343719482, -6.3940815925598145, -6.498584270477295, -6.593432903289795, -7.8692545890808105], [-0.23486259579658508, -2.3354179859161377, -2.947497606277466, -5.052804470062256, -5.078576564788818, -5.102746486663818, -5.116313457489014, -5.737276554107666, -5.768456935882568, -5.893773555755615], [-0.44884783029556274, -2.076716423034668, -3.433588981628418, -3.645073890686035, -3.7033796310424805, -4.046278953552246, -4.097006797790527, -4.209780693054199, -4.518967628479004, -4.961886405944824], [-0.4609467387199402, -2.7377679347991943, -3.158902883529663, -3.686795949935913, -4.416429042816162, -4.448527812957764, -4.571945667266846, -4.669140338897705, -4.695371150970459, -4.748204708099365], [-0.9866101741790771, -1.8349759578704834, -2.460031270980835, -2.511583089828491, -3.1156270503997803, -3.419478178024292, -3.5203635692596436, -3.8049819469451904, -4.003089904785156, -4.258264541625977], [-1.3091926574707031, -1.633901596069336, -1.78680419921875, -2.631612777709961, -3.4602813720703125, -3.615842819213867, -4.020046234130859, -4.093286514282227, -4.177457809448242, -4.309240341186523], [-1.9952350854873657, -2.55653715133667, -2.645555019378662, -2.8948235511779785, -3.11043119430542, -3.3945651054382324, -3.7311148643493652, -3.8203349113464355, -3.967912197113037, -4.026367664337158], [-1.0431095361709595, -1.959466576576233, -3.0507731437683105, -3.4215006828308105, -3.48012113571167, -3.4898486137390137, -3.6020045280456543, -4.332496166229248, -4.492417812347412, -4.517611026763916], [-1.5123934745788574, -2.020272731781006, -2.22328519821167, -2.2493205070495605, -3.376523494720459, -3.92171049118042, -3.9798388481140137, -4.199243068695068, -4.265368938446045, -4.6615891456604], [-3.188736915588379, -3.192061424255371, -3.4252281188964844, -3.494192123413086, -3.7325897216796875, -3.7353057861328125, -3.735713005065918, -3.833963394165039, -4.233694076538086, -4.265708923339844]], \"topKTokens\": [[\"\\n\", \",\", \"Words\", \"Summary\", \"\\n\\n\", \"<|endoftext|>\", \" \", \"Features\", \"Random\", \" the\"], [\",\", \" Tim\", \" bird\", \"!\", \" little\", \" Lily\", \" Tom\", \"!\\\"\", \" Max\", \" tree\"], [\" cat\", \" a\", \" bird\", \" little\", \" I\", \"\\n\", \" dog\", \" but\", \" frog\", \" sw\"], [\" are\", \" I\", \" was\", \" do\", \" fast\", \" big\", \" did\", \" to\", \" he\", \" old\"], [\" you\", \" we\", \" they\", \" the\", \" Lily\", \" your\", \" I\", \" Tim\", \" friends\", \" those\"], [\"?\\\"\", \" today\", \"?\", \" and\", \".\", \",\", \" doing\", \" going\", \" up\", \"?\\\".\"], [\"?\\\"\", \" this\", \" today\", \"?\", \" here\", \" a\", \" it\", \" these\", \".\", \" in\"], [\"?\\\"\", \"?\", \" great\", \" so\", \" job\", \" puzzle\", \" task\", \" very\", \".\", \" amazing\"], [\"\\ufffd\", \" You\", \" I\", \" It\", \" We\", \" Can\", \"\\n\", \" This\", \" Do\", \" Are\"], [\"'m\", \" am\", \" want\", \" like\", \" know\", \" have\", \" love\", \" can\", \" just\", \" hope\"], [\" a\", \" so\", \" very\", \" trying\", \" going\", \" an\", \" looking\", \" just\", \" making\", \" playing\"], [\" good\", \" hungry\", \" enjoying\", \" a\", \" tired\", \" proud\", \" happy\", \" excited\", \" busy\", \" glad\"], [\" the\", \" it\", \" this\", \" my\", \" a\", \" your\", \" playing\", \" some\", \" these\", \" myself\"], [\" meal\", \" food\", \" day\", \" ice\", \" sandwich\", \" delicious\", \" new\", \" dance\", \" cake\", \" time\"]], \"correctTokenRank\": [1718, 0, 675, 0, 0, 6, 1, 1, 2, 0, 23, 2, 5, 10036], \"correctTokenLogProb\": [-13.814948081970215, -1.1883893013000488, -8.445575714111328, -0.9917522668838501, -0.037123147398233414, -5.116313457489014, -2.076716423034668, -2.7377679347991943, -2.460031270980835, -1.3091926574707031, -4.860561847686768, -3.0507731437683105, -3.92171049118042, -17.039505004882812]}\n",
+              "    )\n",
+              "    </script>"
+            ],
+            "text/plain": [
+              "<circuitsvis.utils.render.RenderedHTML at 0x77eaf81e21d0>"
+            ]
+          },
+          "execution_count": 7,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "import circuitsvis as cv  # optional dep, install with pip install circuitsvis\n",
         "\n",
@@ -238,11 +540,40 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 8,
       "metadata": {
         "id": "Nikp2ASlOVHw"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "100%|██████████| 200/200 [00:01<00:00, 103.20it/s]\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "<div id=\"circuits-vis-a7b13498-82dd\" style=\"margin: 15px 0;\"/>\n",
+              "    <script crossorigin type=\"module\">\n",
+              "    import { render, TokenLogProbs } from \"https://unpkg.com/circuitsvis@1.43.2/dist/cdn/esm.js\";\n",
+              "    render(\n",
+              "      \"circuits-vis-a7b13498-82dd\",\n",
+              "      TokenLogProbs,\n",
+              "      {\"prompt\": [\"<|endoftext|>\", \"Once\", \" upon\", \" a\", \" time\", \",\", \" there\", \" was\", \" a\", \" mom\", \"my\", \" and\", \" a\", \" big\", \",\", \" modern\", \" airport\", \".\", \" They\", \" saw\", \" many\", \" people\", \" walking\", \" around\", \" and\", \" waved\", \" at\", \" all\", \" the\", \" planes\", \" too\", \".\", \"\\n\", \"\\n\", \"The\", \" family\", \" asked\", \" the\", \" little\", \" girl\", \" and\", \" daddy\", \" lots\", \" of\", \" questions\", \",\", \" so\", \" the\", \" mom\", \"my\", \" replied\", \",\", \" \\\"\", \"Are\", \" we\", \" just\", \" going\", \" to\", \" explore\", \"?\\\"\", \"\\n\", \"\\n\", \"The\", \" little\", \" girl\", \" smiled\", \" and\", \" nodded\", \" her\", \" head\", \".\", \" She\", \" was\", \" excited\", \" to\", \" see\", \" all\", \" the\", \" planes\", \",\", \" and\", \" she\", \" got\", \" to\", \" the\", \" airport\", \".\", \"\\n\", \"\\n\", \"The\", \" little\", \" girl\", \" climbed\", \" up\", \" the\", \" stairs\", \" and\", \" said\", \",\", \" \\\"\", \"Well\", \" done\", \",\", \" everyone\", \" is\", \" so\", \" proud\", \" of\", \" you\", \"!\", \" Are\", \" you\", \" ready\", \" to\", \" show\", \" your\", \" family\", \"?\\\"\", \"\\n\", \"\\n\", \"The\", \" family\", \" looked\", \",\", \" and\", \" they\", \" all\", \" hugged\", \" each\", \" other\", \",\", \" and\", \" they\", \" agreed\", \"!\", \" The\", \" little\", \" girl\", \" was\", \" so\", \" happy\", \",\", \" she\", \" wanted\", \" to\", \" show\", \" her\", \" new\", \" toy\", \".\", \"\\n\", \"<|endoftext|>\", \".\", \"The\", \" family\", \" went\", \" on\", \" to\", \" the\", \" modern\", \" airport\", \" to\", \" come\", \" on\", \" these\", \" trips\", \".\", \" Every\", \" time\", \" one\", \" of\", \" the\", \" little\", \" girl\", \",\", \" her\", \" family\", \",\", \" would\", \" march\", \" with\", \" enthusiasm\", \" and\", \" hope\", \" for\", \" the\", \" family\", \".\", \"\\n\", \"\\n\", \"Eventually\", \",\", \" the\", \" little\", \" girl\", \" began\", \" to\", \" get\", \" thirsty\", \" again\", \".\", \" She\", \" was\", \" so\", \" tired\"], \"topKLogProbs\": [[-0.020237978547811508, -6.104497909545898, -7.070315361022949, -7.073488235473633, -7.1788787841796875, -7.389683723449707, -7.525317192077637, -7.5885772705078125, -7.763050079345703, -8.224308967590332], [-0.372489869594574, -1.827528953552246, -4.3061370849609375, -4.838803291320801, -4.989402770996094, -5.1117401123046875, -5.187068939208984, -5.250604629516602, -5.662858963012695, -5.694143295288086], [-0.005904730875045061, -5.495633602142334, -7.522153377532959, -8.207719802856445, -8.449649810791016, -8.90707778930664, -9.541242599487305, -9.691887855529785, -10.111766815185547, -10.492300033569336], [-0.011737688444554806, -6.031045913696289, -7.209378242492676, -7.243843078613281, -7.747928619384766, -7.977005958557129, -8.321712493896484, -8.35726547241211, -8.502359390258789, -8.701475143432617], [-0.1878184974193573, -1.8141651153564453, -5.624242782592773, -6.579645156860352, -7.822257995605469, -8.386683464050293, -8.49909496307373, -8.57803726196289, -8.823796272277832, -8.960565567016602], [-0.10250605642795563, -3.402726650238037, -3.946469783782959, -5.159356594085693, -5.411166667938232, -6.1134867668151855, -6.131921291351318, -6.200311183929443, -6.424906253814697, -7.0071539878845215], [-0.1382717341184616, -2.544334888458252, -3.050598621368408, -7.775615215301514, -8.049697875976562, -8.075815200805664, -8.107090950012207, -8.183040618896484, -8.868163108825684, -8.878026962280273], [-0.0321299247443676, -3.6293933391571045, -7.695366382598877, -7.812208652496338, -7.8907952308654785, -8.088072776794434, -8.26425552368164, -9.077349662780762, -9.536680221557617, -9.683252334594727], [-1.843049168586731, -2.4128103256225586, -2.9644947052001953, -3.615971565246582, -4.061753273010254, -4.090791702270508, -4.206696510314941, -4.35806941986084, -4.626364707946777, -4.680699348449707], [-0.5746181607246399, -1.2087008953094482, -3.0003554821014404, -3.3169009685516357, -3.463759183883667, -5.873934745788574, -5.896479606628418, -6.705965995788574, -6.999606132507324, -7.05670166015625], [-0.4282326400279999, -2.4605259895324707, -2.6223549842834473, -3.1082043647766113, -3.1714062690734863, -4.331737995147705, -4.498405933380127, -4.684654712677002, -4.858206272125244, -5.0235915184021], [-0.30328744649887085, -1.6684131622314453, -3.2264671325683594, -4.789220809936523, -5.672477722167969, -5.96893310546875, -6.063676834106445, -6.157503128051758, -6.438438415527344, -6.866046905517578], [-0.8339848518371582, -1.6514954566955566, -2.2883095741271973, -2.579624652862549, -4.058284282684326, -4.280709743499756, -4.445773601531982, -4.500208377838135, -4.837723255157471, -4.86464262008667], [-2.4236323833465576, -3.5779502391815186, -3.747347593307495, -3.8657243251800537, -3.9251821041107178, -4.05374813079834, -4.170233726501465, -4.246179580688477, -4.283749580383301, -4.3532514572143555], [-3.0394649505615234, -3.325023651123047, -3.4425315856933594, -3.482461929321289, -3.5253353118896484, -3.5320472717285156, -3.536334991455078, -3.5988540649414062, -3.6118717193603516, -3.767345428466797], [-0.9748451113700867, -3.445889711380005, -3.5092709064483643, -3.8227875232696533, -4.129208564758301, -4.215622901916504, -4.424491882324219, -4.451557159423828, -4.580273628234863, -4.590479850769043], [-0.042138077318668365, -4.544243812561035, -4.993100166320801, -5.2644147872924805, -6.0866193771362305, -6.586895942687988, -6.6464128494262695, -6.650232315063477, -7.103418350219727, -7.140406608581543], [-1.5009835958480835, -1.9959081411361694, -2.2056002616882324, -2.2982020378112793, -2.3899550437927246, -3.7624144554138184, -4.033015727996826, -4.047290325164795, -4.302596569061279, -4.305589199066162], [-0.8899617195129395, -1.8940033912658691, -2.1531224250793457, -2.6077370643615723, -3.57206392288208, -4.008375644683838, -4.123062610626221, -4.19462251663208, -4.703207492828369, -4.783227443695068], [-0.6426234245300293, -1.8766722679138184, -1.905128002166748, -3.432450771331787, -3.950042247772217, -4.064731121063232, -4.408097743988037, -4.558604717254639, -4.580577373504639, -4.752116680145264], [-0.9895710945129395, -1.2832303047180176, -2.834984302520752, -3.317976474761963, -3.6336464881896973, -3.880049228668213, -4.0731587409973145, -4.0978169441223145, -4.23159646987915, -4.357437610626221], [-1.4843578338623047, -2.5058727264404297, -2.769329071044922, -2.864400863647461, -2.9782562255859375, -3.1617069244384766, -3.1705169677734375, -3.6056766510009766, -3.7407455444335938, -3.7464637756347656], [-0.35668182373046875, -2.5939693450927734, -2.6262855529785156, -3.3084774017333984, -3.5330371856689453, -4.559474945068359, -4.66771125793457, -4.916624069213867, -5.343305587768555, -5.446603775024414], [-0.7548707127571106, -1.6354477405548096, -1.728358507156372, -2.720416307449341, -3.955390214920044, -4.285571575164795, -5.190505504608154, -5.417296886444092, -5.692910671234131, -5.70766019821167], [-2.2759742736816406, -2.447378158569336, -3.347470283508301, -3.3950443267822266, -3.646815299987793, -3.908966064453125, -3.999274253845215, -4.045774459838867, -4.176496505737305, -4.207118988037109], [-0.5946731567382812, -1.5350360870361328, -2.5619163513183594, -2.7717952728271484, -3.591419219970703, -4.287841796875, -4.993228912353516, -5.0897979736328125, -5.203357696533203, -5.756500244140625], [-1.2474231719970703, -1.3613033294677734, -2.1486759185791016, -2.4232349395751953, -2.6617660522460938, -3.4216670989990234, -3.936910629272461, -4.392925262451172, -4.43516731262207, -4.45884895324707], [-0.09278443455696106, -4.05391263961792, -4.28864049911499, -4.342662334442139, -5.320607662200928, -5.534726619720459, -5.752249240875244, -5.930094242095947, -6.490260601043701, -6.499328136444092], [-0.7439873218536377, -1.3276169300079346, -3.703294038772583, -3.781306505203247, -3.9016640186309814, -4.004000663757324, -4.276181221008301, -4.478531837463379, -5.095404624938965, -5.32029914855957], [-0.42122718691825867, -2.8185622692108154, -3.101483106613159, -3.188938856124878, -3.4723422527313232, -3.6380908489227295, -4.098782062530518, -4.378454685211182, -4.669825077056885, -4.683760166168213], [-0.06834664195775986, -2.947554349899292, -5.053656101226807, -6.055807590484619, -6.784708499908447, -7.855190753936768, -8.000382423400879, -8.236613273620605, -8.529458045959473, -8.666272163391113], [-0.889788031578064, -1.9251807928085327, -2.3246865272521973, -2.6348557472229004, -3.4432435035705566, -3.7437500953674316, -3.7707600593566895, -3.8446316719055176, -3.9515233039855957, -4.112818241119385], [-0.0339607410132885, -4.20205545425415, -6.14114236831665, -6.209860324859619, -6.44116735458374, -6.61282205581665, -6.623001575469971, -6.731749057769775, -6.739603519439697, -7.492265224456787], [-0.9458173513412476, -2.315518856048584, -2.3769545555114746, -2.8430075645446777, -2.9142165184020996, -3.017808437347412, -3.0437026023864746, -3.5895285606384277, -4.342545986175537, -4.3558669090271], [-1.278747797012329, -1.6791880130767822, -2.243154764175415, -2.6780436038970947, -2.9457952976226807, -3.3872568607330322, -3.8447487354278564, -3.8778812885284424, -4.074468612670898, -4.290727615356445], [-1.6544435024261475, -2.5747010707855225, -2.6344716548919678, -2.7110230922698975, -2.7909810543060303, -2.9975812435150146, -3.1332204341888428, -3.418637990951538, -3.620513677597046, -3.8255879878997803], [-0.6308883428573608, -2.030764102935791, -2.5993924140930176, -2.7181267738342285, -3.8265442848205566, -4.051680088043213, -4.209730625152588, -4.533054828643799, -4.5756001472473145, -4.666740894317627], [-1.1232950687408447, -2.6717894077301025, -2.8861582279205322, -2.978487253189087, -2.99068284034729, -3.0478804111480713, -3.5038721561431885, -3.5882437229156494, -3.6518137454986572, -3.6590216159820557], [-0.29644775390625, -1.4972686767578125, -4.235706329345703, -6.1917266845703125, -6.29533576965332, -6.35308837890625, -6.444324493408203, -6.609521865844727, -6.973005294799805, -7.24797248840332], [-0.5543937683105469, -2.132505416870117, -2.372150421142578, -2.815032958984375, -3.2442665100097656, -3.6007652282714844, -3.939544677734375, -4.62982177734375, -4.962799072265625, -5.13623046875], [-0.09781533479690552, -3.6927363872528076, -3.9080111980438232, -4.1939592361450195, -5.095574378967285, -5.54971981048584, -6.315329551696777, -6.684817314147949, -6.770205497741699, -6.805417060852051], [-1.2705544233322144, -1.3966587781906128, -2.181002616882324, -2.669062614440918, -2.9217100143432617, -3.26711368560791, -3.413538932800293, -3.883183479309082, -4.156699180603027, -4.2995405197143555], [-0.007819517515599728, -6.477696895599365, -6.587197780609131, -6.76526403427124, -6.938920497894287, -7.21808385848999, -7.243297100067139, -8.155594825744629, -9.206570625305176, -9.448287010192871], [-0.18164236843585968, -2.3973500728607178, -4.385944366455078, -5.066688537597656, -5.128849029541016, -5.214832305908203, -5.355841636657715, -5.812589645385742, -6.640542030334473, -6.641702651977539], [-0.3554803729057312, -1.8430063724517822, -2.637809991836548, -4.232261657714844, -4.252098083496094, -5.280946731567383, -5.581424713134766, -5.7038726806640625, -6.054256439208984, -6.070062637329102], [-1.174675703048706, -1.1769263744354248, -2.3811194896698, -3.574639081954956, -3.6462228298187256, -4.428581237792969, -4.515603065490723, -4.629171371459961, -4.8219404220581055, -4.826866149902344], [-1.2866473197937012, -1.3665785789489746, -1.450305461883545, -2.815547466278076, -3.4129652976989746, -3.5802950859069824, -4.259646892547607, -4.410497188568115, -4.686326503753662, -4.819680690765381], [-1.284511923789978, -1.6003555059432983, -1.8806651830673218, -2.330559253692627, -2.822094440460205, -2.978271961212158, -3.0513157844543457, -3.384586811065674, -4.678221225738525, -4.7065606117248535], [-0.003452533157542348, -6.71318244934082, -7.6864166259765625, -7.866558074951172, -8.395820617675781, -8.797486305236816, -9.47498893737793, -10.08416748046875, -10.08804702758789, -10.378552436828613], [-0.7611655592918396, -1.778707504272461, -2.318513870239258, -3.9089527130126953, -4.0619659423828125, -4.076770782470703, -4.324800491333008, -4.37451171875, -4.376003265380859, -4.376029968261719], [-0.43409377336502075, -2.0129549503326416, -3.2345545291900635, -3.5926191806793213, -3.671215295791626, -3.764213800430298, -4.053459644317627, -4.160521030426025, -4.5704731941223145, -4.845633029937744], [-0.19858379662036896, -1.7301313877105713, -6.944765090942383, -7.969247817993164, -8.085386276245117, -8.19151496887207, -8.449525833129883, -9.160225868225098, -9.588837623596191, -9.984594345092773], [-1.8052799701690674, -2.5017712116241455, -2.586858034133911, -2.8280861377716064, -2.8637516498565674, -2.86714768409729, -3.0553886890411377, -3.674962282180786, -3.7268106937408447, -3.859431505203247], [-0.5131959915161133, -1.924729347229004, -2.9035139083862305, -2.921065330505371, -2.9853811264038086, -3.6293020248413086, -4.092734336853027, -5.427997589111328, -5.525561332702637, -5.584092140197754], [-1.3965237140655518, -1.8620693683624268, -2.1119511127471924, -2.61291241645813, -3.2083914279937744, -3.618133783340454, -3.6918938159942627, -3.8739960193634033, -4.162569046020508, -4.266186714172363], [-1.1275970935821533, -1.3566906452178955, -3.086838960647583, -3.452552080154419, -3.5200703144073486, -3.6519806385040283, -3.8299877643585205, -4.088089942932129, -4.1960649490356445, -4.667270660400391], [-0.4600834846496582, -1.3127923011779785, -3.765890598297119, -4.479422092437744, -4.7028889656066895, -4.795820713043213, -5.030157566070557, -5.468201160430908, -5.512045383453369, -5.532801151275635], [-1.2519283294677734, -1.8291149139404297, -2.6022262573242188, -2.739431381225586, -2.9429473876953125, -3.558868408203125, -3.7740249633789062, -4.217281341552734, -4.22056770324707, -4.551769256591797], [-0.5379815697669983, -2.2664647102355957, -2.621513843536377, -2.739819049835205, -4.065554141998291, -4.312038898468018, -4.4349846839904785, -4.5212788581848145, -4.59879732131958, -4.6161580085754395], [-0.24191200733184814, -2.357131004333496, -2.582028388977051, -5.0107927322387695, -5.192748069763184, -5.412188529968262, -5.6585588455200195, -5.699042320251465, -5.758770942687988, -6.2831010818481445], [-0.050112128257751465, -3.3379197120666504, -5.641367435455322, -6.5138840675354, -6.607151508331299, -6.669376850128174, -7.108365535736084, -7.35122537612915, -8.149378776550293, -8.19774341583252], [-0.2582389712333679, -2.471951723098755, -3.396174669265747, -3.807959794998169, -3.808306932449341, -4.562500953674316, -5.042641639709473, -5.083487510681152, -5.352255821228027, -5.45950984954834], [-0.2726271152496338, -2.563659906387329, -2.6477510929107666, -3.4699480533599854, -3.56538987159729, -4.310754776000977, -6.029754638671875, -6.037014007568359, -6.406652450561523, -6.706075668334961], [-0.0036967527121305466, -5.79642915725708, -8.14664077758789, -9.461244583129883, -9.964672088623047, -10.01359748840332, -10.15234375, -11.19976806640625, -11.389558792114258, -11.412906646728516], [-1.3802396059036255, -1.479141354560852, -1.7676469087600708, -2.4589576721191406, -2.7227916717529297, -3.4538612365722656, -3.771778106689453, -4.313604354858398, -4.533607482910156, -4.815101623535156], [-0.0841309204697609, -3.2644708156585693, -4.228188991546631, -5.183648586273193, -5.896180629730225, -6.039113521575928, -6.171563625335693, -6.318734645843506, -6.36781644821167, -6.743028163909912], [-0.2532283365726471, -2.2264645099639893, -3.574834108352661, -4.541092872619629, -4.606949806213379, -4.977408409118652, -5.337008476257324, -5.488327980041504, -5.846911430358887, -6.011813163757324], [-0.8177525997161865, -1.1370179653167725, -2.4109609127044678, -3.2797353267669678, -4.361211776733398, -4.442050933837891, -4.475515365600586, -4.667137145996094, -4.722789764404297, -4.750612258911133], [-0.0016505202511325479, -7.155468463897705, -8.78048038482666, -9.2362699508667, -9.369120597839355, -9.78431224822998, -9.897849082946777, -10.091814994812012, -10.389817237854004, -10.867840766906738], [-0.2397218495607376, -2.567783832550049, -3.9733777046203613, -4.0111870765686035, -4.055180072784424, -4.189627170562744, -4.272867679595947, -4.729145526885986, -4.93442964553833, -5.013723850250244], [-0.6482017636299133, -1.9497230052947998, -2.744734525680542, -2.9432570934295654, -3.294945478439331, -3.4912421703338623, -3.83294939994812, -3.9057776927948, -4.058547019958496, -4.493414878845215], [-0.9068384170532227, -2.589909553527832, -2.8801259994506836, -2.9731664657592773, -2.984347343444824, -3.2325143814086914, -3.6160764694213867, -3.652400016784668, -3.7895421981811523, -3.795302391052246], [-0.48999089002609253, -2.4962146282196045, -2.6280620098114014, -2.8122050762176514, -3.834912061691284, -4.068287372589111, -4.680363178253174, -4.7075276374816895, -4.838058948516846, -5.1276021003723145], [-0.2099856436252594, -2.396739959716797, -3.7827606201171875, -4.102975845336914, -4.2533721923828125, -4.32032585144043, -5.431510925292969, -5.517753601074219, -5.60052490234375, -5.608646392822266], [-1.2549582719802856, -1.7463961839675903, -1.7806674242019653, -1.8416491746902466, -3.588189125061035, -4.121170997619629, -4.203289985656738, -4.2304792404174805, -4.241677284240723, -4.364739418029785], [-1.1053345203399658, -1.2426445484161377, -2.0468666553497314, -2.3002541065216064, -3.6286728382110596, -3.8687584400177, -4.098183631896973, -4.374518394470215, -4.560420036315918, -4.8011064529418945], [-0.11845239996910095, -3.2166216373443604, -3.4023592472076416, -4.34104585647583, -5.193403720855713, -5.461124897003174, -5.67779016494751, -5.840208530426025, -6.054356098175049, -6.985493183135986], [-1.5437091588974, -1.820473074913025, -1.8392070531845093, -2.431112766265869, -2.857800006866455, -3.21336030960083, -3.728994846343994, -3.8952889442443848, -4.184719562530518, -4.387455463409424], [-1.3339064121246338, -1.81770920753479, -1.9535925388336182, -2.603487253189087, -2.773926019668579, -2.8758318424224854, -2.9865000247955322, -3.459209680557251, -3.6767351627349854, -4.189197540283203], [-1.3527289628982544, -1.379508137702942, -1.534403920173645, -3.617396354675293, -3.7183008193969727, -4.081496238708496, -4.251824378967285, -4.358290672302246, -4.5153303146362305, -4.761061668395996], [-0.8167728185653687, -2.3844056129455566, -3.1805481910705566, -3.185965061187744, -3.335096836090088, -3.715414524078369, -4.072058200836182, -4.121044635772705, -4.12502908706665, -4.142355442047119], [-1.7069964408874512, -1.9735217094421387, -2.8056349754333496, -3.0393080711364746, -3.0493712425231934, -3.0742087364196777, -3.0905661582946777, -3.3695883750915527, -3.4065794944763184, -3.4762511253356934], [-0.7300421595573425, -1.9341418743133545, -1.9869163036346436, -3.066288709640503, -3.2312915325164795, -3.742283582687378, -4.1482014656066895, -4.171370029449463, -4.528067111968994, -4.687931537628174], [-0.9212523698806763, -2.6654233932495117, -2.924098014831543, -3.02217960357666, -3.2975587844848633, -3.307803153991699, -3.327866554260254, -3.7506399154663086, -3.7806501388549805, -3.8221845626831055], [-0.17896340787410736, -3.3544814586639404, -4.439797878265381, -4.7614216804504395, -4.931570529937744, -5.003124713897705, -5.3462090492248535, -5.386402606964111, -5.432949542999268, -5.478273868560791], [-0.18972711265087128, -2.5467216968536377, -4.389462947845459, -4.618724346160889, -4.668748378753662, -5.214345455169678, -5.261071681976318, -5.340207576751709, -5.428225994110107, -5.518626689910889], [-0.7267122864723206, -1.9803457260131836, -2.622727394104004, -2.9903154373168945, -2.998072624206543, -3.2191457748413086, -3.8306455612182617, -3.894312858581543, -3.975949287414551, -4.350564002990723], [-0.020010385662317276, -5.251017093658447, -5.67455530166626, -5.886524677276611, -6.079123020172119, -6.6986165046691895, -6.785297870635986, -7.16776704788208, -7.767279148101807, -8.13803482055664], [-0.7513647675514221, -1.7059850692749023, -2.2032651901245117, -3.0225114822387695, -3.234866142272949, -3.852593421936035, -4.09366512298584, -4.310198783874512, -4.543339729309082, -4.547806739807129], [-0.8763109445571899, -1.5423017740249634, -1.8642431497573853, -2.4766430854797363, -3.7231316566467285, -3.9272236824035645, -4.459427356719971, -4.498356342315674, -4.714981555938721, -4.9154839515686035], [-0.00046659549116156995, -8.833687782287598, -8.882493019104004, -10.132534980773926, -10.212037086486816, -11.393338203430176, -11.408698081970215, -11.511855125427246, -11.546452522277832, -12.114046096801758], [-1.5360831022262573, -2.108304977416992, -2.2858123779296875, -2.431581497192383, -2.8270797729492188, -3.043304443359375, -3.1750946044921875, -3.4131622314453125, -3.6463184356689453, -3.6590728759765625], [-1.2832740545272827, -1.382341742515564, -1.5589431524276733, -2.1615853309631348, -3.652876377105713, -3.682203769683838, -3.9527363777160645, -4.271149158477783, -4.416624546051025, -5.082592487335205], [-1.2193286418914795, -1.3985011577606201, -1.9321696758270264, -2.4761435985565186, -2.815045118331909, -2.8223140239715576, -3.122265577316284, -4.1949872970581055, -4.248230934143066, -4.835675239562988], [-0.8976349234580994, -1.5094265937805176, -2.025852680206299, -2.90116548538208, -3.4823861122131348, -3.9814553260803223, -4.281469821929932, -4.530375003814697, -4.63230562210083, -4.833662509918213], [-0.5017741918563843, -1.735287070274353, -1.9957107305526733, -4.081532001495361, -4.8460869789123535, -4.859695911407471, -5.334743976593018, -5.365137577056885, -5.4647088050842285, -5.676289081573486], [-1.3731788396835327, -2.5181336402893066, -3.075880527496338, -3.1481194496154785, -3.295963764190674, -3.4699769020080566, -3.5433621406555176, -3.543529987335205, -3.7463364601135254, -3.8233399391174316], [-0.109302818775177, -3.57485032081604, -3.641247034072876, -4.436283588409424, -4.822815418243408, -5.259323596954346, -5.804226398468018, -6.221785068511963, -6.34154748916626, -6.366329669952393], [-0.03258315473794937, -3.510650634765625, -7.259967803955078, -7.543283462524414, -7.955667495727539, -8.084836959838867, -8.469537734985352, -10.914033889770508, -12.09141731262207, -12.289237976074219], [-1.7006688117980957, -2.4590649604797363, -2.4627881050109863, -2.4675965309143066, -2.7641453742980957, -3.0670418739318848, -3.1854519844055176, -3.293827533721924, -3.3986802101135254, -3.563965320587158], [-0.291087806224823, -1.8130090236663818, -4.385949611663818, -4.612578868865967, -4.788232326507568, -5.132178783416748, -5.208011150360107, -5.2184977531433105, -5.474307537078857, -5.593989849090576], [-0.719337522983551, -0.9613209366798401, -2.8284263610839844, -4.439323425292969, -4.978010177612305, -5.291358947753906, -5.8328752517700195, -6.318382263183594, -6.369429588317871, -6.421497344970703], [-1.7802187204360962, -2.1415505409240723, -2.3215088844299316, -2.900228977203369, -2.9053006172180176, -3.037436008453369, -3.088858127593994, -3.136812686920166, -3.419224262237549, -3.5133137702941895], [-0.8869927525520325, -1.9240639209747314, -2.519213914871216, -3.080664873123169, -3.316546678543091, -3.67718243598938, -4.00349760055542, -4.049285411834717, -4.193301677703857, -4.2090067863464355], [-2.1019766330718994, -2.2705576419830322, -2.502298593521118, -2.7492926120758057, -2.838247537612915, -2.9228575229644775, -3.1833651065826416, -3.183605432510376, -3.2752249240875244, -3.5262339115142822], [-1.239345908164978, -1.9251044988632202, -2.3547024726867676, -3.1700310707092285, -3.5397591590881348, -3.562124729156494, -3.980347156524658, -4.103588581085205, -4.113361835479736, -4.117939472198486], [-0.09708235412836075, -3.6577627658843994, -4.277444839477539, -4.357044219970703, -4.504293441772461, -4.634490966796875, -5.393272399902344, -5.601314544677734, -6.300224304199219, -6.547618865966797], [-0.3409811854362488, -2.5129554271698, -2.8377902507781982, -3.1673457622528076, -3.9924628734588623, -4.179958820343018, -4.2811970710754395, -4.36254358291626, -4.719370365142822, -5.219009876251221], [-0.7249959707260132, -2.150012493133545, -2.2308554649353027, -2.3732733726501465, -2.387692928314209, -3.5528311729431152, -3.7483725547790527, -4.235726833343506, -4.596755504608154, -5.085277080535889], [-0.9774425029754639, -1.8530681133270264, -2.411982297897339, -2.6150577068328857, -2.7503421306610107, -3.4215800762176514, -3.88342547416687, -3.886164426803589, -4.383856773376465, -4.672942161560059], [-0.04679347202181816, -3.96494722366333, -4.641136646270752, -4.795372486114502, -6.8315510749816895, -6.982397556304932, -7.268871784210205, -7.3839497566223145, -7.489758014678955, -7.857241153717041], [-0.31035345792770386, -2.953282594680786, -3.2479069232940674, -3.761418581008911, -4.112614631652832, -4.1251020431518555, -4.619383811950684, -4.64388370513916, -4.796278953552246, -4.85745906829834], [-0.26277634501457214, -2.1181726455688477, -2.456051826477051, -5.3911638259887695, -5.914050102233887, -6.005043983459473, -6.199723243713379, -6.254300117492676, -6.544785499572754, -6.592713356018066], [-0.5433505773544312, -2.7175393104553223, -2.842195987701416, -3.76133394241333, -3.8736729621887207, -3.977125644683838, -4.064489841461182, -4.075668811798096, -4.150932788848877, -4.173208713531494], [-1.4341762065887451, -1.5299479961395264, -1.9002444744110107, -2.6758296489715576, -2.7313029766082764, -2.8934810161590576, -2.8967063426971436, -3.099146604537964, -3.6680171489715576, -4.6420488357543945], [-1.220574975013733, -1.2577797174453735, -2.7986960411071777, -2.9994406700134277, -3.221740245819092, -3.6000800132751465, -3.7003417015075684, -3.7338995933532715, -4.461680889129639, -4.544175624847412], [-0.23072941601276398, -3.5346615314483643, -3.7650253772735596, -4.1820807456970215, -4.230477809906006, -4.369567394256592, -4.410802364349365, -4.443854808807373, -4.6673903465271, -4.708890438079834], [-0.1686599999666214, -2.5616137981414795, -2.7843387126922607, -5.496199607849121, -5.517390251159668, -6.10916805267334, -7.133181571960449, -7.469878196716309, -7.567255973815918, -7.571822166442871], [-0.03456159308552742, -3.5091214179992676, -6.727182865142822, -7.896212100982666, -7.901533603668213, -8.239046096801758, -8.441466331481934, -8.540034294128418, -8.729623794555664, -8.793167114257812], [-0.041516028344631195, -4.718027591705322, -4.776236057281494, -5.768055438995361, -5.774538516998291, -5.931726932525635, -6.047996997833252, -6.555778980255127, -6.7127366065979, -7.016510486602783], [-0.16437621414661407, -2.4174718856811523, -3.868739128112793, -4.466994285583496, -4.826800346374512, -6.00504207611084, -6.166232109069824, -6.271721839904785, -6.527167320251465, -6.677702903747559], [-0.9411216974258423, -1.8419796228408813, -2.60158109664917, -2.783629894256592, -2.8400683403015137, -3.357631206512451, -3.8178515434265137, -4.265321254730225, -4.463851451873779, -4.62208890914917], [-1.1664338111877441, -1.4504914283752441, -1.9690041542053223, -2.520834445953369, -2.894716739654541, -3.2474350929260254, -4.2989277839660645, -4.555185794830322, -4.65734338760376, -4.66165018081665], [-0.6082432866096497, -3.0967419147491455, -3.2729732990264893, -3.279273271560669, -3.345717668533325, -3.5865261554718018, -3.672593355178833, -3.872760057449341, -4.08180570602417, -4.164279460906982], [-1.1109569072723389, -2.094578504562378, -2.3004329204559326, -2.5082404613494873, -2.611309766769409, -3.4852187633514404, -3.604609251022339, -3.620774030685425, -3.7734286785125732, -3.90071177482605], [-0.6794657707214355, -2.0996108055114746, -2.626533031463623, -2.848659038543701, -3.068887233734131, -3.1731066703796387, -4.053109645843506, -4.519050121307373, -4.584981441497803, -5.241530895233154], [-1.4911484718322754, -1.6066193580627441, -2.219088077545166, -2.6915860176086426, -3.222395420074463, -3.320692539215088, -3.3744683265686035, -3.4580636024475098, -3.7232441902160645, -3.782402515411377], [-0.8373185396194458, -1.3131829500198364, -2.0079612731933594, -3.3321914672851562, -3.5005569458007812, -3.81939697265625, -4.4671783447265625, -4.592998504638672, -4.895544052124023, -5.025142669677734], [-0.0037339269183576107, -6.233117580413818, -7.138347148895264, -8.172951698303223, -8.672486305236816, -9.589415550231934, -9.762785911560059, -10.199854850769043, -10.444130897521973, -10.599610328674316], [-1.1195952892303467, -1.2713172435760498, -2.1387641429901123, -2.4109504222869873, -2.4467132091522217, -3.8374545574188232, -4.7074689865112305, -4.770272254943848, -4.886166572570801, -5.233708381652832], [-1.1844749450683594, -1.3241844177246094, -2.4885005950927734, -2.9741172790527344, -3.170734405517578, -3.346952438354492, -4.089082717895508, -4.3595733642578125, -4.749689102172852, -4.86265754699707], [-1.0153727531433105, -2.1631999015808105, -2.7604708671569824, -2.8160014152526855, -3.2777857780456543, -3.427405834197998, -3.5600314140319824, -4.066863536834717, -4.097205638885498, -4.10852575302124], [-0.6537575125694275, -1.7951149940490723, -2.761894702911377, -3.5116944313049316, -3.9450478553771973, -3.996330738067627, -4.053154468536377, -4.572994709014893, -4.786819934844971, -4.798311710357666], [-0.7565585970878601, -1.420966386795044, -1.9133732318878174, -3.0086758136749268, -3.823054552078247, -4.280617713928223, -4.372963905334473, -5.051085472106934, -5.195734977722168, -5.829764366149902], [-0.8668451309204102, -1.5346269607543945, -1.8505373001098633, -3.065587043762207, -3.361924171447754, -3.3716936111450195, -3.9786481857299805, -4.682793617248535, -4.948409080505371, -5.08335018157959], [-0.14208939671516418, -3.7619662284851074, -3.879697322845459, -4.157227993011475, -4.400855541229248, -4.402129650115967, -4.683719158172607, -6.124358654022217, -6.445504665374756, -6.534369945526123], [-0.0002898749662563205, -9.769424438476562, -9.807418823242188, -10.319768905639648, -10.944612503051758, -10.993598937988281, -11.133970260620117, -11.22636890411377, -11.607852935791016, -11.779435157775879], [-0.5908623933792114, -2.5130348205566406, -2.765125274658203, -3.317281723022461, -3.5334625244140625, -4.115068435668945, -4.217123031616211, -4.22657585144043, -4.311502456665039, -4.336002349853516], [-0.17303341627120972, -3.0575225353240967, -3.922276258468628, -4.146698474884033, -4.581213474273682, -4.748531818389893, -4.849279880523682, -5.442572116851807, -5.494207859039307, -5.858719348907471], [-0.4503314197063446, -1.7294681072235107, -1.9426944255828857, -5.353464603424072, -5.516043186187744, -5.613363742828369, -5.796005725860596, -6.252970218658447, -6.364943027496338, -6.636854648590088], [-1.476321816444397, -1.561669945716858, -1.6821247339248657, -2.082017421722412, -2.1354479789733887, -2.613887310028076, -4.449899196624756, -4.539252758026123, -4.73621129989624, -5.2259039878845215], [-0.47746407985687256, -1.336346983909607, -3.1393275260925293, -4.509477138519287, -4.707902431488037, -4.786777019500732, -5.018617153167725, -5.558213710784912, -5.659398555755615, -6.044270992279053], [-1.6391061544418335, -2.1922335624694824, -2.774014949798584, -2.883714199066162, -2.9082236289978027, -2.9262022972106934, -3.0423922538757324, -3.3938002586364746, -3.4541468620300293, -3.458777904510498], [-0.02097589708864689, -4.337897300720215, -6.0606584548950195, -6.6927900314331055, -7.1293535232543945, -7.134495735168457, -7.27060604095459, -8.015889167785645, -8.567399024963379, -8.801312446594238], [-1.750502109527588, -2.045637607574463, -2.1107640266418457, -2.6134114265441895, -2.8450303077697754, -2.9110817909240723, -3.455650806427002, -3.7072300910949707, -3.708134174346924, -3.850904941558838], [-1.093828558921814, -1.169552206993103, -1.618536353111267, -2.87418794631958, -3.158970355987549, -3.925438404083252, -4.961788654327393, -5.144516468048096, -5.226093769073486, -5.271728992462158], [-0.6131972670555115, -1.8658866882324219, -2.665773391723633, -2.8503684997558594, -3.678285598754883, -4.053487777709961, -4.063518524169922, -4.2359466552734375, -4.6979217529296875, -4.9517364501953125], [-0.5003752112388611, -2.489387035369873, -2.9834933280944824, -3.7678942680358887, -3.893876552581787, -4.040996074676514, -4.1345038414001465, -4.317748546600342, -4.766955852508545, -5.0517497062683105], [-0.4996608793735504, -1.5757774114608765, -3.262159824371338, -3.665256977081299, -3.8835339546203613, -3.91666841506958, -3.9641880989074707, -4.392263889312744, -4.9200921058654785, -5.414328098297119], [-0.5248868465423584, -1.9127118587493896, -2.4744832515716553, -3.450437307357788, -3.669299840927124, -3.9967362880706787, -4.019338607788086, -4.04057502746582, -4.48675537109375, -4.563028335571289], [-0.10120955854654312, -2.9748401641845703, -3.215791702270508, -6.222570419311523, -6.512105941772461, -7.740865707397461, -8.747190475463867, -8.781150817871094, -8.82335090637207, -9.702880859375], [-4.768370445162873e-07, -15.009653091430664, -16.646602630615234, -16.90728187561035, -17.88405990600586, -18.06052589416504, -18.22677230834961, -18.581607818603516, -18.613759994506836, -18.632862091064453], [-0.9634779691696167, -2.034724235534668, -2.7465476989746094, -2.75234317779541, -3.5132293701171875, -3.558762550354004, -3.938291549682617, -3.9997920989990234, -4.167990684509277, -4.170927047729492], [-2.3194642066955566, -2.3558993339538574, -3.1568589210510254, -3.2147393226623535, -3.383246898651123, -3.615540027618408, -3.8907132148742676, -3.9164681434631348, -4.205816745758057, -4.302371501922607], [-1.5479261875152588, -2.8916876316070557, -3.0237953662872314, -3.124971628189087, -3.2257044315338135, -3.494793176651001, -3.6809093952178955, -3.729083299636841, -4.260154724121094, -4.283388137817383], [-0.509814441204071, -2.6297810077667236, -2.7403366565704346, -2.9497692584991455, -3.5784332752227783, -3.7098934650421143, -3.7277138233184814, -4.083382606506348, -4.478659629821777, -4.5356645584106445], [-1.2678115367889404, -1.8093516826629639, -1.8235232830047607, -2.4452788829803467, -3.087082624435425, -3.365715742111206, -3.450453519821167, -3.870945692062378, -3.918255567550659, -4.0189714431762695], [-0.9515280723571777, -1.9114794731140137, -2.832017421722412, -3.0379271507263184, -3.337533473968506, -3.742887020111084, -3.8735642433166504, -4.123470783233643, -4.337832927703857, -4.367605686187744], [-2.171112060546875, -2.391653060913086, -2.7205562591552734, -3.2578487396240234, -3.4893417358398438, -3.6892318725585938, -3.927274703979492, -3.9953651428222656, -4.093984603881836, -4.14581298828125], [-2.4455368518829346, -3.310229539871216, -3.437636613845825, -3.5179636478424072, -3.666815996170044, -3.7329695224761963, -3.8520424365997314, -3.8666489124298096, -3.876596689224243, -3.926138162612915], [-0.7217168211936951, -1.3622503280639648, -2.69547176361084, -3.031609535217285, -3.521487236022949, -4.093205451965332, -4.318233489990234, -4.790000915527344, -4.795565605163574, -4.898306846618652], [-1.4549778699874878, -1.9336727857589722, -3.012394905090332, -3.132857322692871, -3.35697078704834, -3.3767499923706055, -3.4579572677612305, -3.627322196960449, -3.715487480163574, -3.8149595260620117], [-0.928566038608551, -1.7140579223632812, -2.603656768798828, -2.8133316040039062, -2.865650177001953, -3.3873939514160156, -3.5062122344970703, -3.648761749267578, -4.004087448120117, -4.115785598754883], [-1.6350857019424438, -2.005673885345459, -2.3206686973571777, -2.5934481620788574, -2.6495261192321777, -2.7588934898376465, -2.7885851860046387, -2.84981107711792, -3.4200778007507324, -3.8958468437194824], [-1.1532529592514038, -2.265956401824951, -2.4496169090270996, -2.573607921600342, -3.124345302581787, -3.4526925086975098, -3.6920504570007324, -3.8715567588806152, -4.009006023406982, -4.036871433258057], [-0.8533508777618408, -1.472703218460083, -2.4125759601593018, -2.9383881092071533, -2.9886906147003174, -3.2573254108428955, -3.607814073562622, -4.065762519836426, -4.508149147033691, -4.9412126541137695], [-1.1067936420440674, -1.4004642963409424, -2.7948715686798096, -3.1332638263702393, -3.7391464710235596, -3.763219118118286, -4.024504661560059, -4.059776306152344, -4.326423645019531, -4.336912155151367], [-0.3188457787036896, -2.4581260681152344, -3.908651351928711, -4.330844879150391, -4.500626564025879, -4.6908979415893555, -4.797201156616211, -5.000249862670898, -5.222867965698242, -5.3551025390625], [-0.8249402046203613, -1.939770221710205, -2.3196702003479004, -2.598604679107666, -2.8208889961242676, -3.946281909942627, -4.101963520050049, -4.345675945281982, -4.736853122711182, -5.155606746673584], [-0.7935125827789307, -2.069671869277954, -3.827415704727173, -3.9018747806549072, -3.939485788345337, -4.121503829956055, -4.289074897766113, -4.335172653198242, -4.424995422363281, -4.42669677734375], [-0.32164621353149414, -1.5706028938293457, -3.779066562652588, -3.8708577156066895, -5.218122959136963, -5.429989337921143, -6.545755863189697, -6.568811893463135, -6.768260478973389, -6.956370830535889], [-1.003499150276184, -2.687191963195801, -2.758955955505371, -3.494236946105957, -3.5270328521728516, -3.8594179153442383, -3.875253677368164, -3.8897104263305664, -4.095883369445801, -4.098232269287109], [-0.5608841180801392, -2.2943644523620605, -3.1340365409851074, -3.9575676918029785, -4.051490306854248, -4.403706073760986, -4.464087963104248, -4.687629222869873, -4.732462406158447, -5.0428242683410645], [-0.8024858236312866, -2.829089641571045, -3.211879253387451, -3.391812801361084, -3.670717716217041, -3.783231258392334, -3.8098196983337402, -3.9890332221984863, -4.193538188934326, -4.4100661277771], [-1.7592356204986572, -1.9216139316558838, -2.155390977859497, -2.5360605716705322, -3.9073526859283447, -4.033847808837891, -4.157377243041992, -4.460928916931152, -4.479830741882324, -4.518803596496582], [-0.9153603911399841, -1.3166208267211914, -2.6429128646850586, -2.8648805618286133, -3.445237159729004, -3.515660285949707, -3.517298698425293, -4.238512992858887, -4.666811943054199, -4.833283424377441], [-0.8651371002197266, -1.54229736328125, -2.7010269165039062, -2.711465835571289, -3.864898681640625, -4.455408096313477, -4.948841094970703, -5.005257606506348, -5.079590797424316, -5.1587982177734375], [-0.7731188535690308, -3.500368118286133, -3.5553665161132812, -3.5704174041748047, -3.7730846405029297, -4.074491500854492, -4.336993217468262, -4.725822448730469, -4.759237289428711, -4.848731994628906], [-2.587486743927002, -3.220287799835205, -3.265594959259033, -3.3050618171691895, -3.3536458015441895, -3.3653206825256348, -3.451712131500244, -3.5286126136779785, -3.535789966583252, -3.6473450660705566], [-1.966299057006836, -2.095020294189453, -2.0976333618164062, -2.108213424682617, -2.695648193359375, -3.207305908203125, -3.3674163818359375, -3.511425018310547, -3.5776424407958984, -3.7452011108398438], [-0.9955637454986572, -1.3299105167388916, -2.001983880996704, -3.7754557132720947, -3.804365396499634, -3.977809190750122, -4.028399467468262, -4.767955780029297, -4.9801225662231445, -4.997615814208984], [-0.9364675283432007, -1.1900686025619507, -2.182166576385498, -3.261643886566162, -4.056382656097412, -4.183980464935303, -4.428506374359131, -4.436883449554443, -4.692630290985107, -4.720468044281006], [-1.603662371635437, -2.2296066284179688, -2.953805923461914, -3.288052558898926, -3.507845878601074, -3.5429210662841797, -3.5826330184936523, -3.766613006591797, -3.838237762451172, -4.1531877517700195], [-0.8462379574775696, -1.6531550884246826, -1.6723926067352295, -2.608097791671753, -4.05698299407959, -4.460982322692871, -4.486870765686035, -4.9635515213012695, -5.140688896179199, -5.224724769592285], [-1.065568208694458, -2.0943844318389893, -2.1396725177764893, -2.6119625568389893, -3.013720750808716, -3.027959108352661, -3.326831102371216, -3.5820438861846924, -4.101555824279785, -4.295680046081543], [-0.9680259227752686, -1.9687373638153076, -3.1111209392547607, -3.258896589279175, -3.6214053630828857, -3.657491445541382, -3.832491636276245, -3.884389638900757, -4.020899772644043, -4.1801652908325195], [-1.3465521335601807, -1.4725191593170166, -2.9039080142974854, -3.4816410541534424, -3.4852802753448486, -3.6336138248443604, -3.6583521366119385, -3.7884113788604736, -4.055657386779785, -4.09886360168457], [-0.20639769732952118, -2.641072988510132, -3.4455068111419678, -3.996896505355835, -4.044372081756592, -5.007638454437256, -5.136912822723389, -5.801711559295654, -5.996934413909912, -6.154557704925537], [-0.1959957629442215, -1.9113335609436035, -4.625126361846924, -5.152451038360596, -5.468470096588135, -5.7141594886779785, -6.312652111053467, -7.214084148406982, -7.27617883682251, -7.899289608001709], [-0.8202694654464722, -1.7166012525558472, -2.189805030822754, -3.0368471145629883, -3.625382423400879, -3.78731632232666, -3.8527956008911133, -3.9533281326293945, -4.0142107009887695, -4.185128211975098], [-0.052747633308172226, -3.355314254760742, -5.081544876098633, -6.354413986206055, -6.50160026550293, -6.867157936096191, -6.906248092651367, -7.23007869720459, -7.279658317565918, -7.505866050720215], [-0.20388591289520264, -2.4745750427246094, -4.065881729125977, -4.1504974365234375, -4.277046203613281, -4.454280853271484, -4.961923599243164, -5.195476531982422, -5.232025146484375, -5.636140823364258], [-1.143463134765625, -1.4821357727050781, -2.4257431030273438, -3.2510194778442383, -3.505624771118164, -4.219663619995117, -4.289992332458496, -4.553607940673828, -4.5758161544799805, -4.618678092956543], [-0.013952808454632759, -4.854330539703369, -6.877737522125244, -7.375034809112549, -7.783509731292725, -8.322366714477539, -8.37386417388916, -8.458611488342285, -8.699201583862305, -8.817965507507324], [-2.1377251148223877, -2.157262086868286, -2.3355114459991455, -2.47631573677063, -2.66723370552063, -3.0581886768341064, -3.184077501296997, -3.4214928150177, -3.4249699115753174, -3.5257275104522705], [-0.14188191294670105, -2.6490936279296875, -4.979707717895508, -5.086414337158203, -5.405952453613281, -5.988659858703613, -6.213647842407227, -6.231888771057129, -6.284070014953613, -6.319009780883789], [-1.2951977252960205, -2.3064815998077393, -2.596562623977661, -2.7419216632843018, -2.888667345046997, -3.9139444828033447, -4.295862197875977, -4.311478614807129, -4.387269020080566, -4.395474433898926], [-0.8956736922264099, -2.3381900787353516, -2.6800975799560547, -2.9655284881591797, -3.6077518463134766, -3.6957340240478516, -3.8474884033203125, -3.8850440979003906, -3.986766815185547, -4.223785400390625], [-0.4706301987171173, -1.8411955833435059, -1.8983187675476074, -4.480752468109131, -4.5081915855407715, -4.704736232757568, -5.019006252288818, -5.580361843109131, -5.720889568328857, -5.744185924530029], [-0.23656286299228668, -2.469434976577759, -2.4718878269195557, -4.273067474365234, -5.49159049987793, -5.98747444152832, -6.138788223266602, -6.188631057739258, -6.302213668823242, -6.544244766235352], [-0.42535340785980225, -2.7721896171569824, -2.8670668601989746, -3.2564597129821777, -3.717986583709717, -4.038154125213623, -4.114224910736084, -4.304599285125732, -4.427661418914795, -4.567361354827881], [-1.8299320936203003, -2.492424964904785, -3.027043342590332, -3.230889320373535, -3.3434972763061523, -3.3455820083618164, -3.400029182434082, -3.4040613174438477, -3.535750389099121, -3.8819074630737305], [-1.115234136581421, -2.8380773067474365, -3.104570150375366, -3.1606156826019287, -3.4197957515716553, -3.625608205795288, -3.711219549179077, -3.8638875484466553, -4.068668365478516, -4.0814971923828125], [-1.6362870931625366, -1.7251046895980835, -1.7372373342514038, -2.5480685234069824, -2.6541266441345215, -3.454556941986084, -3.4559969902038574, -3.5205302238464355, -4.1082892417907715, -4.346578121185303]], \"topKTokens\": [[\"\\n\", \",\", \"Words\", \"Summary\", \"\\n\\n\", \"<|endoftext|>\", \" \", \"Features\", \"Random\", \" the\"], [\" upon\", \" there\", \",\", \" the\", \" she\", \" on\", \" he\", \" day\", \" in\", \" night\"], [\" a\", \" an\", \" the\", \" one\", \" time\", \" some\", \" upon\", \" it\", \" two\", \" something\"], [\" time\", \" day\", \" night\", \" morning\", \" week\", \" Sunday\", \" long\", \" Wednesday\", \" dark\", \" little\"], [\",\", \" there\", \" in\", \" a\", \" was\", \" upon\", \".\", \" on\", \" the\", \" lived\"], [\" there\", \" a\", \" in\", \" two\", \" the\", \" she\", \" it\", \" Jack\", \" he\", \" Sam\"], [\" was\", \" were\", \" lived\", \" are\", \" is\", \"'s\", \",\", \" a\", \" wasn\", \" had\"], [\" a\", \" an\", \" two\", \" one\", \" something\", \" some\", \" the\", \" three\", \" his\", \" Jack\"], [\" little\", \" girl\", \" boy\", \" big\", \" brave\", \" small\", \" man\", \" young\", \" kind\", \" happy\"], [\"my\", \" and\", \",\", \" who\", \".\", \"ma\", \" bear\", \" named\", \"o\", \" called\"], [\" and\", \" who\", \",\", \" bear\", \".\", \" bird\", \" pig\", \" duck\", \" rabbit\", \" bunny\"], [\" a\", \" her\", \" daddy\", \" she\", \" their\", \" his\", \" baby\", \" the\", \" an\", \" some\"], [\" baby\", \" daddy\", \" little\", \" dad\", \" boy\", \" daughter\", \" girl\", \" Daddy\", \" child\", \" man\"], [\",\", \" house\", \" bear\", \" tree\", \" dog\", \" garden\", \" truck\", \" park\", \" box\", \" mountain\"], [\" strong\", \" green\", \" tall\", \" expensive\", \" brown\", \" long\", \" happy\", \" red\", \" loud\", \" shiny\"], [\" house\", \" castle\", \" building\", \" car\", \" boat\", \" theater\", \" store\", \" TV\", \" train\", \" hotel\"], [\".\", \" with\", \"!\", \" in\", \" to\", \" that\", \",\", \" for\", \"!\\\"\", \" and\"], [\" It\", \" They\", \" She\", \" The\", \" Every\", \" One\", \" People\", \" In\", \" Everywhere\", \" This\"], [\" were\", \" had\", \" went\", \" wanted\", \" all\", \" flew\", \" needed\", \" both\", \" watched\", \" took\"], [\" lots\", \" a\", \" many\", \" planes\", \" the\", \" people\", \" so\", \" that\", \" big\", \" it\"], [\" people\", \" planes\", \" airplanes\", \" things\", \" big\", \" exciting\", \" different\", \" interesting\", \" aer\", \" new\"], [\" and\", \" walking\", \" going\", \" there\", \" rushing\", \",\", \" who\", \" taking\", \" in\", \" working\"], [\" around\", \" and\", \" by\", \" on\", \" in\", \" down\", \",\", \" through\", \" there\", \" from\"], [\" and\", \".\", \",\", \" the\", \" with\", \" in\", \" on\", \" them\", \" looking\", \"!\"], [\" talking\", \" looking\", \" having\", \" singing\", \" pointing\", \" were\", \" heard\", \" around\", \" taking\", \" wondered\"], [\" at\", \" to\", \".\", \" hello\", \" their\", \",\", \" flags\", \" back\", \" and\", \" for\"], [\" everyone\", \" them\", \" all\", \" people\", \" the\", \" their\", \" each\", \" some\", \" planes\", \" it\"], [\" the\", \" of\", \" their\", \".\", \" them\", \",\", \" and\", \" kinds\", \" they\", \" different\"], [\" people\", \" planes\", \" airplanes\", \" little\", \" different\", \" things\", \" big\", \" nice\", \" shops\", \" buildings\"], [\".\", \" in\", \",\", \" they\", \" and\", \" flying\", \" as\", \" before\", \" that\", \" she\"], [\".\", \"!\", \",\", \" and\", \" quickly\", \" before\", \" when\", \" much\", \" far\", \" excited\"], [\"\\n\", \" \", \" The\", \" They\", \" Then\", \" Suddenly\", \" Everyone\", \" One\", \" It\", \" Some\"], [\"\\n\", \"The\", \"Suddenly\", \" \", \"Then\", \"When\", \"Mom\", \"One\", \"\\\"\", \"At\"], [\"The\", \"One\", \"Suddenly\", \"Then\", \"When\", \"\\\"\", \"Mom\", \"At\", \"L\", \"Finally\"], [\" little\", \" family\", \" plane\", \" mom\", \" people\", \" airport\", \" pilot\", \" planes\", \" next\", \" modern\"], [\" was\", \" got\", \" had\", \" wanted\", \" went\", \" decided\", \" saw\", \" watched\", \" kept\", \" were\"], [\" the\", \" their\", \",\", \" a\", \" lots\", \" if\", \" one\", \" her\", \" each\", \" his\"], [\" little\", \" dad\", \" airport\", \" pilot\", \" mom\", \" man\", \" family\", \" people\", \" children\", \" 3\"], [\" girl\", \" boy\", \" 3\", \" one\", \" three\", \" man\", \" child\", \" plane\", \" bird\", \" daughter\"], [\",\", \" if\", \" to\", \" what\", \" who\", \" where\", \" why\", \" for\", \" a\", \" and\"], [\" her\", \" mom\", \" she\", \" the\", \" daddy\", \" Mom\", \" dad\", \" Daddy\", \" mum\", \" his\"], [\" if\", \",\", \" to\", \" what\", \" why\", \" lots\", \" many\", \" where\", \" and\", \" for\"], [\" of\", \"-\", \".\", \" and\", \" about\", \" to\", \",\", \" more\", \" at\", \" so\"], [\" questions\", \" people\", \" planes\", \" luggage\", \" other\", \" things\", \" friends\", \" different\", \" mom\", \" money\"], [\".\", \" about\", \",\", \" to\", \" and\", \" for\", \" in\", \"!\", \" all\", \" every\"], [\" but\", \" and\", \" so\", \" while\", \" like\", \" \\\"\", \" the\", \" pointing\", \" she\", \" then\"], [\" they\", \" the\", \" she\", \" he\", \" her\", \" one\", \" each\", \" many\", \" that\", \" their\"], [\" little\", \" family\", \" mom\", \" dad\", \" parents\", \" daddy\", \" girl\", \" mother\", \" plane\", \" father\"], [\"my\", \"m\", \" and\", \" said\", \"ma\", \" answered\", \" was\", \" took\", \"ie\", \" made\"], [\" and\", \" said\", \" answered\", \" told\", \" was\", \" took\", \" got\", \" showed\", \" made\", \" decided\"], [\",\", \" that\", \".\", \" to\", \" with\", \" \\\"\", \" and\", \":\", \" in\", \",\\\"\"], [\" \\\"\", \" \\ufffd\", \" \\\"'\", \" '\", \" \", \"\\n\", \" ''\", \" pointing\", \" \\\"\\\"\", \"\\ufffd\"], [\"Yes\", \"I\", \"We\", \"That\", \"What\", \"No\", \"Let\", \"You\", \"It\", \"The\"], [\" we\", \" you\", \" they\", \" there\", \" the\", \" these\", \" those\", \" planes\", \" all\", \" an\"], [\" there\", \" going\", \" all\", \" here\", \" the\", \" having\", \" ready\", \" almost\", \" getting\", \" flying\"], [\" having\", \" going\", \" discussing\", \" need\", \" getting\", \" have\", \" like\", \" talking\", \" finished\", \" just\"], [\" to\", \" on\", \" for\", \" out\", \" somewhere\", \" through\", \" home\", \" up\", \" the\", \" in\"], [\" see\", \" the\", \" take\", \" get\", \" go\", \" work\", \" fly\", \" a\", \" have\", \" look\"], [\" the\", \"?\\\"\", \" a\", \" this\", \" some\", \" and\", \" an\", \",\", \" more\", \" something\"], [\"\\n\", \" \", \" The\", \" They\", \" and\", \" Daddy\", \" So\", \" Her\", \" She\", \" And\"], [\"\\n\", \"The\", \"Mom\", \"Daddy\", \"Her\", \" \", \"So\", \"\\\"\", \" The\", \"Finally\"], [\"The\", \"Mom\", \"Her\", \"\\\"\", \"So\", \"Daddy\", \"Finally\", \"One\", \"When\", \"But\"], [\" little\", \" family\", \" mom\", \" girl\", \" daddy\", \" dad\", \" people\", \" parents\", \" father\", \" mother\"], [\" girl\", \" boy\", \" one\", \" 3\", \" child\", \" daughter\", \" family\", \" little\", \" three\", \" girls\"], [\" smiled\", \" nodded\", \" was\", \" said\", \" thought\", \" looked\", \" shook\", \"'s\", \" laughed\", \" didn\"], [\" and\", \",\", \".\", \" as\", \" at\", \" a\", \" back\", \" with\", \" excited\", \" again\"], [\" said\", \" nodded\", \" replied\", \" asked\", \" answered\", \" waved\", \" ran\", \" kept\", \" started\", \" continued\"], [\" her\", \".\", \",\", \" in\", \" yes\", \" his\", \" with\", \" their\", \" and\", \" its\"], [\" head\", \" big\", \" heads\", \" eyes\", \" biggest\", \" face\", \" mom\", \" heart\", \" daddy\", \" ears\"], [\".\", \",\", \" yes\", \" excited\", \" in\", \" eagerly\", \" with\", \" to\", \" and\", \" as\"], [\" She\", \"\\n\", \" They\", \" \\\"\", \" Then\", \" The\", \" So\", \" Daddy\", \" As\", \" \"], [\" was\", \" said\", \" knew\", \" had\", \" loved\", \" wanted\", \" felt\", \" liked\", \" looked\", \" thought\"], [\" so\", \" excited\", \" happy\", \" very\", \" ready\", \" a\", \" feeling\", \" having\", \" getting\", \" already\"], [\" to\", \" and\", \" for\", \",\", \".\", \"!\", \" because\", \" that\", \" by\", \" as\"], [\" explore\", \" see\", \" go\", \" learn\", \" be\", \" get\", \" tell\", \" have\", \" start\", \" meet\"], [\" the\", \" what\", \" all\", \" new\", \" her\", \" so\", \" a\", \" how\", \" everything\", \" where\"], [\" the\", \" of\", \" these\", \" around\", \" kinds\", \" about\", \" sorts\", \" her\", \" this\", \" those\"], [\" new\", \" modern\", \" different\", \" planes\", \" things\", \" people\", \" big\", \" fun\", \" amazing\", \" other\"], [\" and\", \".\", \",\", \" taking\", \" flying\", \" in\", \"!\", \" at\", \" take\", \" up\"], [\" so\", \" and\", \" but\", \" the\", \" planes\", \" especially\", \" which\", \" even\", \" she\", \" big\"], [\" she\", \" the\", \" even\", \" all\", \" they\", \" soon\", \" as\", \" when\", \" so\", \" asked\"], [\" was\", \" couldn\", \" wanted\", \" knew\", \" saw\", \" looked\", \" asked\", \" waved\", \" said\", \" started\"], [\" to\", \" on\", \" ready\", \" out\", \" off\", \" a\", \" up\", \" in\", \" into\", \" dressed\"], [\" see\", \" the\", \" go\", \" meet\", \" work\", \" look\", \" explore\", \" sit\", \" take\", \" know\"], [\" airport\", \" top\", \" modern\", \" very\", \" big\", \" plane\", \" beach\", \" front\", \" park\", \" edge\"], [\".\", \" with\", \" too\", \" and\", \" to\", \" first\", \",\", \" just\", \" for\", \" right\"], [\"\\n\", \" She\", \" When\", \" As\", \" The\", \" \", \" They\", \" Everywhere\", \" It\", \" Her\"], [\"\\n\", \"When\", \"<|endoftext|>\", \"The\", \"Summary\", \"At\", \" \", \"Once\", \"As\", \"They\"], [\"When\", \"The\", \"At\", \"As\", \"Once\", \"On\", \"They\", \"Suddenly\", \"Finally\", \"Soon\"], [\" family\", \" little\", \" mom\", \" airport\", \" plane\", \" dad\", \" parents\", \" girl\", \" next\", \" daddy\"], [\" girl\", \" boy\", \" plane\", \" one\", \" family\", \" girls\", \" town\", \" people\", \" planes\", \" airport\"], [\" was\", \" and\", \" saw\", \" looked\", \" watched\", \" asked\", \" had\", \" got\", \" waved\", \"'s\"], [\" up\", \" onto\", \" into\", \" on\", \" aboard\", \" the\", \" out\", \" down\", \" in\", \" off\"], [\" the\", \" onto\", \" on\", \" into\", \" a\", \" to\", \" and\", \" her\", \" high\", \" in\"], [\" stairs\", \" steps\", \" ladder\", \" plane\", \" big\", \" tall\", \" modern\", \" planes\", \" gate\", \" wall\"], [\" and\", \",\", \" to\", \".\", \" of\", \" with\", \" all\", \" as\", \" in\", \" until\"], [\" looked\", \" saw\", \" went\", \" onto\", \" sat\", \" found\", \" started\", \" waved\", \" when\", \" stepped\"], [\",\", \" to\", \" \\\"\", \" goodbye\", \" hello\", \" in\", \":\", \" with\", \",\\\"\", \" excited\"], [\" \\\"\", \" \\ufffd\", \" '\", \" \\\"'\", \"\\n\", \" \", \" ''\", \"\\ufffd\", \" and\", \" \\\"\\\"\"], [\"I\", \"Mom\", \"Wow\", \"Look\", \"This\", \"It\", \"Hello\", \"We\", \"Can\", \"Let\"], [\",\", \" done\", \" I\", \" you\", \" hello\", \" at\", \" then\", \" come\", \" we\", \" what\"], [\",\", \"!\", \"!\\\"\", \".\", \" everyone\", \" for\", \",\\\"\", \" and\", \" you\", \" at\"], [\" mom\", \" you\", \" sweet\", \" Mom\", \" daddy\", \" I\", \" my\", \" everyone\", \" Daddy\", \" little\"], [\"!\", \"!\\\"\", \" is\", \".\", \" was\", \",\", \" can\", \" in\", \" at\", \" cl\"], [\" so\", \" very\", \" going\", \" ready\", \" getting\", \" welcome\", \" safe\", \" here\", \" waiting\", \" happy\"], [\" happy\", \" proud\", \" excited\", \" big\", \" nice\", \" friendly\", \" busy\", \" full\", \" much\", \" brave\"], [\" of\", \"!\\\"\", \"!\", \" and\", \".\", \" to\", \" that\", \".\\\"\", \",\", \"\\\".\"], [\" you\", \" how\", \" me\", \" the\", \" your\", \" yourself\", \" us\", \" what\", \" my\", \" her\"], [\"!\\\"\", \" for\", \".\", \".\\\"\", \"!\", \",\", \"\\\".\", \" all\", \"!\\\".\", \" and\"], [\" You\", \" We\", \" Now\", \" Let\", \" I\", \" Come\", \" Keep\", \" That\", \" The\", \" It\"], [\" you\", \" we\", \" your\", \" the\", \" all\", \" everyone\", \" there\", \" they\", \" these\", \" I\"], [\" ready\", \" all\", \" excited\", \" going\", \" a\", \" sure\", \" getting\", \" the\", \" brave\", \" feeling\"], [\" to\", \" for\", \"?\\\"\", \",\", \"?\\\".\", \" yet\", \" and\", \" now\", \"?\", \" enough\"], [\" go\", \" explore\", \" fly\", \" have\", \" take\", \" start\", \" get\", \" learn\", \" see\", \" climb\"], [\" everyone\", \" off\", \" them\", \" me\", \" it\", \" the\", \" us\", \" your\", \" everybody\", \" her\"], [\" mom\", \" family\", \" new\", \" friends\", \" modern\", \" passport\", \" parents\", \" dad\", \" grandma\", \" daddy\"], [\"?\\\"\", \" the\", \" what\", \" your\", \" and\", \"?\\\".\", \" how\", \",\", \" around\", \" when\"], [\"\\n\", \" \", \" The\", \" And\", \" She\", \" Everyone\", \" Her\", \" Daddy\", \" Mom\", \" They\"], [\"\\n\", \"The\", \" \", \"Summary\", \"<|endoftext|>\", \"Mom\", \"So\", \"\\\"\", \"Everyone\", \"Her\"], [\"The\", \"So\", \"\\\"\", \"Her\", \"Everyone\", \"Mom\", \"And\", \"She\", \"All\", \"At\"], [\" little\", \" family\", \" girl\", \" plane\", \" mom\", \" people\", \" modern\", \" dad\", \" pilot\", \" airport\"], [\" smiled\", \" all\", \" laughed\", \" was\", \" nodded\", \" cheered\", \" said\", \" thanked\", \" looked\", \" replied\"], [\" at\", \" up\", \" around\", \" down\", \" out\", \" surprised\", \" over\", \" back\", \" so\", \" a\"], [\" and\", \" who\", \" so\", \" a\", \" but\", \" then\", \" the\", \" smiled\", \" feeling\", \" proud\"], [\" the\", \" they\", \" sure\", \" then\", \" smiled\", \" after\", \" everyone\", \" soon\", \" nodded\", \" finally\"], [\" all\", \" were\", \" saw\", \" said\", \" nodded\", \" smiled\", \" had\", \" couldn\", \" looked\", \" watched\"], [\" said\", \" smiled\", \" cheered\", \" nodded\", \" agreed\", \" cl\", \" hugged\", \" laughed\", \" shouted\", \" had\"], [\" and\", \" each\", \" the\", \".\", \",\", \" their\", \" her\", \" mom\", \" for\", \" one\"], [\" other\", \"other\", \" one\", \" others\", \" of\", \" family\", \" and\", \" so\", \" plane\", \" mom\"], [\".\", \" tightly\", \" and\", \" tight\", \",\", \" with\", \" again\", \" in\", \" when\", \" before\"], [\" and\", \" happy\", \" then\", \" grateful\", \" thankful\", \" feeling\", \" proud\", \" excited\", \" smiled\", \" \\\"\"], [\" then\", \" the\", \" they\", \" said\", \" thanked\", \" from\", \" smiled\", \" were\", \" everyone\", \" had\"], [\" all\", \" were\", \" said\", \" smiled\", \" knew\", \" laughed\", \" had\", \" thanked\", \" enjoyed\", \" started\"], [\" that\", \" to\", \".\", \" it\", \" with\", \"!\", \" they\", \" the\", \" on\", \",\"], [\"\\n\", \" They\", \" The\", \" From\", \" \", \" Everyone\", \" It\", \" Then\", \" And\", \" Together\"], [\" little\", \" family\", \" girl\", \" modern\", \" mom\", \" plane\", \" airport\", \" dad\", \" daughter\", \" three\"], [\" girl\", \" one\", \" boy\", \" 3\", \" plane\", \" little\", \" family\", \" daughter\", \" three\", \" town\"], [\" was\", \" smiled\", \" felt\", \" had\", \" said\", \"'s\", \" knew\", \" hugged\", \" couldn\", \" ran\"], [\" so\", \" very\", \" the\", \" happy\", \" a\", \" over\", \" thrilled\", \" excited\", \" really\", \" now\"], [\" happy\", \" excited\", \" proud\", \" glad\", \" brave\", \" thrilled\", \" pleased\", \" grateful\", \" eager\", \" thankful\"], [\" that\", \" to\", \" and\", \" she\", \",\", \".\", \"!\", \" with\", \" they\", \" because\"], [\" she\", \" and\", \" they\", \" that\", \" it\", \" because\", \" the\", \" but\", \" with\", \" knowing\"], [\" couldn\", \" had\", \" jumped\", \" was\", \" smiled\", \" ran\", \" felt\", \" hugged\", \" knew\", \" started\"], [\" to\", \" everyone\", \" the\", \" her\", \" them\", \" more\", \" a\", \" all\", \" everything\", \" one\"], [\" explore\", \" go\", \" show\", \" be\", \" stay\", \" keep\", \" make\", \" come\", \" take\", \" share\"], [\" everyone\", \" her\", \" them\", \" the\", \" off\", \" all\", \" that\", \" it\", \" how\", \" their\"], [\" family\", \" mom\", \" appreciation\", \" new\", \" parents\", \" modern\", \" friends\", \" dad\", \" daddy\", \" passport\"], [\" modern\", \" family\", \" toy\", \" friends\", \" friend\", \" passport\", \" plane\", \" adventure\", \" world\", \" things\"], [\" to\", \".\", \" plane\", \" planes\", \" and\", \" at\", \",\", \" car\", \" on\", \" with\"], [\"\\n\", \" She\", \" \", \" The\", \" So\", \" Everyone\", \" From\", \" They\", \" And\", \" Her\"], [\"\\n\", \"Summary\", \"<|endoftext|>\", \"The\", \" \", \"So\", \"She\", \"And\", \"They\", \"From\"], [\"\\n\", \",\", \".\", \"\\n\\n\", \" \\\"\", \".\\\"\", \" The\", \" \", \":\", \" the\"], [\" She\", \" The\", \" Her\", \"\\n\", \" So\", \" Everyone\", \"\\ufffd\", \" It\", \" He\", \" But\"], [\" little\", \" end\", \" moral\", \" End\", \" next\", \" girl\", \" old\", \" sun\", \" man\", \" bear\"], [\" was\", \" took\", \" laughed\", \" went\", \" all\", \" smiled\", \" got\", \" had\", \" watched\", \" arrived\"], [\" to\", \" home\", \" on\", \" outside\", \" back\", \" for\", \" into\", \" out\", \" shopping\", \" inside\"], [\" a\", \" to\", \" the\", \",\", \" their\", \" and\", \" lots\", \" an\", \" many\", \" another\"], [\" the\", \" have\", \" play\", \" explore\", \" find\", \" get\", \" buy\", \" see\", \" enjoy\", \" celebrate\"], [\" beach\", \" store\", \" park\", \" playground\", \" kitchen\", \" big\", \" car\", \" toy\", \" pond\", \" pool\"], [\" store\", \" mall\", \" house\", \" park\", \" playground\", \" restaurant\", \" hospital\", \" shop\", \" market\", \" bathroom\"], [\" and\", \",\", \".\", \" with\", \" to\", \" every\", \" for\", \" had\", \" where\", \" was\"], [\" show\", \" go\", \" see\", \" visit\", \" pick\", \" get\", \" take\", \" meet\", \" celebrate\", \" watch\"], [\" and\", \" back\", \" down\", \" with\", \" on\", \" out\", \" along\", \" up\", \".\", \" together\"], [\" the\", \" their\", \" time\", \",\", \" a\", \" and\", \" board\", \".\", \" big\", \" vacation\"], [\" special\", \" big\", \" fun\", \" trips\", \" new\", \" adventures\", \" trip\", \" amazing\", \" exciting\", \" wonderful\"], [\".\", \",\", \" and\", \" with\", \" too\", \"!\", \" to\", \" every\", \" together\", \" again\"], [\"\\n\", \" They\", \" When\", \" The\", \" Everyone\", \" As\", \" Everywhere\", \" \", \" One\", \" After\"], [\" time\", \" day\", \" morning\", \" week\", \" now\", \" year\", \" person\", \" member\", \" night\", \" step\"], [\" they\", \" the\", \" someone\", \" she\", \",\", \" a\", \" one\", \" he\", \" there\", \" it\"], [\" of\", \" member\", \" day\", \",\", \" little\", \" went\", \" person\", \" was\", \" arrived\", \" item\"], [\" the\", \" them\", \" these\", \" her\", \" those\", \" their\", \" a\", \" it\", \" his\", \" this\"], [\" planes\", \" family\", \" people\", \" luggage\", \" little\", \" flights\", \" new\", \" kids\", \" children\", \" passengers\"], [\" girl\", \" boy\", \" kids\", \" children\", \" girls\", \" people\", \" bird\", \" birds\", \" boys\", \" ones\"], [\"'s\", \",\", \" went\", \" looked\", \" would\", \" was\", \"\\ufffd\", \" and\", \" got\", \" goes\"], [\" the\", \" she\", \" her\", \" they\", \" Lucy\", \" Sarah\", \" Amy\", \" Jane\", \" it\", \" a\"], [\" dad\", \" mom\", \" daddy\", \" parents\", \" family\", \" father\", \" mum\", \" mother\", \" eyes\", \" Dad\"], [\" would\", \",\", \" and\", \" was\", \" went\", \" got\", \" came\", \" smiled\", \" waved\", \" looked\"], [\" would\", \" they\", \" the\", \" she\", \" went\", \" was\", \" her\", \" and\", \" asked\", \" asks\"], [\" always\", \" pick\", \" take\", \" go\", \" watch\", \" wave\", \" come\", \" bring\", \" look\", \" give\"], [\" around\", \" up\", \" and\", \" to\", \" with\", \" together\", \" in\", \" back\", \" on\", \" through\"], [\" her\", \" them\", \" the\", \" joy\", \" excitement\", \" their\", \" a\", \" it\", \" tw\", \" all\"], [\" and\", \".\", \",\", \" to\", \"!\", \" for\", \" in\", \" as\", \" until\", \" on\"], [\" joy\", \" excitement\", \" love\", \" anticipation\", \" cheer\", \" help\", \" make\", \" enthusiasm\", \" smile\", \" have\"], [\" for\", \".\", \" to\", \" that\", \",\", \" she\", \" of\", \"!\", \" they\", \" wherever\"], [\" the\", \" their\", \" her\", \" more\", \" new\", \" something\", \" a\", \" all\", \" many\", \" everyone\"], [\" best\", \" next\", \" new\", \" family\", \" future\", \" fun\", \" surprise\", \" special\", \" airport\", \" rest\"], [\".\", \" to\", \" soon\", \" would\", \" and\", \" of\", \"'s\", \" again\", \" one\", \" on\"], [\"\\n\", \" \", \" She\", \" The\", \" They\", \" And\", \" It\", \" Finally\", \" Her\", \" Everyone\"], [\"Summary\", \"\\n\", \"Story\", \"Words\", \"The\", \"Features\", \"Finally\", \" \", \"One\", \"Eventually\"], [\"The\", \"Finally\", \"One\", \"At\", \"Eventually\", \"When\", \"Soon\", \"And\", \"After\", \"They\"], [\",\", \" the\", \" they\", \" she\", \" her\", \" it\", \" one\", \" their\", \" enough\", \" after\"], [\" the\", \" they\", \" it\", \" their\", \" her\", \" after\", \" she\", \" one\", \" when\", \" a\"], [\" little\", \" family\", \" plane\", \" airport\", \" day\", \" march\", \" girl\", \" dad\", \" whole\", \" mom\"], [\" girl\", \" boy\", \" one\", \" plane\", \" bird\", \" brother\", \" duck\", \" march\", \" ant\", \" family\"], [\" had\", \"'s\", \" got\", \" grew\", \" and\", \" was\", \" found\", \" said\", \" became\", \" started\"], [\" to\", \" her\", \" marching\", \" their\", \" the\", \" singing\", \" his\", \" walking\", \" exploring\", \" looking\"], [\" feel\", \" understand\", \" get\", \" march\", \" miss\", \" grow\", \" realize\", \" sing\", \" enjoy\", \" learn\"], [\" tired\", \" ready\", \" a\", \" excited\", \" bored\", \" off\", \" sleepy\", \" hungry\", \" better\", \" very\"], [\".\", \" and\", \",\", \" from\", \" so\", \" again\", \" too\", \"!\", \" in\", \" at\"], [\".\", \",\", \" and\", \"!\", \" so\", \" -\", \" in\", \" but\", \" because\", \" when\"], [\" She\", \" Her\", \" But\", \" The\", \" They\", \" This\", \" As\", \" So\", \"\\n\", \" When\"], [\" asked\", \" was\", \" looked\", \" wanted\", \" knew\", \" said\", \" saw\", \" had\", \" grabbed\", \" went\"], [\" so\", \" filled\", \" very\", \" getting\", \" a\", \" sad\", \" determined\", \" happy\", \" looking\", \" feeling\"], [\" excited\", \" thirsty\", \" happy\", \" proud\", \" tired\", \" eager\", \" full\", \" glad\", \" close\", \" grateful\"]], \"correctTokenRank\": [294, 0, 0, 0, 0, 0, 0, 0, 73, 0, 0, 0, 13, 0, 28, 95, 0, 1, 22, 2, 0, 1, 0, 0, 24, 0, 2, 0, 1, 38, 0, 0, 0, 0, 1, 11, 0, 0, 0, 9, 4, 5, 0, 0, 2, 2, 1, 2, 0, 15, 0, 0, 20, 0, 15, 1, 0, 13, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 2, 0, 3, 2, 1, 0, 16, 0, 1, 0, 0, 0, 0, 1, 1, 0, 42, 0, 0, 0, 0, 16, 0, 0, 49, 1, 0, 7, 2, 0, 1, 0, 0, 4, 23, 0, 0, 0, 13, 7, 1, 0, 0, 0, 0, 1, 8, 42, 0, 1, 0, 6, 1, 0, 4, 0, 2, 13, 5, 2, 0, 0, 0, 0, 0, 4, 0, 10, 0, 2, 1, 3, 2, 1, 0, 2, 2, 20, 10, 3, 2, 1, 0, 97, 39, 4, 23, 4, 73, 3, 0, 19, 0, 6, 0, 0, 4, 0, 1, 2, 4, 1, 0, 152, 4, 25, 0, 42, 0, 0, 3, 0, 0, 1, 4, 0, 0, 0, 0, 12, 0, 2, 79, 5, 0, 0, 1, 0, 4], \"correctTokenLogProb\": [-11.625454902648926, -0.372489869594574, -0.005904730875045061, -0.011737688444554806, -0.1878184974193573, -0.10250605642795563, -0.1382717341184616, -0.0321299247443676, -5.991303443908691, -0.5746181607246399, -0.4282326400279999, -0.30328744649887085, -5.65919828414917, -2.4236323833465576, -4.623875617980957, -6.590052604675293, -0.042138077318668365, -1.9959081411361694, -5.6335530281066895, -1.905128002166748, -0.9895710945129395, -2.5058727264404297, -0.35668182373046875, -0.7548707127571106, -4.735068321228027, -0.5946731567382812, -2.1486759185791016, -0.09278443455696106, -1.3276169300079346, -7.621414661407471, -0.06834664195775986, -0.889788031578064, -0.0339607410132885, -0.9458173513412476, -1.6791880130767822, -3.8572938442230225, -0.6308883428573608, -1.1232950687408447, -0.29644775390625, -5.13623046875, -5.095574378967285, -3.26711368560791, -0.007819517515599728, -0.18164236843585968, -2.637809991836548, -2.3811194896698, -1.3665785789489746, -1.8806651830673218, -0.003452533157542348, -5.287355422973633, -0.43409377336502075, -0.19858379662036896, -4.841526985168457, -0.5131959915161133, -4.908519744873047, -1.3566906452178955, -0.4600834846496582, -4.830392837524414, -2.2664647102355957, -0.24191200733184814, -0.050112128257751465, -0.2582389712333679, -0.2726271152496338, -0.0036967527121305466, -1.3802396059036255, -0.0841309204697609, -2.2264645099639893, -0.8177525997161865, -0.0016505202511325479, -0.2397218495607376, -0.6482017636299133, -0.9068384170532227, -2.4962146282196045, -0.2099856436252594, -1.7463961839675903, -2.0468666553497314, -0.11845239996910095, -2.431112766265869, -1.9535925388336182, -1.379508137702942, -0.8167728185653687, -4.419667720794678, -0.7300421595573425, -2.6654233932495117, -0.17896340787410736, -0.18972711265087128, -0.7267122864723206, -0.020010385662317276, -1.7059850692749023, -1.5423017740249634, -0.00046659549116156995, -6.36149787902832, -1.2832740545272827, -1.2193286418914795, -0.8976349234580994, -0.5017741918563843, -4.201582431793213, -0.109302818775177, -0.03258315473794937, -6.803769588470459, -1.8130090236663818, -0.719337522983551, -3.136812686920166, -2.519213914871216, -2.1019766330718994, -1.9251044988632202, -0.09708235412836075, -0.3409811854362488, -2.387692928314209, -5.753907203674316, -0.04679347202181816, -0.31035345792770386, -0.26277634501457214, -4.80776834487915, -3.099146604537964, -1.2577797174453735, -0.23072941601276398, -0.1686599999666214, -0.03456159308552742, -0.041516028344631195, -2.4174718856811523, -4.463851451873779, -7.481378078460693, -0.6082432866096497, -2.094578504562378, -0.6794657707214355, -3.3744683265686035, -1.3131829500198364, -0.0037339269183576107, -2.4467132091522217, -1.1844749450683594, -2.7604708671569824, -5.0592942237854, -4.280617713928223, -1.8505373001098633, -0.14208939671516418, -0.0002898749662563205, -0.5908623933792114, -0.17303341627120972, -0.4503314197063446, -2.1354479789733887, -0.47746407985687256, -3.7926478385925293, -0.02097589708864689, -2.1107640266418457, -1.169552206993103, -2.8503684997558594, -2.9834933280944824, -1.5757774114608765, -0.5248868465423584, -3.215791702270508, -16.646602630615234, -5.324350357055664, -4.359821796417236, -3.124971628189087, -2.7403366565704346, -1.8093516826629639, -0.9515280723571777, -6.534619331359863, -5.37264347076416, -3.521487236022949, -4.890478134155273, -2.865650177001953, -8.220545768737793, -2.573607921600342, -0.8533508777618408, -5.152226448059082, -0.3188457787036896, -4.101963520050049, -0.7935125827789307, -0.32164621353149414, -3.5270328521728516, -0.5608841180801392, -2.829089641571045, -2.155390977859497, -3.445237159729004, -1.54229736328125, -0.7731188535690308, -7.224249362945557, -2.695648193359375, -6.626819610595703, -0.9364675283432007, -5.718138694763184, -0.8462379574775696, -1.065568208694458, -3.258896589279175, -1.3465521335601807, -0.20639769732952118, -1.9113335609436035, -3.625382423400879, -0.052747633308172226, -0.20388591289520264, -1.143463134765625, -0.013952808454632759, -4.204424858093262, -0.14188191294670105, -2.596562623977661, -7.299529075622559, -4.704736232757568, -0.23656286299228668, -0.42535340785980225, -2.492424964904785, -1.115234136581421, -2.6541266441345215]}\n",
+              "    )\n",
+              "    </script>"
+            ],
+            "text/plain": [
+              "<circuitsvis.utils.render.RenderedHTML at 0x77ec1eb34310>"
+            ]
+          },
+          "execution_count": 8,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "example_prompt = model.generate(\n",
         "    \"Once upon a time\",\n",
@@ -295,11 +626,146 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 9,
       "metadata": {
         "id": "oAsZCAdJOVHw"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Run name: 16384-L1-5-LR-5e-05-Tokens-1.229e+08\n",
+            "n_tokens_per_buffer (millions): 0.262144\n",
+            "Lower bound: n_contexts_per_buffer (millions): 0.001024\n",
+            "Total training steps: 30000\n",
+            "Total wandb updates: 1000\n",
+            "n_tokens_per_feature_sampling_window (millions): 1048.576\n",
+            "n_tokens_per_dead_feature_window (millions): 1048.576\n",
+            "We will reset the sparsity calculation 30 times.\n",
+            "Number tokens in sparsity calculation window: 4.10e+06\n",
+            "Loaded pretrained model tiny-stories-1L-21M into HookedTransformer\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Downloading readme: 100%|██████████| 415/415 [00:00<00:00, 4.30MB/s]\n",
+            "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcurt-tigges\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "Tracking run with wandb version 0.17.1"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Run data is saved locally in <code>/home/curttigges/projects/SAELens/tutorials/wandb/run-20240610_114538-yr1gvjdc</code>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Syncing run <strong><a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc' target=\"_blank\">16384-L1-5-LR-5e-05-Tokens-1.229e+08</a></strong> to <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View project at <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial</a>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View run at <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc</a>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Estimating norm scaling factor: 100%|██████████| 1000/1000 [00:33<00:00, 29.79it/s]\n",
+            "30000| MSE Loss 187.703 | L1 156.883:   1%|          | 1228800/122880000 [42:59<70:56:28, 476.34it/s]\n"
+          ]
+        },
+        {
+          "data": {
+            "text/html": [
+              "<style>\n",
+              "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
+              "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
+              "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
+              "    </style>\n",
+              "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>details/current_l1_coefficient</td><td>▁▅██████████████████████████████████████</td></tr><tr><td>details/current_learning_rate</td><td>████████████████████████████████▇▇▅▅▄▃▂▁</td></tr><tr><td>details/n_training_tokens</td><td>▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███</td></tr><tr><td>losses/ghost_grad_loss</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/l1_loss</td><td>█▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/mse_loss</td><td>▄▅█▇▆▆▅▄▄▄▄▃▃▃▃▃▃▃▂▃▂▂▂▂▂▂▂▂▁▁▂▁▁▁▁▁▁▁▁▁</td></tr><tr><td>losses/overall_loss</td><td>▁▅█▇▇▆▆▆▆▆▅▅▅▅▅▅▅▅▅▅▅▄▅▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄</td></tr><tr><td>metrics/CE_loss_score</td><td>█▂▁▂▃▃▄▄▄▄▄▄▄▄▅▅▅▅▅▅▅▅▅▆▅▆▅▅▅▅▆▅▆▆▆▆▆▆▆▅</td></tr><tr><td>metrics/ce_loss_with_ablation</td><td>█▆▇▃▄▄▃▆▆▅▃▅▅▅▅▃▄▅▅▆▄▅▃▅▃▄▅▁▅▅▆▅▆▄▆▄▆▄▅▄</td></tr><tr><td>metrics/ce_loss_with_sae</td><td>▂▇█▆▅▅▅▅▄▄▅▆▄▄▃▄▄▄▁▃▃▃▁▄▄▂▃▂▄▂▁▃▄▃▄▂▄▂▃▄</td></tr><tr><td>metrics/ce_loss_without_sae</td><td>▇▇▇▅▄▄▆▆▅▄▇█▆▆▄▅▆▇▁▄▄▅▁▇▆▄▅▄▆▄▂▅▇▆▇▄▇▄▅█</td></tr><tr><td>metrics/explained_variance</td><td>▅▄▁▂▃▄▄▅▅▅▅▆▆▆▆▆▆▆▆▇▆▇▇▇▇▇▇▇▇█▇▇▇███████</td></tr><tr><td>metrics/explained_variance_std</td><td>▁▅███▇▇▇▇▇▇▆▆▇▇▆▆▆▆▆▆▆▆▆▆▆▆▆▆▆▇▆▆▆▆▆▆▆▅▆</td></tr><tr><td>metrics/l0</td><td>█▃▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>metrics/l2_norm</td><td>▆▅▁▃▃▂▂▃▁▅▃▃▅▆▂█▄▃▄▆▇▃▆▃▃▁▂▅▅▆▇▅▄▅▃██▃▃▃</td></tr><tr><td>metrics/l2_norm_in</td><td>▅▆▅▆▆▄▅▅▃▆▄▄▆▇▄█▄▃▅▆▇▃▆▃▃▁▂▅▅▆▇▄▅▆▃▇█▄▄▃</td></tr><tr><td>metrics/l2_ratio</td><td>▇▅▁▃▃▃▃▄▃▆▅▅▆▆▄▇▆▆▅▆▇▅▇▆▅▅▅▆▆▇▇▇▆▆▆██▆▆▆</td></tr><tr><td>metrics/mean_log10_feature_sparsity</td><td>█▄▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>sparsity/below_1e-5</td><td>▁▁▁▅▆███▆▆▆▆▆▆▆▆▅▅▅▅▅▅▅▅▅▅▅▅▅▅</td></tr><tr><td>sparsity/below_1e-6</td><td>▁▁▁▁▁▃▆▆▆█▆█████▆▆▆▆▆▆▆▆▆▆▆▆▆▆</td></tr><tr><td>sparsity/dead_features</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▅▅▅▁▁▅█▁▁▁▅▅█▅█▅▅████▅▅▅████</td></tr><tr><td>sparsity/mean_passes_since_fired</td><td>▂▁▁▁▁▁▁▁▂▂▁▁▂▂▂▂▂▂▄▂▂▂▃▃▄▄▄▃▃▃▄▅▆▄▄▅▆▇▇█</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>details/current_l1_coefficient</td><td>5</td></tr><tr><td>details/current_learning_rate</td><td>0.0</td></tr><tr><td>details/n_training_tokens</td><td>122880000</td></tr><tr><td>losses/ghost_grad_loss</td><td>0.0</td></tr><tr><td>losses/l1_loss</td><td>31.37665</td></tr><tr><td>losses/mse_loss</td><td>187.70346</td></tr><tr><td>losses/overall_loss</td><td>344.5867</td></tr><tr><td>metrics/CE_loss_score</td><td>0.90369</td></tr><tr><td>metrics/ce_loss_with_ablation</td><td>8.30545</td></tr><tr><td>metrics/ce_loss_with_sae</td><td>2.62867</td></tr><tr><td>metrics/ce_loss_without_sae</td><td>2.02306</td></tr><tr><td>metrics/explained_variance</td><td>0.66377</td></tr><tr><td>metrics/explained_variance_std</td><td>0.13242</td></tr><tr><td>metrics/l0</td><td>192.95703</td></tr><tr><td>metrics/l2_norm</td><td>24.64933</td></tr><tr><td>metrics/l2_norm_in</td><td>31.38967</td></tr><tr><td>metrics/l2_ratio</td><td>0.77361</td></tr><tr><td>metrics/mean_log10_feature_sparsity</td><td>-2.66941</td></tr><tr><td>sparsity/below_1e-5</td><td>2</td></tr><tr><td>sparsity/below_1e-6</td><td>2</td></tr><tr><td>sparsity/dead_features</td><td>2</td></tr><tr><td>sparsity/mean_passes_since_fired</td><td>0.86823</td></tr></table><br/></div></div>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              " View run <strong style=\"color:#cdcd00\">16384-L1-5-LR-5e-05-Tokens-1.229e+08</strong> at: <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial/runs/yr1gvjdc</a><br/> View project at: <a href='https://wandb.ai/curt-tigges/sae_lens_tutorial' target=\"_blank\">https://wandb.ai/curt-tigges/sae_lens_tutorial</a><br/>Synced 6 W&B file(s), 0 media file(s), 3 artifact file(s) and 0 other file(s)"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        },
+        {
+          "data": {
+            "text/html": [
+              "Find logs at: <code>./wandb/run-20240610_114538-yr1gvjdc/logs</code>"
+            ],
+            "text/plain": [
+              "<IPython.core.display.HTML object>"
+            ]
+          },
+          "metadata": {},
+          "output_type": "display_data"
+        }
+      ],
       "source": [
         "total_training_steps = 30_000  # probably we should do more\n",
         "batch_size = 4096\n",
@@ -425,7 +891,7 @@
       "name": "python",
       "nbconvert_exporter": "python",
       "pygments_lexer": "ipython3",
-      "version": "3.11.7"
+      "version": "3.10.14"
     }
   },
   "nbformat": 4,