sae-lens 6.0.0rc3__py3-none-any.whl → 6.0.0rc5__py3-none-any.whl

sae_lens/saes/sae.py

@@ -1,5 +1,6 @@
 """Base classes for Sparse Autoencoders (SAEs)."""
 
+import copy
 import json
 import warnings
 from abc import ABC, abstractmethod
@@ -59,23 +60,91 @@ T_SAE = TypeVar("T_SAE", bound="SAE")  # type: ignore
 T_TRAINING_SAE = TypeVar("T_TRAINING_SAE", bound="TrainingSAE")  # type: ignore
 
 
-@dataclass
 class SAEMetadata:
     """Core metadata about how this SAE should be used, if known."""
 
-    model_name: str | None = None
-    hook_name: str | None = None
-    model_class_name: str | None = None
-    hook_head_index: int | None = None
-    model_from_pretrained_kwargs: dict[str, Any] | None = None
-    prepend_bos: bool | None = None
-    exclude_special_tokens: bool | list[int] | None = None
-    neuronpedia_id: str | None = None
-    context_size: int | None = None
-    seqpos_slice: tuple[int | None, ...] | None = None
-    dataset_path: str | None = None
-    sae_lens_version: str = field(default_factory=lambda: __version__)
-    sae_lens_training_version: str = field(default_factory=lambda: __version__)
+    def __init__(self, **kwargs: Any):
+        # Set default version fields with their current behavior
+        self.sae_lens_version = kwargs.pop("sae_lens_version", __version__)
+        self.sae_lens_training_version = kwargs.pop(
+            "sae_lens_training_version", __version__
+        )
+
+        # Set all other attributes dynamically
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __getattr__(self, name: str) -> None:
+        """Return None for any missing attribute (like defaultdict)"""
+        return
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Allow setting any attribute"""
+        super().__setattr__(name, value)
+
+    def __getitem__(self, key: str) -> Any:
+        """Allow dictionary-style access: metadata['key'] - returns None for missing keys"""
+        return getattr(self, key)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Allow dictionary-style assignment: metadata['key'] = value"""
+        setattr(self, key, value)
+
+    def __contains__(self, key: str) -> bool:
+        """Allow 'in' operator: 'key' in metadata"""
+        # Only return True if the attribute was explicitly set (not just defaulting to None)
+        return key in self.__dict__
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Dictionary-style get with default"""
+        value = getattr(self, key)
+        # If the attribute wasn't explicitly set and we got None from __getattr__,
+        # use the provided default instead
+        if key not in self.__dict__ and value is None:
+            return default
+        return value
+
+    def keys(self):
+        """Return all explicitly set attribute names"""
+        return self.__dict__.keys()
+
+    def values(self):
+        """Return all explicitly set attribute values"""
+        return self.__dict__.values()
+
+    def items(self):
+        """Return all explicitly set attribute name-value pairs"""
+        return self.__dict__.items()
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization"""
+        return self.__dict__.copy()
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SAEMetadata":
+        """Create from dictionary"""
+        return cls(**data)
+
+    def __repr__(self) -> str:
+        return f"SAEMetadata({self.__dict__})"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, SAEMetadata):
+            return False
+        return self.__dict__ == other.__dict__
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> "SAEMetadata":
+        """Support for deep copying"""
+
+        return SAEMetadata(**copy.deepcopy(self.__dict__, memo))
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Support for pickling"""
+        return self.__dict__
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Support for unpickling"""
+        self.__dict__.update(state)
 
 
 @dataclass
@@ -99,7 +168,7 @@ class SAEConfig(ABC):
 
     def to_dict(self) -> dict[str, Any]:
         res = {field.name: getattr(self, field.name) for field in fields(self)}
-        res["metadata"] = asdict(self.metadata)
+        res["metadata"] = self.metadata.to_dict()
         res["architecture"] = self.architecture()
         return res
 
