sae-lens 6.32.1__py3-none-any.whl → 6.34.1__py3-none-any.whl

sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.32.1"
+__version__ = "6.34.1"
 
 import logging
 

sae_lens/analysis/hooked_sae_transformer.py

@@ -1,13 +1,14 @@
-import logging
 from contextlib import contextmanager
 from typing import Any, Callable
 
 import torch
+from torch import nn
 from transformer_lens.ActivationCache import ActivationCache
 from transformer_lens.components.mlps.can_be_used_as_mlp import CanBeUsedAsMLP
 from transformer_lens.hook_points import HookPoint  # Hooking utilities
 from transformer_lens.HookedTransformer import HookedTransformer
 
+from sae_lens import logger
 from sae_lens.saes.sae import SAE
 
 SingleLoss = torch.Tensor  # Type alias for a single element tensor
@@ -15,6 +16,72 @@ LossPerToken = torch.Tensor
 Loss = SingleLoss | LossPerToken
 
 
+class _SAEWrapper(nn.Module):
+    """Wrapper for SAE/Transcoder that handles error term and hook coordination.
+
+    For SAEs (input_hook == output_hook), _captured_input stays None and we use
+    the forward argument directly. For transcoders, _captured_input is set at
+    the input hook via capture_input().
+
+    Implementation Note:
+        The SAE is stored in __dict__ directly rather than as a registered submodule.
+        This is intentional: PyTorch's module registration would add a ".sae." prefix
+        to all hook names in the cache (e.g., "blocks.0.hook_mlp_out.sae.hook_sae_input"
+        instead of "blocks.0.hook_mlp_out.hook_sae_input"). By storing in __dict__ and
+        copying hooks directly to the wrapper, we preserve the expected cache paths
+        for backwards compatibility.
+    """
+
+    def __init__(self, sae: SAE[Any], use_error_term: bool = False):
+        super().__init__()
+        # Store SAE in __dict__ to avoid registering as submodule. This keeps cache
+        # paths clean by avoiding a ".sae." prefix on hook names. See class docstring.
+        self.__dict__["_sae"] = sae
+        # Copy SAE's hooks directly to wrapper so they appear at the right path
+        for name, hook in sae.hook_dict.items():
+            setattr(self, name, hook)
+        self.use_error_term = use_error_term
+        self._captured_input: torch.Tensor | None = None
+
+    @property
+    def sae(self) -> SAE[Any]:
+        return self.__dict__["_sae"]
+
+    def capture_input(self, x: torch.Tensor) -> None:
+        """Capture input at input hook (for transcoders).
+
+        Note: We don't clone the tensor here - the input should not be modified
+        in-place between capture and use, and avoiding clone preserves memory.
+        """
+        self._captured_input = x
+
+    def forward(self, original_output: torch.Tensor) -> torch.Tensor:
+        """Run SAE/transcoder at output hook location."""
+        # For SAE: use original_output as input (same hook for input/output)
+        # For transcoder: use captured input from earlier hook
+        sae_input = (
+            self._captured_input
+            if self._captured_input is not None
+            else original_output
+        )
+
+        # Temporarily disable SAE's internal use_error_term - we handle it here
+        # Use _use_error_term directly to avoid triggering deprecation warning
+        sae_use_error_term = self.sae._use_error_term
+        self.sae._use_error_term = False
+        try:
+            sae_out = self.sae(sae_input)
+
+            if self.use_error_term:
+                error = original_output - sae_out.detach()
+                sae_out = sae_out + error
+
+            return sae_out
+        finally:
+            self.sae._use_error_term = sae_use_error_term
+            self._captured_input = None
+
+
 def get_deep_attr(obj: Any, path: str):
     """Helper function to get a nested attribute from a object.
     In practice used to access HookedTransformer HookPoints (eg model.blocks[0].attn.hook_z)
@@ -78,88 +145,130 @@ class HookedSAETransformer(HookedTransformer):
             add_hook_in_to_mlp(block.mlp)  # type: ignore
         self.setup()
 
-        self.acts_to_saes: dict[str, SAE] = {}  # type: ignore
+        self._acts_to_saes: dict[str, _SAEWrapper] = {}
+        # Track output hooks used by transcoders for cleanup
+        self._transcoder_output_hooks: dict[str, str] = {}
+
+    @property
+    def acts_to_saes(self) -> dict[str, SAE[Any]]:
+        """Returns a dict mapping hook names to attached SAEs."""
+        return {name: wrapper.sae for name, wrapper in self._acts_to_saes.items()}
 
     def add_sae(self, sae: SAE[Any], use_error_term: bool | None = None):
-        """Attaches an SAE to the model
+        """Attaches an SAE or Transcoder to the model.
 
-        WARNING: This sae will be permanantly attached until you remove it with reset_saes. This function will also overwrite any existing SAE attached to the same hook point.
+        WARNING: This SAE will be permanently attached until you remove it with
+        reset_saes. This function will also overwrite any existing SAE attached
+        to the same hook point.
 
         Args:
-            sae: SparseAutoencoderBase. The SAE to attach to the model
-            use_error_term: (bool | None) If provided, will set the use_error_term attribute of the SAE to this value. Determines whether the SAE returns input or reconstruction. Defaults to None.
+            sae: The SAE or Transcoder to attach to the model.
+            use_error_term: If True, computes error term so output matches what the
+                model would have produced without the SAE. This works for both SAEs
+                (where input==output hook) and transcoders (where they differ).
+                Defaults to None (uses SAE's existing setting).
         """
