sae-lens 5.9.1__py3-none-any.whl → 6.0.0rc1__py3-none-any.whl

sae_lens/saes/standard_sae.py

@@ -0,0 +1,167 @@
+import numpy as np
+import torch
+from jaxtyping import Float
+from numpy.typing import NDArray
+from torch import nn
+
+from sae_lens.saes.sae import SAE, SAEConfig, TrainingSAE, TrainStepInput
+
+
+class StandardSAE(SAE):
+    """
+    StandardSAE is an inference-only implementation of a Sparse Autoencoder (SAE)
+    using a simple linear encoder and decoder.
+
+    It implements the required abstract methods from BaseSAE:
+      - initialize_weights: sets up simple parameter initializations for W_enc, b_enc, W_dec, and b_dec.
+      - encode: computes the feature activations from an input.
+      - decode: reconstructs the input from the feature activations.
+
+    The BaseSAE.forward() method automatically calls encode and decode,
+    including any error-term processing if configured.
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(self, cfg: SAEConfig, use_error_term: bool = False):
+        super().__init__(cfg, use_error_term)
+
+    def initialize_weights(self) -> None:
+        # Initialize encoder weights and bias.
+        self.b_enc = nn.Parameter(
+            torch.zeros(self.cfg.d_sae, dtype=self.dtype, device=self.device)
+        )
+        self.b_dec = nn.Parameter(
+            torch.zeros(self.cfg.d_in, dtype=self.dtype, device=self.device)
+        )
+
+        # Use Kaiming Uniform for W_enc
+        w_enc_data = torch.empty(
+            self.cfg.d_in, self.cfg.d_sae, dtype=self.dtype, device=self.device
+        )
+        nn.init.kaiming_uniform_(w_enc_data)
+        self.W_enc = nn.Parameter(w_enc_data)
+
+        # Use Kaiming Uniform for W_dec
+        w_dec_data = torch.empty(
+            self.cfg.d_sae, self.cfg.d_in, dtype=self.dtype, device=self.device
+        )
+        nn.init.kaiming_uniform_(w_dec_data)
+        self.W_dec = nn.Parameter(w_dec_data)
+
+    def encode(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        """
+        Encode the input tensor into the feature space.
+        For inference, no noise is added.
+        """
+        # Preprocess the SAE input (casting type, applying hooks, normalization)
+        sae_in = self.process_sae_in(x)
+        # Compute the pre-activation values
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)
+        # Apply the activation function (e.g., ReLU, tanh-relu, depending on config)
+        return self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+
+    def decode(
+        self, feature_acts: Float[torch.Tensor, "... d_sae"]
+    ) -> Float[torch.Tensor, "... d_in"]:
+        """
+        Decode the feature activations back to the input space.
+        Now, if hook_z reshaping is turned on, we reverse the flattening.
+        """
+        # 1) apply finetuning scaling if configured.
+        scaled_features = self.apply_finetuning_scaling_factor(feature_acts)
+        # 2) linear transform
+        sae_out_pre = scaled_features @ self.W_dec + self.b_dec
+        # 3) hook reconstruction
+        sae_out_pre = self.hook_sae_recons(sae_out_pre)
+        # 4) optional out-normalization (e.g. constant_norm_rescale or layer_norm)
+        sae_out_pre = self.run_time_activation_norm_fn_out(sae_out_pre)
+        # 5) if hook_z is enabled, rearrange back to (..., n_heads, d_head).
+        return self.reshape_fn_out(sae_out_pre, self.d_head)
+
+
+class StandardTrainingSAE(TrainingSAE):
+    """
+    StandardTrainingSAE is a concrete implementation of BaseTrainingSAE using the "standard" SAE architecture.
+    It implements:
+      - initialize_weights: basic weight initialization for encoder/decoder.
+      - encode: inference encoding (invokes encode_with_hidden_pre).
+      - decode: a simple linear decoder.
+      - encode_with_hidden_pre: computes pre-activations, adds noise when training, and then activates.
+      - calculate_aux_loss: computes a sparsity penalty based on the (optionally scaled) p-norm of feature activations.
+    """
+
+    b_enc: nn.Parameter
+
+    def initialize_weights(self) -> None:
+        # Basic init
+        # In Python MRO, this calls StandardSAE.initialize_weights()
+        StandardSAE.initialize_weights(self)  # type: ignore
+
+        # Complex init logic from original TrainingSAE
+        if self.cfg.decoder_orthogonal_init:
+            self.W_dec.data = nn.init.orthogonal_(self.W_dec.data.T).T
+
+        elif self.cfg.decoder_heuristic_init:
+            self.W_dec.data = torch.rand(  # Changed from Parameter to data assignment
+                self.cfg.d_sae, self.cfg.d_in, dtype=self.dtype, device=self.device
+            )
+            self.initialize_decoder_norm_constant_norm()
+
+        if self.cfg.init_encoder_as_decoder_transpose:
+            self.W_enc.data = self.W_dec.data.T.clone().contiguous()  # type: ignore
+
+        if self.cfg.normalize_sae_decoder:
+            with torch.no_grad():
+                self.set_decoder_norm_to_unit_norm()
+
+    @torch.no_grad()
+    def initialize_decoder_norm_constant_norm(self, norm: float = 0.1):
+        self.W_dec.data /= torch.norm(self.W_dec.data, dim=1, keepdim=True)  # type: ignore
+        self.W_dec.data *= norm
+
+    def encode_with_hidden_pre(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+        # Process the input (including dtype conversion, hook call, and any activation normalization)
+        sae_in = self.process_sae_in(x)
+        # Compute the pre-activation (and allow for a hook if desired)
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)  # type: ignore
+        # Add noise during training for robustness (scaled by noise_scale from the configuration)
+        if self.training and self.cfg.noise_scale > 0:
+            hidden_pre_noised = (
+                hidden_pre + torch.randn_like(hidden_pre) * self.cfg.noise_scale
+            )
+        else:
+            hidden_pre_noised = hidden_pre
+        # Apply the activation function (and any post-activation hook)
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre_noised))
+        return feature_acts, hidden_pre_noised
+
+    def calculate_aux_loss(
+        self,
+        step_input: TrainStepInput,
+        feature_acts: torch.Tensor,
+        hidden_pre: torch.Tensor,
+        sae_out: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        # The "standard" auxiliary loss is a sparsity penalty on the feature activations
+        weighted_feature_acts = feature_acts
+        if self.cfg.scale_sparsity_penalty_by_decoder_norm:
+            weighted_feature_acts = feature_acts * self.W_dec.norm(dim=1)
+
+        # Compute the p-norm (set by cfg.lp_norm) over the feature dimension
+        sparsity = weighted_feature_acts.norm(p=self.cfg.lp_norm, dim=-1)
+        l1_loss = (step_input.current_l1_coefficient * sparsity).mean()
+
+        return {"l1_loss": l1_loss}
+
+    def log_histograms(self) -> dict[str, NDArray[np.generic]]:
+        """Log histograms of the weights and biases."""
+        b_e_dist = self.b_enc.detach().float().cpu().numpy()
+        return {
+            **super().log_histograms(),
+            "weights/b_e": b_e_dist,
+        }

sae_lens/saes/topk_sae.py

@@ -0,0 +1,305 @@
+"""Inference-only TopKSAE variant, similar in spirit to StandardSAE but using a TopK-based activation."""
+
+from typing import Callable
+
+import torch
+from jaxtyping import Float
+from torch import nn
+
+from sae_lens.saes.sae import (
+    SAE,
+    SAEConfig,
+    TrainingSAE,
+    TrainingSAEConfig,
+    TrainStepInput,
+)
+
+
+class TopK(nn.Module):
+    """
+    A simple TopK activation that zeroes out all but the top K elements along the last dimension,
+    then optionally applies a post-activation function (e.g., ReLU).
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(
+        self,
+        k: int,
+        postact_fn: Callable[[torch.Tensor], torch.Tensor] = nn.ReLU(),
+    ):
+        super().__init__()
+        self.k = k
+        self.postact_fn = postact_fn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        1) Select top K elements along the last dimension.
+        2) Apply post-activation (often ReLU).
+        3) Zero out all other entries.
+        """
+        topk = torch.topk(x, k=self.k, dim=-1)
+        values = self.postact_fn(topk.values)
+        result = torch.zeros_like(x)
+        result.scatter_(-1, topk.indices, values)
+        return result
+
+
+class TopKSAE(SAE):
+    """
+    An inference-only sparse autoencoder using a "topk" activation function.
+    It uses linear encoder and decoder layers, applying the TopK activation
+    to the hidden pre-activation in its encode step.
+    """
+
+    def __init__(self, cfg: SAEConfig, use_error_term: bool = False):
+        """
+        Args:
+            cfg: SAEConfig defining model size and behavior.
+            use_error_term: Whether to apply the error-term approach in the forward pass.
+        """
+        super().__init__(cfg, use_error_term)
+
+        if self.cfg.activation_fn != "topk":
+            raise ValueError("TopKSAE must use a TopK activation function.")
+
+    def initialize_weights(self) -> None:
+        """
+        Initializes weights and biases for encoder/decoder similarly to the standard SAE,
+        that is:
+          - b_enc, b_dec are zero-initialized
+          - W_enc, W_dec are Kaiming Uniform
+        """
+        # encoder bias
+        self.b_enc = nn.Parameter(
+            torch.zeros(self.cfg.d_sae, dtype=self.dtype, device=self.device)
+        )
+        # decoder bias
+        self.b_dec = nn.Parameter(
+            torch.zeros(self.cfg.d_in, dtype=self.dtype, device=self.device)
+        )
+
+        # encoder weight
+        w_enc_data = torch.empty(
+            self.cfg.d_in, self.cfg.d_sae, dtype=self.dtype, device=self.device
+        )
+        nn.init.kaiming_uniform_(w_enc_data)
+        self.W_enc = nn.Parameter(w_enc_data)
+
+        # decoder weight
+        w_dec_data = torch.empty(
+            self.cfg.d_sae, self.cfg.d_in, dtype=self.dtype, device=self.device
+        )
+        nn.init.kaiming_uniform_(w_dec_data)
+        self.W_dec = nn.Parameter(w_dec_data)
+
+    def encode(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        """
+        Converts input x into feature activations.
+        Uses topk activation from the config (cfg.activation_fn == "topk")
+        under the hood.
+        """
+        sae_in = self.process_sae_in(x)
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)
+        # The BaseSAE already sets self.activation_fn to TopK(...) if config requests topk.
+        return self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+
+    def decode(
+        self, feature_acts: Float[torch.Tensor, "... d_sae"]
+    ) -> Float[torch.Tensor, "... d_in"]:
+        """
+        Reconstructs the input from topk feature activations.
+        Applies optional finetuning scaling, hooking to recons, out normalization,
+        and optional head reshaping.
+        """
+        scaled_features = self.apply_finetuning_scaling_factor(feature_acts)
+        sae_out_pre = scaled_features @ self.W_dec + self.b_dec
+        sae_out_pre = self.hook_sae_recons(sae_out_pre)
+        sae_out_pre = self.run_time_activation_norm_fn_out(sae_out_pre)
+        return self.reshape_fn_out(sae_out_pre, self.d_head)
+
+    def _get_activation_fn(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        if self.cfg.activation_fn == "topk":
+            if "k" not in self.cfg.activation_fn_kwargs:
+                raise ValueError("TopK activation function requires a k value.")
+            k = self.cfg.activation_fn_kwargs.get(
+                "k", 1
+            )  # Default k to 1 if not provided
+            postact_fn = self.cfg.activation_fn_kwargs.get(
+                "postact_fn", nn.ReLU()
+            )  # Default post-activation to ReLU if not provided
+            return TopK(k, postact_fn)
+        # Otherwise, return the "standard" handling from BaseSAE
+        return super()._get_activation_fn()
+
+
+class TopKTrainingSAE(TrainingSAE):
+    """
+    TopK variant with training functionality. Injects noise during training, optionally
+    calculates a topk-related auxiliary loss, etc.
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(self, cfg: TrainingSAEConfig, use_error_term: bool = False):
+        super().__init__(cfg, use_error_term)
+
+        if self.cfg.activation_fn != "topk":
+            raise ValueError("TopKSAE must use a TopK activation function.")
+
+    def initialize_weights(self) -> None:
+        """Very similar to TopKSAE, using zero biases + Kaiming Uniform weights."""
+        self.b_enc = nn.Parameter(
+            torch.zeros(self.cfg.d_sae, dtype=self.dtype, device=self.device)
+        )
+        self.b_dec = nn.Parameter(
+            torch.zeros(self.cfg.d_in, dtype=self.dtype, device=self.device)
+        )
+
+        w_enc_data = torch.empty(
+            self.cfg.d_in, self.cfg.d_sae, dtype=self.dtype, device=self.device
+        )
+        nn.init.kaiming_uniform_(w_enc_data)
+        self.W_enc = nn.Parameter(w_enc_data)
+
+        w_dec_data = torch.empty(
+            self.cfg.d_sae, self.cfg.d_in, dtype=self.dtype, device=self.device
+        )
+        nn.init.kaiming_uniform_(w_dec_data)
+        self.W_dec = nn.Parameter(w_dec_data)
+
+    def encode_with_hidden_pre(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+        """
+        Similar to the base training method: cast input, optionally add noise, then apply TopK.
+        """
+        sae_in = self.process_sae_in(x)
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)
+
+        # Inject noise if training
+        if self.training and self.cfg.noise_scale > 0:
+            hidden_pre_noised = (
+                hidden_pre + torch.randn_like(hidden_pre) * self.cfg.noise_scale
+            )
+        else:
+            hidden_pre_noised = hidden_pre
+
+        # Apply the TopK activation function (already set in self.activation_fn if config is "topk")
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre_noised))
+        return feature_acts, hidden_pre_noised
+
+    def calculate_aux_loss(
+        self,
+        step_input: TrainStepInput,
+        feature_acts: torch.Tensor,
+        hidden_pre: torch.Tensor,
+        sae_out: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        # Calculate the auxiliary loss for dead neurons
+        topk_loss = self.calculate_topk_aux_loss(
+            sae_in=step_input.sae_in,
+            sae_out=sae_out,
+            hidden_pre=hidden_pre,
+            dead_neuron_mask=step_input.dead_neuron_mask,
+        )
+        return {"auxiliary_reconstruction_loss": topk_loss}
+
+    def _get_activation_fn(self):
+        if self.cfg.activation_fn == "topk":
+            if "k" not in self.cfg.activation_fn_kwargs:
+                raise ValueError("TopK activation function requires a k value.")
+            k = self.cfg.activation_fn_kwargs.get("k", 1)
+            postact_fn = self.cfg.activation_fn_kwargs.get("postact_fn", nn.ReLU())
+            return TopK(k, postact_fn)
+        return super()._get_activation_fn()
+
+    def calculate_topk_aux_loss(
+        self,
+        sae_in: torch.Tensor,
+        sae_out: torch.Tensor,
+        hidden_pre: torch.Tensor,
+        dead_neuron_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """
+        Calculate TopK auxiliary loss.
+
+        This auxiliary loss encourages dead neurons to learn useful features by having
+        them reconstruct the residual error from the live neurons. It's a key part of
+        preventing neuron death in TopK SAEs.
+        """
+        # Mostly taken from https://github.com/EleutherAI/sae/blob/main/sae/sae.py, except without variance normalization
+        # NOTE: checking the number of dead neurons will force a GPU sync, so performance can likely be improved here
+        if dead_neuron_mask is None or (num_dead := int(dead_neuron_mask.sum())) == 0:
+            return sae_out.new_tensor(0.0)
+        residual = (sae_in - sae_out).detach()
+
+        # Heuristic from Appendix B.1 in the paper
+        k_aux = sae_in.shape[-1] // 2
+
+        # Reduce the scale of the loss if there are a small number of dead latents
+        scale = min(num_dead / k_aux, 1.0)
+        k_aux = min(k_aux, num_dead)
+
+        auxk_acts = _calculate_topk_aux_acts(
+            k_aux=k_aux,
+            hidden_pre=hidden_pre,
+            dead_neuron_mask=dead_neuron_mask,
+        )
+
+        # Encourage the top ~50% of dead latents to predict the residual of the
+        # top k living latents
+        recons = self.decode(auxk_acts)
+        auxk_loss = (recons - residual).pow(2).sum(dim=-1).mean()
+        return scale * auxk_loss
+
+    def _calculate_topk_aux_acts(
+        self,
+        k_aux: int,
+        hidden_pre: torch.Tensor,
+        dead_neuron_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Helper method to calculate activations for the auxiliary loss.
+
+        Args:
+            k_aux: Number of top dead neurons to select
+            hidden_pre: Pre-activation values from encoder
+            dead_neuron_mask: Boolean mask indicating which neurons are dead
+
+        Returns:
+            Tensor with activations for only the top-k dead neurons, zeros elsewhere
+        """
+        # Don't include living latents in this loss (set them to -inf so they won't be selected)
+        auxk_latents = torch.where(
+            dead_neuron_mask[None],
+            hidden_pre,
+            torch.tensor(-float("inf"), device=hidden_pre.device),
+        )
+
+        # Find topk values among dead neurons
+        auxk_topk = auxk_latents.topk(k_aux, dim=-1, sorted=False)
+
+        # Create a tensor of zeros, then place the topk values at their proper indices
+        auxk_acts = torch.zeros_like(hidden_pre)
+        auxk_acts.scatter_(-1, auxk_topk.indices, auxk_topk.values)
+
+        return auxk_acts
+
+
+def _calculate_topk_aux_acts(
+    k_aux: int,
+    hidden_pre: torch.Tensor,
+    dead_neuron_mask: torch.Tensor,
+) -> torch.Tensor:
+    # Don't include living latents in this loss
+    auxk_latents = torch.where(dead_neuron_mask[None], hidden_pre, -torch.inf)
+    # Top-k dead latents
+    auxk_topk = auxk_latents.topk(k_aux, sorted=False)
+    # Set the activations to zero for all but the top k_aux dead latents
+    auxk_acts = torch.zeros_like(hidden_pre)
+    auxk_acts.scatter_(-1, auxk_topk.indices, auxk_topk.values)
+    # Set activations to zero for all but top k_aux dead latents
+    return auxk_acts

sae_lens/training/activations_store.py

@@ -28,7 +28,7 @@ from sae_lens.config import (
     HfDataset,
     LanguageModelSAERunnerConfig,
 )
-from sae_lens.sae import SAE
+from sae_lens.saes.sae import SAE
 from sae_lens.tokenization_and_batching import concat_and_batch_sequences
 
 
@@ -177,7 +177,7 @@ class ActivationsStore:
             dataset_trust_remote_code=sae.cfg.dataset_trust_remote_code,
             dtype=sae.cfg.dtype,
             device=torch.device(device),
-            seqpos_slice=sae.cfg.seqpos_slice,
+            seqpos_slice=sae.cfg.seqpos_slice or (None,),
         )
 
     def __init__(

sae_lens/training/sae_trainer.py

@@ -11,9 +11,9 @@ from transformer_lens.hook_points import HookedRootModule
 from sae_lens import __version__
 from sae_lens.config import LanguageModelSAERunnerConfig
 from sae_lens.evals import EvalConfig, run_evals
+from sae_lens.saes.sae import TrainingSAE, TrainStepInput, TrainStepOutput
 from sae_lens.training.activations_store import ActivationsStore
 from sae_lens.training.optim import L1Scheduler, get_lr_scheduler
-from sae_lens.training.training_sae import TrainingSAE, TrainStepOutput
 
 # used to map between parameters which are updated during finetuning and the config str.
 FINETUNING_PARAMETERS = {
@@ -186,7 +186,7 @@ class SAETrainer:
 
             step_output = self._train_step(sae=self.sae, sae_in=layer_acts)
 
-            if self.cfg.log_to_wandb:
+            if self.cfg.logger.log_to_wandb:
                 self._log_train_step(step_output)
                 self._run_and_log_evals()
 
@@ -226,7 +226,7 @@ class SAETrainer:
 
         # log and then reset the feature sparsity every feature_sampling_window steps
         if (self.n_training_steps + 1) % self.cfg.feature_sampling_window == 0:
-            if self.cfg.log_to_wandb:
+            if self.cfg.logger.log_to_wandb:
                 sparsity_log_dict = self._build_sparsity_log_dict()
                 wandb.log(sparsity_log_dict, step=self.n_training_steps)
             self._reset_running_sparsity_stats()
@@ -235,9 +235,11 @@ class SAETrainer:
         # https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html
         with self.autocast_if_enabled:
             train_step_output = self.sae.training_forward_pass(
-                sae_in=sae_in,
-                dead_neuron_mask=self.dead_neurons,
-                current_l1_coefficient=self.current_l1_coefficient,
+                step_input=TrainStepInput(
+                    sae_in=sae_in,
+                    dead_neuron_mask=self.dead_neurons,
+                    current_l1_coefficient=self.current_l1_coefficient,
+                ),
             )
 
             with torch.no_grad():
@@ -270,7 +272,7 @@ class SAETrainer:
 
     @torch.no_grad()
     def _log_train_step(self, step_output: TrainStepOutput):
-        if (self.n_training_steps + 1) % self.cfg.wandb_log_frequency == 0:
+        if (self.n_training_steps + 1) % self.cfg.logger.wandb_log_frequency == 0:
             wandb.log(
                 self._build_train_step_log_dict(
                     output=step_output,
@@ -331,7 +333,8 @@ class SAETrainer:
     def _run_and_log_evals(self):
         # record loss frequently, but not all the time.
         if (self.n_training_steps + 1) % (
-            self.cfg.wandb_log_frequency * self.cfg.eval_every_n_wandb_logs
+            self.cfg.logger.wandb_log_frequency
+            * self.cfg.logger.eval_every_n_wandb_logs
         ) == 0:
             self.sae.eval()
             ignore_tokens = set()
@@ -358,17 +361,8 @@ class SAETrainer:
             # Remove metrics that are not useful for wandb logging
             eval_metrics.pop("metrics/total_tokens_evaluated", None)
 
-            W_dec_norm_dist = self.sae.W_dec.detach().float().norm(dim=1).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().float().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().float().cpu().numpy()
-                eval_metrics["weights/b_gate"] = wandb.Histogram(b_gate_dist)  # type: ignore
-                b_mag_dist = self.sae.b_mag.detach().float().cpu().numpy()
-                eval_metrics["weights/b_mag"] = wandb.Histogram(b_mag_dist)  # type: ignore
+            for key, value in self.sae.log_histograms().items():
+                eval_metrics[key] = wandb.Histogram(value)  # type: ignore
 
             wandb.log(
                 eval_metrics,

sae_lens/training/upload_saes_to_huggingface.py

@@ -14,7 +14,7 @@ from sae_lens.config import (
     SAE_WEIGHTS_FILENAME,
     SPARSITY_FILENAME,
 )
-from sae_lens.sae import SAE
+from sae_lens.saes.sae import SAE
 
 
 def upload_saes_to_huggingface(

{sae_lens-5.9.1.dist-info → sae_lens-6.0.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 5.9.1
+Version: 6.0.0rc1
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
@@ -89,7 +89,7 @@ Please cite the package as follows:
 ```
 @misc{bloom2024saetrainingcodebase,
    title = {SAELens},
-   author = {Joseph Bloom, Curt Tigges, Anthony Duong and David Chanin},
+   author = {Bloom, Joseph and Tigges, Curt and Duong, Anthony and Chanin, David},
    year = {2024},
    howpublished = {\url{https://github.com/jbloomAus/SAELens}},
 }

sae_lens-6.0.0rc1.dist-info/RECORD

@@ -0,0 +1,32 @@
+sae_lens/__init__.py,sha256=ofQyurU7LtxIsg89QFCZe13QsdYpxErRI0x0tiCpB04,2074
+sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/analysis/hooked_sae_transformer.py,sha256=RK0mcLhymXdJInXHcagQggxW9Qf4ptePnH7sKXvGGaU,13727
+sae_lens/analysis/neuronpedia_integration.py,sha256=dFiKRWfuT5iUfTPBPmZydSaNG3VwqZ1asuNbbQv_NCM,18488
+sae_lens/cache_activations_runner.py,sha256=dGK5EHJMHAKDAFyr25fy1COSm-61q-q6kpWENHFMaKk,12561
+sae_lens/config.py,sha256=SPjziXrTyOBjObSi-3s0_mza3Z7WH8gd9NT9pVUfosg,34375
+sae_lens/evals.py,sha256=tjDKmkUM4fBbP9LHZuBLCx37ux8Px9CliTMme3Wjt1A,38898
+sae_lens/load_model.py,sha256=tE70sXsyyyGYW7o506O3eiw1MXyyW6DCQojLG49hWYI,6771
+sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/loading/pretrained_sae_loaders.py,sha256=NcqyH2KDL8Dg66-hjXsBAq1-IwdLEpYfKwbkHxSQbrg,29961
+sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
+sae_lens/pretokenize_runner.py,sha256=0nHQq3s_d80VS8iVK4-e6y_orAYVO8c4RrLGtIDfK_E,6885
+sae_lens/pretrained_saes.yaml,sha256=C_z-7Lxz6ZIy2V-c-4Xw45eAQ926O9aGjocSNuki0xs,573557
+sae_lens/regsitry.py,sha256=yCse5NmVH-ZaPET3jW8r7C_py2DL3yoox40GxGzJ0TI,1098
+sae_lens/sae_training_runner.py,sha256=VRNSAIsZLfcQMfZB8qdnK45PUXwoNvJ-rKt9BVYjMMY,8244
+sae_lens/saes/gated_sae.py,sha256=l5ucq7AZHya6ZClWNNE7CionGSf1ms5m1Ah3IoN6SH4,9916
+sae_lens/saes/jumprelu_sae.py,sha256=DRWgY58894cNh_sYAlefObI4rr0Eb6KHu1WuhTCcvB4,13468
+sae_lens/saes/sae.py,sha256=fd7OEsSXbmVii6QoYI_TRti6dwaxAQyrBcKyX7PxERw,36779
+sae_lens/saes/standard_sae.py,sha256=m2eNL_w6ave-_g7F1eQiwI4qbjMwwjzvxp96RN_WVAw,7110
+sae_lens/saes/topk_sae.py,sha256=aBET4F55A4xMIvZ8AazPtyl3oL-9S7krKx78li0uKGk,11370
+sae_lens/tokenization_and_batching.py,sha256=oUAscjy_LPOrOb8_Ty6eLAcZ0B3HB_wiWjWktgolhG0,4314
+sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/training/activations_store.py,sha256=ilJdcnZWfTDus1bdoqIb1wF_7H8_HWLmf8OCGrybmlA,35998
+sae_lens/training/geometric_median.py,sha256=3kH8ZJAgKStlnZgs6s1uYGDYh004Bl0r4RLhuwT3lBY,3719
+sae_lens/training/optim.py,sha256=AImcc-MAaGDLOBP2hJ4alDFCtaqqgm4cc2eBxIxiQAo,5784
+sae_lens/training/sae_trainer.py,sha256=6TkqbzA0fYluRM8ouI_nU9sz-FaP63axxcnDrVfw37E,16279
+sae_lens/training/upload_saes_to_huggingface.py,sha256=tVC-2Txw7-9XttGlKzM0OSqU8CK7HDO9vIzDMqEwAYU,4366
+sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
+sae_lens-6.0.0rc1.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.0.0rc1.dist-info/METADATA,sha256=wHH-VRtquu-FjZEOHdPJi3zYW3ns7MCT1fVerbPEylc,5326
+sae_lens-6.0.0rc1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+sae_lens-6.0.0rc1.dist-info/RECORD,,