@@ -124,7 +193,7 @@ class SAEConfig(ABC):
             "layer_norm",
         ]:
             raise ValueError(
-                f"normalize_activations must be none, expected_average_only_in, constant_norm_rescale, or layer_norm. Got {self.normalize_activations}"
+                f"normalize_activations must be none, expected_average_only_in, layer_norm, or constant_norm_rescale. Got {self.normalize_activations}"
             )
 
 
@@ -238,9 +307,8 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
 
             self.run_time_activation_norm_fn_in = run_time_activation_norm_fn_in
             self.run_time_activation_norm_fn_out = run_time_activation_norm_fn_out
-
         elif self.cfg.normalize_activations == "layer_norm":
-
+            #  we need to scale the norm of the input and store the scaling factor
             def run_time_activation_ln_in(
                 x: torch.Tensor, eps: float = 1e-5
             ) -> torch.Tensor:
@@ -522,7 +590,7 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
         device: str = "cpu",
         force_download: bool = False,
         converter: PretrainedSaeHuggingfaceLoader | None = None,
-    ) -> tuple[T_SAE, dict[str, Any], torch.Tensor | None]:
+    ) -> T_SAE:
         """
         Load a pretrained SAE from the Hugging Face model hub.
 
@@ -530,7 +598,28 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
             release: The release name. This will be mapped to a huggingface repo id based on the pretrained_saes.yaml file.
             id: The id of the SAE to load. This will be mapped to a path in the huggingface repo.
             device: The device to load the SAE on.
-            return_sparsity_if_present: If True, will return the log sparsity tensor if it is present in the model directory in the Hugging Face model hub.
+        """
+        return cls.from_pretrained_with_cfg_and_sparsity(
+            release, sae_id, device, force_download, converter=converter
+        )[0]
+
+    @classmethod
+    def from_pretrained_with_cfg_and_sparsity(
+        cls: Type[T_SAE],
+        release: str,
+        sae_id: str,
+        device: str = "cpu",
+        force_download: bool = False,
+        converter: PretrainedSaeHuggingfaceLoader | None = None,
+    ) -> tuple[T_SAE, dict[str, Any], torch.Tensor | None]:
+        """
+        Load a pretrained SAE from the Hugging Face model hub, along with its config dict and sparsity, if present.
+        In SAELens <= 5.x.x, this was called SAE.from_pretrained().
+
+        Args:
+            release: The release name. This will be mapped to a huggingface repo id based on the pretrained_saes.yaml file.
+            id: The id of the SAE to load. This will be mapped to a path in the huggingface repo.
+            device: The device to load the SAE on.
         """
 
         # get sae directory
@@ -643,11 +732,67 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
     ) -> type[SAEConfig]:
         return SAEConfig
 