-        act_name = sae.cfg.metadata.hook_name
-        if (act_name not in self.acts_to_saes) and (act_name not in self.hook_dict):
-            logging.warning(
-                f"No hook found for {act_name}. Skipping. Check model.hook_dict for available hooks."
+        input_hook = sae.cfg.metadata.hook_name
+        output_hook = sae.cfg.metadata.hook_name_out or input_hook
+
+        if (input_hook not in self._acts_to_saes) and (
+            input_hook not in self.hook_dict
+        ):
+            logger.warning(
+                f"No hook found for {input_hook}. Skipping. Check model.hook_dict for available hooks."
             )
             return
 
-        if use_error_term is not None:
-            if not hasattr(sae, "_original_use_error_term"):
-                sae._original_use_error_term = sae.use_error_term  # type: ignore
-            sae.use_error_term = use_error_term
-        self.acts_to_saes[act_name] = sae
-        set_deep_attr(self, act_name, sae)
+        # Check if output hook exists (either as hook_dict entry or already has SAE attached)
+        output_hook_exists = (
+            output_hook in self.hook_dict
+            or output_hook in self._acts_to_saes
+            or any(v == output_hook for v in self._transcoder_output_hooks.values())
+        )
+        if not output_hook_exists:
+            logger.warning(f"No hook found for output {output_hook}. Skipping.")
+            return
+
+        # Always use wrapper - it handles both SAEs and transcoders uniformly
+        # If use_error_term not specified, respect SAE's existing setting
+        effective_use_error_term = (
+            use_error_term if use_error_term is not None else sae.use_error_term
+        )
+        wrapper = _SAEWrapper(sae, use_error_term=effective_use_error_term)
+
+        # For transcoders (input != output), capture input at input hook
+        if input_hook != output_hook:
+            input_hook_point = get_deep_attr(self, input_hook)
+            if isinstance(input_hook_point, HookPoint):
+                input_hook_point.add_hook(
+                    lambda tensor, hook: (wrapper.capture_input(tensor), tensor)[1],  # noqa: ARG005
+                    dir="fwd",
+                    is_permanent=True,
+                )
+            self._transcoder_output_hooks[input_hook] = output_hook
+
+        # Store wrapper in _acts_to_saes and at output hook
+        self._acts_to_saes[input_hook] = wrapper
+        set_deep_attr(self, output_hook, wrapper)
         self.setup()
 
-    def _reset_sae(self, act_name: str, prev_sae: SAE[Any] | None = None):
-        """Resets an SAE that was attached to the model
+    def _reset_sae(
+        self, act_name: str, prev_wrapper: _SAEWrapper | None = None
+    ) -> None:
+        """Resets an SAE that was attached to the model.
 
         By default will remove the SAE from that hook_point.
-        If prev_sae is provided, will replace the current SAE with the provided one.
-        This is mainly used to restore previously attached SAEs after temporarily running with different SAEs (eg with run_with_saes)
+        If prev_wrapper is provided, will restore that wrapper's SAE with its settings.
 
         Args:
-            act_name: str. The hook_name of the SAE to reset
-            prev_sae: SAE | None. The SAE to replace the current one with. If None, will just remove the SAE from this hook point. Defaults to None
+            act_name: The hook_name of the SAE to reset.
+            prev_wrapper: The previous wrapper to restore. If None, will just
+                remove the SAE from this hook point. Defaults to None.
         """
-        if act_name not in self.acts_to_saes:
-            logging.warning(
+        if act_name not in self._acts_to_saes:
+            logger.warning(
                 f"No SAE is attached to {act_name}. There's nothing to reset."
             )
             return
 
-        current_sae = self.acts_to_saes[act_name]
-        if hasattr(current_sae, "_original_use_error_term"):
-            current_sae.use_error_term = current_sae._original_use_error_term  # type: ignore
-            delattr(current_sae, "_original_use_error_term")
+        # Determine output hook location (different from input for transcoders)
+        output_hook = self._transcoder_output_hooks.pop(act_name, act_name)
+
+        # For transcoders, clear permanent hooks from input hook point
+        if output_hook != act_name:
+            input_hook_point = get_deep_attr(self, act_name)
+            if isinstance(input_hook_point, HookPoint):
+                input_hook_point.remove_hooks(dir="fwd", including_permanent=True)
+
+        # Reset output hook location
+        set_deep_attr(self, output_hook, HookPoint())
+        del self._acts_to_saes[act_name]
 
-        if prev_sae is not None:
-            set_deep_attr(self, act_name, prev_sae)
-            self.acts_to_saes[act_name] = prev_sae
-        else:
-            set_deep_attr(self, act_name, HookPoint())
-            del self.acts_to_saes[act_name]
+        if prev_wrapper is not None:
+            # Rebuild hook_dict before adding new SAE
+            self.setup()
+            self.add_sae(prev_wrapper.sae, use_error_term=prev_wrapper.use_error_term)
 
     def reset_saes(
         self,
         act_names: str | list[str] | None = None,
-        prev_saes: list[SAE[Any] | None] | None = None,
-    ):
-        """Reset the SAEs attached to the model
+    ) -> None:
+        """Reset the SAEs attached to the model.
 
-        If act_names are provided will just reset SAEs attached to those hooks. Otherwise will reset all SAEs attached to the model.
-        Optionally can provide a list of prev_saes to reset to. This is mainly used to restore previously attached SAEs after temporarily running with different SAEs (eg with run_with_saes).
+        If act_names are provided will just reset SAEs attached to those hooks.
+        Otherwise will reset all SAEs attached to the model.
 
         Args:
-            act_names (str | list[str] | None): The act_names of the SAEs to reset. If None, will reset all SAEs attached to the model. Defaults to None.
-            prev_saes (list[SAE | None] | None): List of SAEs to replace the current ones with. If None, will just remove the SAEs. Defaults to None.
+            act_names: The act_names of the SAEs to reset. If None, will reset
+                all SAEs attached to the model. Defaults to None.
         """
         if isinstance(act_names, str):
             act_names = [act_names]
         elif act_names is None:
-            act_names = list(self.acts_to_saes.keys())
-
-        if prev_saes:
-            if len(act_names) != len(prev_saes):
-                raise ValueError("act_names and prev_saes must have the same length")
-        else:
-            prev_saes = [None] * len(act_names)  # type: ignore
+            act_names = list(self._acts_to_saes.keys())
 
-        for act_name, prev_sae in zip(act_names, prev_saes):  # type: ignore
-            self._reset_sae(act_name, prev_sae)
+        for act_name in act_names:
+            self._reset_sae(act_name)
 
         self.setup()
 
@@ -269,41 +378,32 @@ class HookedSAETransformer(HookedTransformer):
         reset_saes_end: bool = True,
         use_error_term: bool | None = None,
     ):
-        """
-        A context manager for adding temporary SAEs to the model.
-        See HookedTransformer.hooks for a similar context manager for hooks.
-        By default will keep track of previously attached SAEs, and restore them when the context manager exits.
-
-        Example:
-
-        .. code-block:: python
-
-            from transformer_lens import HookedSAETransformer
-            from sae_lens.saes.sae import SAE
-
-            model = HookedSAETransformer.from_pretrained('gpt2-small')
-            sae_cfg = SAEConfig(...)
-            sae = SAE(sae_cfg)
-            with model.saes(saes=[sae]):
-                spliced_logits = model(text)
+        """A context manager for adding temporary SAEs to the model.
 
+        See HookedTransformer.hooks for a similar context manager for hooks.
+        By default will keep track of previously attached SAEs, and restore
+        them when the context manager exits.
 
         Args:
-            saes (SAE | list[SAE]): SAEs to be attached.
-            reset_saes_end (bool): If True, removes all SAEs added by this context manager when the context manager exits, returning previously attached SAEs to their original state.
-            use_error_term (bool | None): If provided, will set the use_error_term attribute of all SAEs attached during this run to this value. Defaults to None.
+            saes: SAEs to be attached.
+            reset_saes_end: If True, removes all SAEs added by this context
+                manager when the context manager exits, returning previously
+                attached SAEs to their original state.
+            use_error_term: If provided, will set the use_error_term attribute
+                of all SAEs attached during this run to this value.
         """
-        act_names_to_reset = []
-        prev_saes = []
+        saes_to_restore: list[tuple[str, _SAEWrapper | None]] = []
         if isinstance(saes, SAE):
             saes = [saes]
         try:
             for sae in saes:
-                act_names_to_reset.append(sae.cfg.metadata.hook_name)
-                prev_sae = self.acts_to_saes.get(sae.cfg.metadata.hook_name, None)
-                prev_saes.append(prev_sae)
+                act_name = sae.cfg.metadata.hook_name
+                prev_wrapper = self._acts_to_saes.get(act_name, None)
+                saes_to_restore.append((act_name, prev_wrapper))
                 self.add_sae(sae, use_error_term=use_error_term)
             yield self
         finally:
             if reset_saes_end:
-                self.reset_saes(act_names_to_reset, prev_saes)
+                for act_name, prev_wrapper in saes_to_restore:
+                    self._reset_sae(act_name, prev_wrapper)
+                self.setup()

sae_lens/analysis/sae_transformer_bridge.py

@@ -8,7 +8,11 @@ from transformer_lens.hook_points import HookPoint
 from transformer_lens.model_bridge import TransformerBridge
 
 from sae_lens import logger
-from sae_lens.analysis.hooked_sae_transformer import set_deep_attr
+from sae_lens.analysis.hooked_sae_transformer import (
+    _SAEWrapper,
+    get_deep_attr,
+    set_deep_attr,
+)
 from sae_lens.saes.sae import SAE
 
 SingleLoss = torch.Tensor  # Type alias for a single element tensor
@@ -30,11 +34,18 @@ class SAETransformerBridge(TransformerBridge):  # type: ignore[misc,no-untyped-c
     useful for models not natively supported by HookedTransformer, such as Gemma 3.
     """
 