+    ### Methods to support deprecated usage of SAE.from_pretrained() ###
+
+    def __getitem__(self, index: int) -> Any:
+        """
+        Support indexing for backward compatibility with tuple unpacking.
+        DEPRECATED: SAE.from_pretrained() no longer returns a tuple.
+        Use SAE.from_pretrained_with_cfg_and_sparsity() instead.
+        """
+        warnings.warn(
+            "Indexing SAE objects is deprecated. SAE.from_pretrained() now returns "
+            "only the SAE object. Use SAE.from_pretrained_with_cfg_and_sparsity() "
+            "to get the config dict and sparsity as well.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        if index == 0:
+            return self
+        if index == 1:
+            return self.cfg.to_dict()
+        if index == 2:
+            return None
+        raise IndexError(f"SAE tuple index {index} out of range")
+
+    def __iter__(self):
+        """
+        Support unpacking for backward compatibility with tuple unpacking.
+        DEPRECATED: SAE.from_pretrained() no longer returns a tuple.
+        Use SAE.from_pretrained_with_cfg_and_sparsity() instead.
+        """
+        warnings.warn(
+            "Unpacking SAE objects is deprecated. SAE.from_pretrained() now returns "
+            "only the SAE object. Use SAE.from_pretrained_with_cfg_and_sparsity() "
+            "to get the config dict and sparsity as well.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        yield self
+        yield self.cfg.to_dict()
+        yield None
+
+    def __len__(self) -> int:
+        """
+        Support len() for backward compatibility with tuple unpacking.
+        DEPRECATED: SAE.from_pretrained() no longer returns a tuple.
+        Use SAE.from_pretrained_with_cfg_and_sparsity() instead.
+        """
+        warnings.warn(
+            "Getting length of SAE objects is deprecated. SAE.from_pretrained() now returns "
+            "only the SAE object. Use SAE.from_pretrained_with_cfg_and_sparsity() "
+            "to get the config dict and sparsity as well.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        return 3
+
 
 @dataclass(kw_only=True)
 class TrainingSAEConfig(SAEConfig, ABC):
-    noise_scale: float = 0.0
-    mse_loss_normalization: str | None = None
     # https://transformer-circuits.pub/2024/april-update/index.html#training-saes
     # 0.1 corresponds to the "heuristic" initialization, use None to disable
     decoder_init_norm: float | None = 0.1
@@ -680,9 +825,6 @@ class TrainingSAEConfig(SAEConfig, ABC):
     def from_dict(
         cls: type[T_TRAINING_SAE_CONFIG], config_dict: dict[str, Any]
     ) -> T_TRAINING_SAE_CONFIG:
-        # remove any keys that are not in the dataclass
-        # since we sometimes enhance the config with the whole LM runner config
-        valid_config_dict = filter_valid_dataclass_fields(config_dict, cls)
         cfg_class = cls
         if "architecture" in config_dict:
             cfg_class = get_sae_training_class(config_dict["architecture"])[1]
@@ -690,6 +832,9 @@ class TrainingSAEConfig(SAEConfig, ABC):
             raise ValueError(
                 f"SAE config class {cls} does not match dict config class {type(cfg_class)}"
             )
+        # remove any keys that are not in the dataclass
+        # since we sometimes enhance the config with the whole LM runner config
+        valid_config_dict = filter_valid_dataclass_fields(config_dict, cfg_class)
         if "metadata" in config_dict:
             valid_config_dict["metadata"] = SAEMetadata(**config_dict["metadata"])
         return cfg_class(**valid_config_dict)
@@ -698,6 +843,7 @@ class TrainingSAEConfig(SAEConfig, ABC):
         return {
             **super().to_dict(),
             **asdict(self),
+            "metadata": self.metadata.to_dict(),
             "architecture": self.architecture(),
         }
 
@@ -708,12 +854,14 @@ class TrainingSAEConfig(SAEConfig, ABC):
         Creates a dictionary containing attributes corresponding to the fields
         defined in the base SAEConfig class.
         """
-        base_config_field_names = {f.name for f in fields(SAEConfig)}
+        base_sae_cfg_class = get_sae_class(self.architecture())[1]
+        base_config_field_names = {f.name for f in fields(base_sae_cfg_class)}
         result_dict = {
             field_name: getattr(self, field_name)
             for field_name in base_config_field_names
         }
         result_dict["architecture"] = self.architecture()
+        result_dict["metadata"] = self.metadata.to_dict()
         return result_dict
 
 
@@ -726,7 +874,7 @@ class TrainingSAE(SAE[T_TRAINING_SAE_CONFIG], ABC):
         # Turn off hook_z reshaping for training mode - the activation store
         # is expected to handle reshaping before passing data to the SAE
         self.turn_off_forward_pass_hook_z_reshaping()
-        self.mse_loss_fn = self._get_mse_loss_fn()
+        self.mse_loss_fn = mse_loss
 
     @abstractmethod
     def get_coefficients(self) -> dict[str, float | TrainCoefficientConfig]: ...
@@ -861,27 +1009,6 @@ class TrainingSAE(SAE[T_TRAINING_SAE_CONFIG], ABC):
         """
         return self.process_state_dict_for_saving(state_dict)
 
-    def _get_mse_loss_fn(self) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
-        """Get the MSE loss function based on config."""
-
-        def standard_mse_loss_fn(
-            preds: torch.Tensor, target: torch.Tensor
-        ) -> torch.Tensor:
-            return torch.nn.functional.mse_loss(preds, target, reduction="none")
-
-        def batch_norm_mse_loss_fn(
-            preds: torch.Tensor, target: torch.Tensor
-        ) -> torch.Tensor:
-            target_centered = target - target.mean(dim=0, keepdim=True)
-            normalization = target_centered.norm(dim=-1, keepdim=True)
-            return torch.nn.functional.mse_loss(preds, target, reduction="none") / (
-                normalization + 1e-6
-            )
-
-        if self.cfg.mse_loss_normalization == "dense_batch":
-            return batch_norm_mse_loss_fn
-        return standard_mse_loss_fn
-
     @torch.no_grad()
     def remove_gradient_parallel_to_decoder_directions(self) -> None:
         """Remove gradient components parallel to decoder directions."""
@@ -943,3 +1070,7 @@ def _disable_hooks(sae: SAE[Any]):
     finally:
         for hook_name, hook in sae.hook_dict.items():
             setattr(sae, hook_name, hook)
+
+
+def mse_loss(preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.mse_loss(preds, target, reduction="none")

sae_lens/saes/standard_sae.py

@@ -67,7 +67,7 @@ class StandardSAE(SAE[StandardSAEConfig]):
         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)
+        # Apply the activation function (e.g., ReLU, depending on config)
         return self.hook_sae_acts_post(self.activation_fn(hidden_pre))
 
     def decode(
@@ -81,7 +81,7 @@ class StandardSAE(SAE[StandardSAEConfig]):
         sae_out_pre = feature_acts @ self.W_dec + self.b_dec
         # 2) hook reconstruction
         sae_out_pre = self.hook_sae_recons(sae_out_pre)
-        # 4) optional out-normalization (e.g. constant_norm_rescale or layer_norm)
+        # 4) optional out-normalization (e.g. constant_norm_rescale)
         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)
@@ -136,16 +136,9 @@ class StandardTrainingSAE(TrainingSAE[StandardTrainingSAEConfig]):
         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
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+        return feature_acts, hidden_pre
 
     def calculate_aux_loss(
         self,

sae_lens/saes/topk_sae.py

@@ -91,8 +91,7 @@ class TopKSAE(SAE[TopKSAEConfig]):
     ) -> Float[torch.Tensor, "... d_sae"]:
         """
         Converts input x into feature activations.
-        Uses topk activation from the config (cfg.activation_fn == "topk")
-        under the hood.
+        Uses topk activation 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)
@@ -116,6 +115,13 @@ class TopKSAE(SAE[TopKSAEConfig]):
     def get_activation_fn(self) -> Callable[[torch.Tensor], torch.Tensor]:
         return TopK(self.cfg.k)
 
+    @override
+    @torch.no_grad()
+    def fold_W_dec_norm(self) -> None:
+        raise NotImplementedError(
+            "Folding W_dec_norm is not safe for TopKSAEs, as this may change the topk activations"
+        )
+
 
 @dataclass
 class TopKTrainingSAEConfig(TrainingSAEConfig):
@@ -156,18 +162,11 @@ class TopKTrainingSAE(TrainingSAE[TopKTrainingSAEConfig]):
         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
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+        return feature_acts, hidden_pre
 
+    @override
     def calculate_aux_loss(
         self,
         step_input: TrainStepInput,
@@ -184,6 +183,13 @@ class TopKTrainingSAE(TrainingSAE[TopKTrainingSAEConfig]):
         )
         return {"auxiliary_reconstruction_loss": topk_loss}
 
+    @override
+    @torch.no_grad()
+    def fold_W_dec_norm(self) -> None:
+        raise NotImplementedError(
+            "Folding W_dec_norm is not safe for TopKSAEs, as this may change the topk activations"
+        )
+
     @override
     def get_activation_fn(self) -> Callable[[torch.Tensor], torch.Tensor]:
         return TopK(self.cfg.k)

sae_lens/training/activation_scaler.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from statistics import mean
 
 import torch
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 from sae_lens.training.types import DataProvider
 

sae_lens/training/activations_store.py

@@ -161,8 +161,6 @@ class ActivationsStore:
     ) -> ActivationsStore:
         if sae.cfg.metadata.hook_name is None:
             raise ValueError("hook_name is required")
-        if sae.cfg.metadata.hook_head_index is None:
-            raise ValueError("hook_head_index is required")
         if sae.cfg.metadata.context_size is None:
             raise ValueError("context_size is required")
         if sae.cfg.metadata.prepend_bos is None:
@@ -430,7 +428,7 @@ class ActivationsStore:
         ):
             # temporalily set estimated_norm_scaling_factor to 1.0 so the dataloader works
             self.estimated_norm_scaling_factor = 1.0
-            acts = self.next_batch()[0]
+            acts = self.next_batch()[:, 0]
             self.estimated_norm_scaling_factor = None
             norms_per_batch.append(acts.norm(dim=-1).mean().item())
         mean_norm = np.mean(norms_per_batch)

sae_lens/training/sae_trainer.py

@@ -7,7 +7,7 @@ import torch
 import wandb
 from safetensors.torch import save_file
 from torch.optim import Adam
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 from sae_lens import __version__
 from sae_lens.config import SAETrainerConfig
@@ -161,6 +161,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
         return (self.n_forward_passes_since_fired > self.cfg.dead_feature_window).bool()
 
     def fit(self) -> T_TRAINING_SAE:
+        self.sae.to(self.cfg.device)
         pbar = tqdm(total=self.cfg.total_training_samples, desc="Training SAE")
 
         if self.sae.cfg.normalize_activations == "expected_average_only_in":
@@ -194,10 +195,11 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             )
             self.activation_scaler.scaling_factor = None
 
-        # save final sae group to checkpoints folder
+        # save final inference sae group to checkpoints folder
         self.save_checkpoint(
             checkpoint_name=f"final_{self.n_training_samples}",
             wandb_aliases=["final_model"],
+            save_inference_model=True,
         )
 
         pbar.close()
@@ -207,11 +209,17 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
         self,
         checkpoint_name: str,
         wandb_aliases: list[str] | None = None,
+        save_inference_model: bool = False,
     ) -> None:
         checkpoint_path = Path(self.cfg.checkpoint_path) / checkpoint_name
         checkpoint_path.mkdir(exist_ok=True, parents=True)
 
-        weights_path, cfg_path = self.sae.save_model(str(checkpoint_path))
+        save_fn = (
+            self.sae.save_inference_model
+            if save_inference_model
+            else self.sae.save_model
+        )
+        weights_path, cfg_path = save_fn(str(checkpoint_path))
 
         sparsity_path = checkpoint_path / SPARSITY_FILENAME
         save_file({"sparsity": self.log_feature_sparsity}, sparsity_path)

sae_lens/training/upload_saes_to_huggingface.py

@@ -88,7 +88,7 @@ def _create_default_readme(repo_id: str, sae_ids: Iterable[str]) -> str:
         ```python
         from sae_lens import SAE
 
-        sae, cfg_dict, sparsity = SAE.from_pretrained("{repo_id}", "<sae_id>")
+        sae = SAE.from_pretrained("{repo_id}", "<sae_id>")
         ```
         """
     )

{sae_lens-6.0.0rc3.dist-info → sae_lens-6.0.0rc5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.0.0rc3
+Version: 6.0.0rc5
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch
@@ -80,7 +80,7 @@ Pre-trained SAEs for various models can be imported via SAE Lens. See this [page
 
 ## Join the Slack!
 
-Feel free to join the [Open Source Mechanistic Interpretability Slack](https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-2o756ku1c-_yKBeUQMVfS_p_qcK6QLeA) for support!
+Feel free to join the [Open Source Mechanistic Interpretability Slack](https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-375zalm04-GFd5tdBU1yLKlu_T_JSqZQ) for support!
 
 ## Citation
 

sae_lens-6.0.0rc5.dist-info/RECORD

@@ -0,0 +1,37 @@
+sae_lens/__init__.py,sha256=hiHDLT9_1V7iVulw5hwqDqDj2HVxUR9I88xOfYx6X94,2861
+sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/analysis/hooked_sae_transformer.py,sha256=Eyg1Y2hVIHNuiiLOCTgzstOuW6iA-7hPHqaGR8y_vMs,13809
+sae_lens/analysis/neuronpedia_integration.py,sha256=MrENqc81Mc2SMbxGjbwHzpkGUCAFKSf0i4EdaUF2Oj4,18707
+sae_lens/cache_activations_runner.py,sha256=L5hhuU2-zPQr2S3L64GMKKLeMQfqXxwDl8NbuOtrybI,12567
+sae_lens/config.py,sha256=9Lg4HkQvj1t9QZJdmC071lyJMc_iqNQknosT7zOYfwM,27278
+sae_lens/constants.py,sha256=CSjmiZ-bhjQeVLyRvWxAjBokCgkfM8mnvd7-vxLIWTY,639
+sae_lens/evals.py,sha256=kQyrzczKaVD9rHwfFa_DxL_gMXDxsoIVHmsFIPIU2bY,38696
+sae_lens/llm_sae_training_runner.py,sha256=58XbDylw2fPOD7C-ZfSAjeNqJLXB05uHGTuiYVVbXXY,13354
+sae_lens/load_model.py,sha256=dBB_9gO6kWyQ4sXHq7qB8T3YUlXm3PGwYcpR4UVW4QY,8633
+sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sae_lens/loading/pretrained_sae_loaders.py,sha256=5XEU4uFFeGCePwqDwhlE7CrFGRSI0U9Cu-UQVa33Y1E,36432
+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=nhHW1auhyi4GHYrjUnHQqbNVhI5cMJv-HThzbzU1xG0,574145
+sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
+sae_lens/saes/__init__.py,sha256=v6mfeDzyGYtT6x5SszAQtkldTXwPE-V_iwOlrT_pDwQ,1008
+sae_lens/saes/gated_sae.py,sha256=0zd66bH04nsaGk3bxHk10hsZofa2GrFbMo15LOsuqgU,9233
+sae_lens/saes/jumprelu_sae.py,sha256=iwmPQJ4XpIxzgosty680u8Zj7x1uVZhM75kPOT3obi0,12060
+sae_lens/saes/sae.py,sha256=ZEXEXFVtrtFrzuOV3nyweTBleNCV4EDGh1ImaF32uqg,39618
+sae_lens/saes/standard_sae.py,sha256=PfkGLsw_6La3PXHOQL0u7qQsaZsXCJqYCeCcRDj5n64,6274
+sae_lens/saes/topk_sae.py,sha256=kmry1FE1H06OvCfn84V-j2JfWGKcU5b2urwAq_Oq5j4,9893
+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/activation_scaler.py,sha256=seEE-2Qd2JMHxqgnsNWPt-DGtYGZxWPnOwCGuVNSOtI,1719
+sae_lens/training/activations_store.py,sha256=z8erbiB6ODbsqlu-bwEWbyj4XZvgsVgjCRBuQovqp2Q,32612
+sae_lens/training/mixing_buffer.py,sha256=vDpYG5ZE70szDvBsRKcNHEES3h_WTKJ16qDYk5jPOVA,2015
+sae_lens/training/optim.py,sha256=KXdOym-Ly3f2aFbndRc0JEH0Wa7u1BE5ljxGN3YtouQ,6836
+sae_lens/training/sae_trainer.py,sha256=9K0VudwSTJp9OlCVzaU_ngZ0WlYNrN6-ozTCCAxR9_k,15421
+sae_lens/training/types.py,sha256=qSjmGzXf3MLalygG0psnVjmhX_mpLmL47MQtZfe7qxg,81
+sae_lens/training/upload_saes_to_huggingface.py,sha256=r_WzI1zLtGZ5TzAxuG3xa_8T09j3zXJrWd_vzPsPGkQ,4469
+sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
+sae_lens/util.py,sha256=mCwLAilGMVo8Scm7CIsCafU7GsfmBvCcjwmloI4Ly7Y,1718
+sae_lens-6.0.0rc5.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.0.0rc5.dist-info/METADATA,sha256=ZrBaBFeIuM-ZJ9r0HHKakxnx3tGv7Zf6l_Z2OIdBxIU,5326
+sae_lens-6.0.0rc5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.0.0rc5.dist-info/RECORD,,

sae_lens/training/geometric_median.py

@@ -1,101 +0,0 @@
-from types import SimpleNamespace
-
-import torch
-import tqdm
-
-
-def weighted_average(points: torch.Tensor, weights: torch.Tensor):
-    weights = weights / weights.sum()
-    return (points * weights.view(-1, 1)).sum(dim=0)
-
-
-@torch.no_grad()
-def geometric_median_objective(
-    median: torch.Tensor, points: torch.Tensor, weights: torch.Tensor
-) -> torch.Tensor:
-    norms = torch.linalg.norm(points - median.view(1, -1), dim=1)  # type: ignore
-
-    return (norms * weights).sum()
-
-
-def compute_geometric_median(
-    points: torch.Tensor,
-    weights: torch.Tensor | None = None,
-    eps: float = 1e-6,
-    maxiter: int = 100,
-    ftol: float = 1e-20,
-    do_log: bool = False,
-):
-    """
-    :param points: ``torch.Tensor`` of shape ``(n, d)``
-    :param weights: Optional ``torch.Tensor`` of shape :math:``(n,)``.
-    :param eps: Smallest allowed value of denominator, to avoid divide by zero.
-        Equivalently, this is a smoothing parameter. Default 1e-6.
-    :param maxiter: Maximum number of Weiszfeld iterations. Default 100
-    :param ftol: If objective value does not improve by at least this `ftol` fraction, terminate the algorithm. Default 1e-20.
-    :param do_log: If true will return a log of function values encountered through the course of the algorithm
-    :return: SimpleNamespace object with fields
-        - `median`: estimate of the geometric median, which is a ``torch.Tensor`` object of shape :math:``(d,)``
-        - `termination`: string explaining how the algorithm terminated.
-        - `logs`: function values encountered through the course of the algorithm in a list (None if do_log is false).
-    """
-    with torch.no_grad():
-        if weights is None:
-            weights = torch.ones((points.shape[0],), device=points.device)
-        # initialize median estimate at mean
-        new_weights = weights
-        median = weighted_average(points, weights)
-        objective_value = geometric_median_objective(median, points, weights)
-        logs = [objective_value] if do_log else None
-
-        # Weiszfeld iterations
-        early_termination = False
-        pbar = tqdm.tqdm(range(maxiter))
-        for _ in pbar:
-            prev_obj_value = objective_value
-
-            norms = torch.linalg.norm(points - median.view(1, -1), dim=1)  # type: ignore
-            new_weights = weights / torch.clamp(norms, min=eps)
-            median = weighted_average(points, new_weights)
-            objective_value = geometric_median_objective(median, points, weights)
-
-            if logs is not None:
-                logs.append(objective_value)
-            if abs(prev_obj_value - objective_value) <= ftol * objective_value:
-                early_termination = True
-                break
-
-            pbar.set_description(f"Objective value: {objective_value:.4f}")
-
-    median = weighted_average(points, new_weights)  # allow autodiff to track it
-    return SimpleNamespace(
-        median=median,
-        new_weights=new_weights,
-        termination=(
-            "function value converged within tolerance"
-            if early_termination
-            else "maximum iterations reached"
-        ),
-        logs=logs,
-    )
-
-
-if __name__ == "__main__":
-    import time
-
-    TOLERANCE = 1e-2
-
-    dim1 = 10000
-    dim2 = 768
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    sample = (
-        torch.randn((dim1, dim2), device=device) * 100
-    )  # seems to be the order of magnitude of the actual use case
-    weights = torch.randn((dim1,), device=device)
-
-    torch.tensor(weights, device=device)
-
-    tic = time.perf_counter()
-    new = compute_geometric_median(sample, weights=weights, maxiter=100)
-    print(f"new code takes {time.perf_counter()-tic} seconds!")  # noqa: T201