-    acts_to_saes: dict[str, SAE[Any]]
+    _acts_to_saes: dict[str, _SAEWrapper]
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.acts_to_saes = {}
+        self._acts_to_saes = {}
+        # Track output hooks used by transcoders for cleanup
+        self._transcoder_output_hooks: dict[str, str] = {}
+
+    @property
+    def acts_to_saes(self) -> dict[str, SAE[Any]]:
+        """Returns a dict mapping hook names to attached SAEs."""
+        return {name: wrapper.sae for name, wrapper in self._acts_to_saes.items()}
 
     @classmethod
     def boot_transformers(  # type: ignore[override]
@@ -56,7 +67,8 @@ class SAETransformerBridge(TransformerBridge):  # type: ignore[misc,no-untyped-c
         # Convert to our class
         # NOTE: this is super hacky and scary, but I don't know how else to achieve this given TLens' internal code
         bridge.__class__ = cls
-        bridge.acts_to_saes = {}  # type: ignore[attr-defined]
+        bridge._acts_to_saes = {}  # type: ignore[attr-defined]
+        bridge._transcoder_output_hooks = {}  # type: ignore[attr-defined]
         return bridge  # type: ignore[return-value]
 
     def _resolve_hook_name(self, hook_name: str) -> str:
@@ -75,110 +87,129 @@ class SAETransformerBridge(TransformerBridge):  # type: ignore[misc,no-untyped-c
         return resolved if isinstance(resolved, str) else hook_name
 
     def add_sae(self, sae: SAE[Any], use_error_term: bool | None = None) -> None:
-        """Attaches an SAE to the model.
+        """Attaches an SAE or Transcoder to the model.
 
         WARNING: This SAE will be permanently attached until you remove it with
         reset_saes. This function will also overwrite any existing SAE attached
         to the same hook point.
 
         Args:
-            sae: The SAE to attach to the model
-            use_error_term: If provided, will set the use_error_term attribute of
-                the SAE to this value. Determines whether the SAE returns input
-                or reconstruction. Defaults to None.
+            sae: The SAE or Transcoder to attach to the model.
+            use_error_term: If True, computes error term so output matches what the
+                model would have produced without the SAE. This works for both SAEs
+                (where input==output hook) and transcoders (where they differ).
+                Defaults to None (uses SAE's existing setting).
         """
-        alias_name = sae.cfg.metadata.hook_name
-        actual_name = self._resolve_hook_name(alias_name)
-
-        # Check if hook exists (either as alias or actual name)
-        if (alias_name not in self.acts_to_saes) and (
-            actual_name not in self._hook_registry
+        input_hook_alias = sae.cfg.metadata.hook_name
+        output_hook_alias = sae.cfg.metadata.hook_name_out or input_hook_alias
+        input_hook_actual = self._resolve_hook_name(input_hook_alias)
+        output_hook_actual = self._resolve_hook_name(output_hook_alias)
+
+        # Check if hooks exist
+        if (input_hook_alias not in self._acts_to_saes) and (
+            input_hook_actual not in self._hook_registry
         ):
             logger.warning(
-                f"No hook found for {alias_name}. Skipping. "
+                f"No hook found for {input_hook_alias}. Skipping. "
                 f"Check model._hook_registry for available hooks."
             )
             return
 
-        if use_error_term is not None:
-            if not hasattr(sae, "_original_use_error_term"):
-                sae._original_use_error_term = sae.use_error_term  # type: ignore[attr-defined]
-            sae.use_error_term = use_error_term
-
-        # Replace hook and update registry
-        set_deep_attr(self, actual_name, sae)
-        self._hook_registry[actual_name] = sae  # type: ignore[assignment]
-        self.acts_to_saes[alias_name] = sae
+        # Check if output hook exists (either as registry entry or already has SAE attached)
+        output_hook_exists = (
+            output_hook_actual in self._hook_registry
+            or input_hook_alias in self._acts_to_saes
+            or any(
+                v == output_hook_actual for v in self._transcoder_output_hooks.values()
+            )
+        )
+        if not output_hook_exists:
+            logger.warning(f"No hook found for output {output_hook_alias}. Skipping.")
+            return
 
-    def _reset_sae(self, act_name: str, prev_sae: SAE[Any] | None = None) -> None:
+        # Always use wrapper - it handles both SAEs and transcoders uniformly
+        # If use_error_term not specified, respect SAE's existing setting
+        effective_use_error_term = (
+            use_error_term if use_error_term is not None else sae.use_error_term
+        )
+        wrapper = _SAEWrapper(sae, use_error_term=effective_use_error_term)
+
+        # For transcoders (input != output), capture input at input hook
+        if input_hook_alias != output_hook_alias:
+            input_hook_point = get_deep_attr(self, input_hook_actual)
+            if isinstance(input_hook_point, HookPoint):
+                input_hook_point.add_hook(
+                    lambda tensor, hook: (wrapper.capture_input(tensor), tensor)[1],  # noqa: ARG005
+                    dir="fwd",
+                    is_permanent=True,
+                )
+            self._transcoder_output_hooks[input_hook_alias] = output_hook_actual
+
+        # Store wrapper in _acts_to_saes and at output hook
+        set_deep_attr(self, output_hook_actual, wrapper)
+        self._hook_registry[output_hook_actual] = wrapper  # type: ignore[assignment]
+        self._acts_to_saes[input_hook_alias] = wrapper
+
+    def _reset_sae(
+        self, act_name: str, prev_wrapper: _SAEWrapper | None = None
+    ) -> None:
         """Resets an SAE that was attached to the model.
 
         By default will remove the SAE from that hook_point.
-        If prev_sae is provided, will replace the current SAE with the provided one.
-        This is mainly used to restore previously attached SAEs after temporarily
-        running with different SAEs (e.g., with run_with_saes).
+        If prev_wrapper is provided, will restore that wrapper's SAE with its settings.
 
         Args:
             act_name: The hook_name of the SAE to reset
-            prev_sae: The SAE to replace the current one with. If None, will just
+            prev_wrapper: The previous wrapper to restore. If None, will just
                 remove the SAE from this hook point. Defaults to None.
         """
-        if act_name not in self.acts_to_saes:
+        if act_name not in self._acts_to_saes:
             logger.warning(
                 f"No SAE is attached to {act_name}. There's nothing to reset."
             )
             return
 
         actual_name = self._resolve_hook_name(act_name)
-        current_sae = self.acts_to_saes[act_name]
-
-        if hasattr(current_sae, "_original_use_error_term"):
-            current_sae.use_error_term = current_sae._original_use_error_term  # type: ignore[attr-defined]
-            delattr(current_sae, "_original_use_error_term")
-
-        if prev_sae is not None:
-            set_deep_attr(self, actual_name, prev_sae)
-            self._hook_registry[actual_name] = prev_sae  # type: ignore[assignment]
-            self.acts_to_saes[act_name] = prev_sae
-        else:
-            new_hook = HookPoint()
-            new_hook.name = actual_name
-            set_deep_attr(self, actual_name, new_hook)
-            self._hook_registry[actual_name] = new_hook
-            del self.acts_to_saes[act_name]
+
+        # Determine output hook location (different from input for transcoders)
+        output_hook = self._transcoder_output_hooks.pop(act_name, actual_name)
+
+        # For transcoders, clear permanent hooks from input hook point
+        if output_hook != actual_name:
+            input_hook_point = get_deep_attr(self, actual_name)
+            if isinstance(input_hook_point, HookPoint):
+                input_hook_point.remove_hooks(dir="fwd", including_permanent=True)
+
+        # Reset output hook location
+        new_hook = HookPoint()
+        new_hook.name = output_hook
+        set_deep_attr(self, output_hook, new_hook)
+        self._hook_registry[output_hook] = new_hook
+        del self._acts_to_saes[act_name]
+
+        if prev_wrapper is not None:
+            self.add_sae(prev_wrapper.sae, use_error_term=prev_wrapper.use_error_term)
 
     def reset_saes(
         self,
         act_names: str | list[str] | None = None,
-        prev_saes: list[SAE[Any] | None] | None = None,
     ) -> None:
         """Reset the SAEs attached to the model.
 
         If act_names are provided will just reset SAEs attached to those hooks.
         Otherwise will reset all SAEs attached to the model.
-        Optionally can provide a list of prev_saes to reset to. This is mainly
-        used to restore previously attached SAEs after temporarily running with
-        different SAEs (e.g., with run_with_saes).
 
         Args:
             act_names: The act_names of the SAEs to reset. If None, will reset all
                 SAEs attached to the model. Defaults to None.
-            prev_saes: List of SAEs to replace the current ones with. If None, will
-                just remove the SAEs. Defaults to None.
         """
         if isinstance(act_names, str):
             act_names = [act_names]
         elif act_names is None:
-            act_names = list(self.acts_to_saes.keys())
-
-        if prev_saes:
-            if len(act_names) != len(prev_saes):
-                raise ValueError("act_names and prev_saes must have the same length")
-        else:
-            prev_saes = [None] * len(act_names)  # type: ignore[assignment]
+            act_names = list(self._acts_to_saes.keys())
 
-        for act_name, prev_sae in zip(act_names, prev_saes):  # type: ignore[arg-type]
-            self._reset_sae(act_name, prev_sae)
+        for act_name in act_names:
+            self._reset_sae(act_name)
 
     def run_with_saes(
         self,
@@ -310,20 +341,20 @@ class SAETransformerBridge(TransformerBridge):  # type: ignore[misc,no-untyped-c
             use_error_term: If provided, will set the use_error_term attribute of
                 all SAEs attached during this run to this value. Defaults to None.
         """
-        act_names_to_reset: list[str] = []
-        prev_saes: list[SAE[Any] | None] = []
+        saes_to_restore: list[tuple[str, _SAEWrapper | None]] = []
         if isinstance(saes, SAE):
             saes = [saes]
         try:
             for sae in saes:
-                act_names_to_reset.append(sae.cfg.metadata.hook_name)
-                prev_sae = self.acts_to_saes.get(sae.cfg.metadata.hook_name, None)
-                prev_saes.append(prev_sae)
+                act_name = sae.cfg.metadata.hook_name
+                prev_wrapper = self._acts_to_saes.get(act_name, None)
+                saes_to_restore.append((act_name, prev_wrapper))
                 self.add_sae(sae, use_error_term=use_error_term)
             yield self
         finally:
             if reset_saes_end:
-                self.reset_saes(act_names_to_reset, prev_saes)
+                for act_name, prev_wrapper in saes_to_restore:
+                    self._reset_sae(act_name, prev_wrapper)
 
     @property
     def hook_dict(self) -> dict[str, HookPoint]:
@@ -337,9 +368,9 @@ class SAETransformerBridge(TransformerBridge):  # type: ignore[misc,no-untyped-c
         hooks: dict[str, HookPoint] = {}
 
         for name, hook_or_sae in self._hook_registry.items():
-            if isinstance(hook_or_sae, SAE):
+            if isinstance(hook_or_sae, _SAEWrapper):
                 # Include SAE's internal hooks with full path names
-                for sae_hook_name, sae_hook in hook_or_sae.hook_dict.items():
+                for sae_hook_name, sae_hook in hook_or_sae.sae.hook_dict.items():
                     full_name = f"{name}.{sae_hook_name}"
                     hooks[full_name] = sae_hook
             else:

sae_lens/loading/pretrained_saes_directory.py

@@ -57,28 +57,6 @@ def get_pretrained_saes_directory() -> dict[str, PretrainedSAELookup]:
     return directory
 
 
-def get_norm_scaling_factor(release: str, sae_id: str) -> float | None:
-    """
-    Retrieve the norm_scaling_factor for a specific SAE if it exists.
-
-    Args:
-        release (str): The release name of the SAE.
-        sae_id (str): The ID of the specific SAE.
-
-    Returns:
-        float | None: The norm_scaling_factor if it exists, None otherwise.
-    """
-    package = "sae_lens"
-    yaml_file = files(package).joinpath("pretrained_saes.yaml")
-    with yaml_file.open("r") as file:
-        data = yaml.safe_load(file)
-        if release in data:
-            for sae_info in data[release]["saes"]:
-                if sae_info["id"] == sae_id:
-                    return sae_info.get("norm_scaling_factor")
-    return None
-
-
 def get_repo_id_and_folder_name(release: str, sae_id: str) -> tuple[str, str]:
     saes_directory = get_pretrained_saes_directory()
     sae_info = saes_directory.get(release, None)

sae_lens/saes/sae.py

@@ -45,7 +45,6 @@ from sae_lens.loading.pretrained_sae_loaders import (
 )
 from sae_lens.loading.pretrained_saes_directory import (
     get_config_overrides,
-    get_norm_scaling_factor,
     get_pretrained_saes_directory,
     get_releases_for_repo_id,
     get_repo_id_and_folder_name,
@@ -230,7 +229,7 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
     cfg: T_SAE_CONFIG
     dtype: torch.dtype
     device: torch.device
-    use_error_term: bool
+    _use_error_term: bool
 
     # For type checking only - don't provide default values
     # These will be initialized by subclasses
@@ -255,7 +254,9 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
 
         self.dtype = str_to_dtype(cfg.dtype)
         self.device = torch.device(cfg.device)
-        self.use_error_term = use_error_term
+        self._use_error_term = False  # Set directly to avoid warning during init
+        if use_error_term:
+            self.use_error_term = True  # Use property setter to trigger warning
 
         # Set up activation function
         self.activation_fn = self.get_activation_fn()
@@ -282,6 +283,22 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
 
         self.setup()  # Required for HookedRootModule
 
+    @property
+    def use_error_term(self) -> bool:
+        return self._use_error_term
+
+    @use_error_term.setter
+    def use_error_term(self, value: bool) -> None:
+        if value and not self._use_error_term:
+            warnings.warn(
+                "Setting use_error_term directly on SAE is deprecated. "
+                "Use HookedSAETransformer.add_sae(sae, use_error_term=True) instead. "
+                "This will be removed in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        self._use_error_term = value
+
     @torch.no_grad()
     def fold_activation_norm_scaling_factor(self, scaling_factor: float):
         self.W_enc.data *= scaling_factor  # type: ignore
@@ -638,24 +655,6 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
                     stacklevel=2,
                 )
         elif sae_id not in sae_directory[release].saes_map:
-            # Handle special cases like Gemma Scope
-            if (
-                "gemma-scope" in release
-                and "canonical" not in release
-                and f"{release}-canonical" in sae_directory
-            ):
-                canonical_ids = list(
-                    sae_directory[release + "-canonical"].saes_map.keys()
-                )
-                # Shorten the lengthy string of valid IDs
-                if len(canonical_ids) > 5:
-                    str_canonical_ids = str(canonical_ids[:5])[:-1] + ", ...]"
-                else:
-                    str_canonical_ids = str(canonical_ids)
-                value_suffix = f" If you don't want to specify an L0 value, consider using release {release}-canonical which has valid IDs {str_canonical_ids}"
-            else:
-                value_suffix = ""
-
             valid_ids = list(sae_directory[release].saes_map.keys())
             # Shorten the lengthy string of valid IDs
             if len(valid_ids) > 5:
@@ -665,7 +664,6 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
 
             raise ValueError(
                 f"ID {sae_id} not found in release {release}. Valid IDs are {str_valid_ids}."
-                + value_suffix
             )
 
         conversion_loader = (
@@ -702,17 +700,6 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
         sae.process_state_dict_for_loading(state_dict)
         sae.load_state_dict(state_dict, assign=True)
 
-        # Apply normalization if needed
-        if cfg_dict.get("normalize_activations") == "expected_average_only_in":
-            norm_scaling_factor = get_norm_scaling_factor(release, sae_id)
-            if norm_scaling_factor is not None:
-                sae.fold_activation_norm_scaling_factor(norm_scaling_factor)
-                cfg_dict["normalize_activations"] = "none"
-            else:
-                warnings.warn(
-                    f"norm_scaling_factor not found for {release} and {sae_id}, but normalize_activations is 'expected_average_only_in'. Skipping normalization folding."
-                )
-
         # the loaders should already handle the dtype / device conversion
         # but this is a fallback to guarantee the SAE is on the correct device and dtype
         return (

sae_lens/synthetic/__init__.py

@@ -50,6 +50,13 @@ from sae_lens.synthetic.plotting import (
     find_best_feature_ordering_from_sae,
     plot_sae_feature_similarity,
 )
+from sae_lens.synthetic.stats import (
+    CorrelationMatrixStats,
+    SuperpositionStats,
+    compute_correlation_matrix_stats,
+    compute_low_rank_correlation_matrix_stats,
+    compute_superposition_stats,
+)
 from sae_lens.synthetic.training import (
     SyntheticActivationIterator,
     train_toy_sae,
@@ -80,6 +87,12 @@ __all__ = [
     "orthogonal_initializer",
     "FeatureDictionaryInitializer",
     "cosine_similarities",
+    # Statistics
+    "compute_correlation_matrix_stats",
+    "compute_low_rank_correlation_matrix_stats",
+    "compute_superposition_stats",
+    "CorrelationMatrixStats",
+    "SuperpositionStats",
     # Training utilities
     "SyntheticActivationIterator",
     "SyntheticDataEvalResult",

sae_lens/synthetic/correlation.py

@@ -3,6 +3,7 @@ from typing import NamedTuple
 
 import torch
 
+from sae_lens import logger
 from sae_lens.util import str_to_dtype
 
 
@@ -268,7 +269,7 @@ def generate_random_correlation_matrix(
 def generate_random_low_rank_correlation_matrix(
     num_features: int,
     rank: int,
-    correlation_scale: float = 0.1,
+    correlation_scale: float = 0.075,
     seed: int | None = None,
     device: torch.device | str = "cpu",
     dtype: torch.dtype | str = torch.float32,
@@ -331,20 +332,17 @@ def generate_random_low_rank_correlation_matrix(
     factor_sq_sum = (factor**2).sum(dim=1)
     diag_term = 1 - factor_sq_sum
 
-    # Ensure diagonal terms are at least _MIN_DIAG for numerical stability
-    # If any diagonal term is too small, scale down the factor matrix
-    if torch.any(diag_term < _MIN_DIAG):
-        # Scale factor so max row norm squared is at most (1 - _MIN_DIAG)
-        # This ensures all diagonal terms are >= _MIN_DIAG
-        max_factor_contribution = 1 - _MIN_DIAG
-        max_sq_sum = factor_sq_sum.max()
-        scale = torch.sqrt(
-            torch.tensor(max_factor_contribution, device=device, dtype=dtype)
-            / max_sq_sum
+    # alternatively, we can rescale each row independently to ensure the diagonal is 1
+    mask = diag_term < _MIN_DIAG
+    factor[mask, :] *= torch.sqrt((1 - _MIN_DIAG) / factor_sq_sum[mask].unsqueeze(1))
+    factor_sq_sum = (factor**2).sum(dim=1)
+    diag_term = 1 - factor_sq_sum
+
+    total_rescaled = mask.sum().item()
+    if total_rescaled > 0:
+        logger.warning(
+            f"{total_rescaled} / {num_features} rows were capped. Either reduce the rank or reduce the correlation_scale to avoid rescaling."
         )
-        factor = factor * scale
-        factor_sq_sum = (factor**2).sum(dim=1)
-        diag_term = 1 - factor_sq_sum
 
     return LowRankCorrelationMatrix(
         correlation_factor=factor, correlation_diag=diag_term

sae_lens/synthetic/stats.py

@@ -0,0 +1,205 @@
+from dataclasses import dataclass
+
+import torch
+
+from sae_lens.synthetic.correlation import LowRankCorrelationMatrix
+from sae_lens.synthetic.feature_dictionary import FeatureDictionary
+
+
+@dataclass
+class CorrelationMatrixStats:
+    """Statistics computed from a correlation matrix."""
+
+    rms_correlation: float  # Root mean square of off-diagonal correlations
+    mean_correlation: float  # Mean of off-diagonal correlations (not absolute)
+    num_features: int
+
+
+@torch.no_grad()
+def compute_correlation_matrix_stats(
+    correlation_matrix: torch.Tensor,
+) -> CorrelationMatrixStats:
+    """Compute correlation statistics from a dense correlation matrix.
+
+    Args:
+        correlation_matrix: Dense correlation matrix of shape (n, n)
+
+    Returns:
+        CorrelationMatrixStats with correlation statistics
+    """
+    num_features = correlation_matrix.shape[0]
+
+    # Extract off-diagonal elements
+    mask = ~torch.eye(num_features, dtype=torch.bool, device=correlation_matrix.device)
+    off_diag = correlation_matrix[mask]
+
+    rms_correlation = (off_diag**2).mean().sqrt().item()
+    mean_correlation = off_diag.mean().item()
+
+    return CorrelationMatrixStats(
+        rms_correlation=rms_correlation,
+        mean_correlation=mean_correlation,
+        num_features=num_features,
+    )
+
+
+@torch.no_grad()
+def compute_low_rank_correlation_matrix_stats(
+    correlation_matrix: LowRankCorrelationMatrix,
+) -> CorrelationMatrixStats:
+    """Compute correlation statistics from a LowRankCorrelationMatrix.
+
+    The correlation matrix is represented as:
+        correlation = factor @ factor.T + diag(diag_term)
+
+    The off-diagonal elements are simply factor @ factor.T (the diagonal term
+    only affects the diagonal).
+
+    All statistics are computed efficiently in O(n*r²) time and O(r²) memory
+    without materializing the full n×n correlation matrix.
+
+    Args:
+        correlation_matrix: Low-rank correlation matrix
+
+    Returns:
+        CorrelationMatrixStats with correlation statistics
+    """
+
+    factor = correlation_matrix.correlation_factor
+    num_features = factor.shape[0]
+    num_off_diag = num_features * (num_features - 1)
+
+    # RMS correlation: uses ||F @ F.T||_F² = ||F.T @ F||_F²
+    # This avoids computing the (num_features, num_features) matrix
+    G = factor.T @ factor  # (rank, rank) - small!
+    frobenius_sq = (G**2).sum()
+    row_norms_sq = (factor**2).sum(dim=1)  # ||F[i]||² for each row
+    diag_sq_sum = (row_norms_sq**2).sum()  # Σᵢ ||F[i]||⁴
+    off_diag_sq_sum = frobenius_sq - diag_sq_sum
+    rms_correlation = (off_diag_sq_sum / num_off_diag).sqrt().item()
+
+    # Mean correlation (not absolute): sum(C) = ||col_sums(F)||², trace(C) = Σ||F[i]||²
+    col_sums = factor.sum(dim=0)  # (rank,)
+    sum_all = (col_sums**2).sum()  # 1ᵀ C 1
+    trace_C = row_norms_sq.sum()
+    mean_correlation = ((sum_all - trace_C) / num_off_diag).item()
+
+    return CorrelationMatrixStats(
+        rms_correlation=rms_correlation,
+        mean_correlation=mean_correlation,
+        num_features=num_features,
+    )
+
+
+@dataclass
+class SuperpositionStats:
+    """Statistics measuring superposition in a feature dictionary."""
+
+    # Per-latent statistics: for each latent, max and percentile of |cos_sim| with others
+    max_abs_cos_sims: torch.Tensor  # Shape: (num_features,)
+    percentile_abs_cos_sims: dict[int, torch.Tensor]  # {percentile: (num_features,)}
+
+    # Summary statistics (means of the per-latent values)
+    mean_max_abs_cos_sim: float
+    mean_percentile_abs_cos_sim: dict[int, float]
+    mean_abs_cos_sim: float  # Mean |cos_sim| across all pairs
+
+    # Metadata
+    num_features: int
+    hidden_dim: int
+
+
+@torch.no_grad()
+def compute_superposition_stats(
+    feature_dictionary: FeatureDictionary,
+    batch_size: int = 1024,
+    device: str | torch.device | None = None,
+    percentiles: list[int] | None = None,
+) -> SuperpositionStats:
+    """Compute superposition statistics for a feature dictionary.
+
+    Computes pairwise cosine similarities in batches to handle large dictionaries.
+
+    For each latent i, computes:
+
+    - max |cos_sim(i, j)| over all j != i
+    - kth percentile of |cos_sim(i, j)| over all j != i (for each k in percentiles)
+
+    Args:
+        feature_dictionary: FeatureDictionary containing the feature vectors
+        batch_size: Number of features to process per batch
+        device: Device for computation (defaults to feature dictionary's device)
+        percentiles: List of percentiles to compute per latent (default: [95, 99])
+
+    Returns:
+        SuperpositionStats with superposition metrics
+    """
+    if percentiles is None:
+        percentiles = [95, 99]
+
+    feature_vectors = feature_dictionary.feature_vectors
+    num_features, hidden_dim = feature_vectors.shape
+
+    if num_features < 2:
+        raise ValueError("Need at least 2 features to compute superposition stats")
+    if device is None:
+        device = feature_vectors.device
+
+    # Normalize features to unit norm for cosine similarity
+    features_normalized = feature_vectors.to(device).float()
+    norms = torch.linalg.norm(features_normalized, dim=1, keepdim=True)
+    features_normalized = features_normalized / norms.clamp(min=1e-8)
+
+    # Track per-latent statistics
+    max_abs_cos_sims = torch.zeros(num_features, device=device)
+    percentile_abs_cos_sims = {
+        p: torch.zeros(num_features, device=device) for p in percentiles
+    }
+    sum_abs_cos_sim = 0.0
+    n_pairs = 0
+
+    # Process in batches: for each batch of features, compute similarities with all others
+    for i in range(0, num_features, batch_size):
+        batch_end = min(i + batch_size, num_features)
+        batch = features_normalized[i:batch_end]  # (batch_size, hidden_dim)
+
+        # Compute cosine similarities with all features: (batch_size, num_features)
+        cos_sims = batch @ features_normalized.T
+
+        # Absolute cosine similarities
+        abs_cos_sims = cos_sims.abs()
+
+        # Process each latent in the batch
+        for j, idx in enumerate(range(i, batch_end)):
+            # Get similarities with all other features (exclude self)
+            row = abs_cos_sims[j].clone()
+            row[idx] = 0.0  # Exclude self for max
+            max_abs_cos_sims[idx] = row.max()
+
+            # For percentiles, exclude self and compute
+            other_sims = torch.cat([abs_cos_sims[j, :idx], abs_cos_sims[j, idx + 1 :]])
+            for p in percentiles:
+                percentile_abs_cos_sims[p][idx] = torch.quantile(other_sims, p / 100.0)
+
+            # Sum for mean computation (only count pairs once - with features after this one)
+            sum_abs_cos_sim += abs_cos_sims[j, idx + 1 :].sum().item()
+            n_pairs += num_features - idx - 1
+
+    # Compute summary statistics
+    mean_max_abs_cos_sim = max_abs_cos_sims.mean().item()
+    mean_percentile_abs_cos_sim = {
+        p: percentile_abs_cos_sims[p].mean().item() for p in percentiles
+    }
+    mean_abs_cos_sim = sum_abs_cos_sim / n_pairs if n_pairs > 0 else 0.0
+
+    return SuperpositionStats(
+        max_abs_cos_sims=max_abs_cos_sims.cpu(),
+        percentile_abs_cos_sims={
+            p: v.cpu() for p, v in percentile_abs_cos_sims.items()
+        },
+        mean_max_abs_cos_sim=mean_max_abs_cos_sim,
+        mean_percentile_abs_cos_sim=mean_percentile_abs_cos_sim,
+        mean_abs_cos_sim=mean_abs_cos_sim,
+        num_features=num_features,
+        hidden_dim=hidden_dim,
+    )

{sae_lens-6.32.1.dist-info → sae_lens-6.34.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.32.1
+Version: 6.34.1
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE

{sae_lens-6.32.1.dist-info → sae_lens-6.34.1.dist-info}/RECORD

@@ -1,9 +1,9 @@
-sae_lens/__init__.py,sha256=Y_TVKGehpnTvQw8tvIn0fjo8uAw-XAYi7carZS_cRjQ,5168
+sae_lens/__init__.py,sha256=U5I6dWw2_x6ZXwL8BF72vB0HGPEUt-_BzHq-btLq2pw,5168
 sae_lens/analysis/__init__.py,sha256=FZExlMviNwWR7OGUSGRbd0l-yUDGSp80gglI_ivILrY,412
 sae_lens/analysis/compat.py,sha256=cgE3nhFcJTcuhppxbL71VanJS7YqVEOefuneB5eOaPw,538
-sae_lens/analysis/hooked_sae_transformer.py,sha256=LpnjxSAcItqqXA4SJyZuxY4Ki0UOuWV683wg9laYAsY,14050
+sae_lens/analysis/hooked_sae_transformer.py,sha256=-LY9CKYEziSTt-H7MeLdTx6ErfvMqgNEF709wV7tEs4,17826
 sae_lens/analysis/neuronpedia_integration.py,sha256=Gx1W7hUBEuMoasNcnOnZ1wmqbXDd1pSZ1nqKEya1HQc,4962
-sae_lens/analysis/sae_transformer_bridge.py,sha256=xpJRRcB0g47EOQcmNCwMyrJJsbqMsGxVViDrV6C3upU,14916
+sae_lens/analysis/sae_transformer_bridge.py,sha256=y_ZdvaxuUM_-7ywSCeFl6f6cq1_FiqRstMAo8mxZtYE,16191
 sae_lens/cache_activations_runner.py,sha256=TjqNWIc46Nw09jHWFjzQzgzG5wdu_87Ahe-iFjI5_0Q,13117
 sae_lens/config.py,sha256=V0BXV8rvpbm5YuVukow9FURPpdyE4HSflbdymAo0Ycg,31205
 sae_lens/constants.py,sha256=CM-h9AjZNAl2aP7hVpKk7YsFHpu-_Lfhhmq2d5qPEVc,887
@@ -12,7 +12,7 @@ sae_lens/llm_sae_training_runner.py,sha256=M7BK55gSFYu2qFQKABHX3c8i46P1LfODCeyHF
 sae_lens/load_model.py,sha256=C8AMykctj6H7tz_xRwB06-EXj6TfW64PtSJZR5Jxn1Y,8649
 sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/loading/pretrained_sae_loaders.py,sha256=kshvA0NivOc7B3sL19lHr_zrC_DDfW2T6YWb5j0hgAk,63930
-sae_lens/loading/pretrained_saes_directory.py,sha256=1at_aQbD8WFywchQCKuwfP-yvCq_Z2aUYrpKDnSN5Nc,4283
+sae_lens/loading/pretrained_saes_directory.py,sha256=lSnHl77IO5dd7iO21ynCzZNMrzuJAT8Za4W5THNq0qw,3554
 sae_lens/pretokenize_runner.py,sha256=amJwIz3CKi2s2wNQn-10E7eAV7VFhNqtFDNTeTkwEI8,7133
 sae_lens/pretrained_saes.yaml,sha256=IVBLLR8_XNllJ1O-kVv9ED4u0u44Yn8UOL9R-f8Idp4,1511936
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
@@ -22,20 +22,21 @@ sae_lens/saes/gated_sae.py,sha256=V_2ZNlV4gRD-rX5JSx1xqY7idT8ChfdQ5yxWDdu_6hg,88
 sae_lens/saes/jumprelu_sae.py,sha256=miiF-xI_yXdV9EkKjwAbU9zSMsx9KtKCz5YdXEzkN8g,13313
 sae_lens/saes/matching_pursuit_sae.py,sha256=08_G9p1YMLnE5qZVCPp6gll-iG6nHRbMMASf4_bkFt8,13207
 sae_lens/saes/matryoshka_batchtopk_sae.py,sha256=Qr6htt1HHOuO9FXI9hyaPSnGFIiJG-v7y1t1CEmkFzM,5995
-sae_lens/saes/sae.py,sha256=xRmgiLuaFlDCv8SyLbL-5TwdrWHpNLqSGe8mC1L6WcI,40942
+sae_lens/saes/sae.py,sha256=2FHhLoTZYOJUGjkWH7eF2EAc1fPxQsF4KYr5TZ8IUIU,40155
 sae_lens/saes/standard_sae.py,sha256=_hldNZkFPAf9VGrxouR1-tN8T2OEk8IkWBcXoatrC1o,5749
 sae_lens/saes/temporal_sae.py,sha256=S44sPddVj2xujA02CC8gT1tG0in7c_CSAhspu9FHbaA,13273
 sae_lens/saes/topk_sae.py,sha256=vrMRPrCQR1o8G_kXqY_EAoGZARupkQNFB2dNZVLsusE,21073
 sae_lens/saes/transcoder.py,sha256=CTpJs8ASOK06npih7gZHygZuxqTR7HICWlOYfTiKjI4,13501
-sae_lens/synthetic/__init__.py,sha256=MtTnGkTfHV2WjkIgs7zZyx10EK9U5fjOHXy69Aq3uKw,3095
+sae_lens/synthetic/__init__.py,sha256=hRRA3xhEQUacGyFbJXkLVYg_8A1bbSYYWlVovb0g4KU,3503
 sae_lens/synthetic/activation_generator.py,sha256=8L9nwC4jFRv_wg3QN-n1sFwX8w1NqwJMysWaJ41lLlY,15197
-sae_lens/synthetic/correlation.py,sha256=tMTLo9fBfDpeXwqhyUgFqnTipj9x2W0t4oEtNxB7AG0,13256
+sae_lens/synthetic/correlation.py,sha256=tD8J9abWfuFtGZrEbbFn4P8FeTcNKF2V5JhBLwDUmkg,13146
 sae_lens/synthetic/evals.py,sha256=Nhi314ZnRgLfhBj-3tm_zzI-pGyFTcwllDXbIpPFXeU,4584
 sae_lens/synthetic/feature_dictionary.py,sha256=Nd4xjSTxKMnKilZ3uYi8Gv5SS5D4bv4wHiSL1uGB69E,6933
 sae_lens/synthetic/firing_probabilities.py,sha256=yclz1pWl5gE1r8LAxFvzQS88Lxwk5-3r8BCX9HLVejA,3370
 sae_lens/synthetic/hierarchy.py,sha256=nm7nwnTswktVJeKUsRZ0hLOdXcFWGbxnA1b6lefHm-4,33592
 sae_lens/synthetic/initialization.py,sha256=orMGW-786wRDHIS2W7bEH0HmlVFQ4g2z4bnnwdv5w4s,1386
 sae_lens/synthetic/plotting.py,sha256=5lFrej1QOkGAcImFNo5-o-8mI_rUVqvEI57KzUQPPtQ,8208
+sae_lens/synthetic/stats.py,sha256=BoDPKDx8pgFF5Ko_IaBRZTczm7-ANUIRjjF5W5Qh3Lk,7441
 sae_lens/synthetic/training.py,sha256=fHcX2cZ6nDupr71GX0Gk17f1NvQ0SKIVXIA6IuAb2dw,5692
 sae_lens/tokenization_and_batching.py,sha256=uoHtAs9z3XqG0Fh-iQVYVlrbyB_E3kFFhrKU30BosCo,5438
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -48,7 +49,7 @@ sae_lens/training/types.py,sha256=1FpLx_Doda9vZpmfm-x1e8wGBYpyhe9Kpb_JuM5nIFM,90
 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=oIMoeyEP2IzcPFmRbKUzOAycgEyMcOasGeO_BGVZbc4,4846
-sae_lens-6.32.1.dist-info/METADATA,sha256=TcO6hFEXKdbLp32UTiVluHcMXFetfYJDqTHNCsx9PRw,6566
-sae_lens-6.32.1.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
-sae_lens-6.32.1.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.32.1.dist-info/RECORD,,
+sae_lens-6.34.1.dist-info/METADATA,sha256=vAlbWT90NggKoxjII6dnEj7wlC-PmuBn8zzqL6r8dRg,6566
+sae_lens-6.34.1.dist-info/WHEEL,sha256=kJCRJT_g0adfAJzTx2GUMmS80rTJIVHRCfG0DQgLq3o,88
+sae_lens-6.34.1.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.34.1.dist-info/RECORD,,

{sae_lens-6.32.1.dist-info → sae_lens-6.34.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 2.3.0
+Generator: poetry-core 2.3.1
 Root-Is-Purelib: true
 Tag: py3-none-any