sae-lens 6.31.0__tar.gz → 6.34.0__tar.gz

{sae_lens-6.31.0 → sae_lens-6.34.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.31.0
+Version: 6.34.0
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE
@@ -27,7 +27,7 @@ Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
 Requires-Dist: safetensors (>=0.4.2,<1.0.0)
 Requires-Dist: simple-parsing (>=0.1.6,<0.2.0)
 Requires-Dist: tenacity (>=9.0.0)
-Requires-Dist: transformer-lens (>=2.16.1,<3.0.0)
+Requires-Dist: transformer-lens (>=2.16.1)
 Requires-Dist: transformers (>=4.38.1,<5.0.0)
 Requires-Dist: typing-extensions (>=4.10.0,<5.0.0)
 Project-URL: Homepage, https://decoderesearch.github.io/SAELens

{sae_lens-6.31.0 → sae_lens-6.34.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.31.0"
+version = "6.34.0"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"
@@ -19,7 +19,7 @@ classifiers = ["Topic :: Scientific/Engineering :: Artificial Intelligence"]
 
 [tool.poetry.dependencies]
 python = "^3.10"
-transformer-lens = "^2.16.1"
+transformer-lens = ">=2.16.1"
 transformers = "^4.38.1"
 plotly = ">=5.19.0"
 plotly-express = ">=0.4.1"
@@ -59,6 +59,7 @@ mike = "^2.0.0"
 trio = "^0.30.0"
 dictionary-learning = "^0.1.0"
 kaleido = "^1.2.0"
+transformer-lens = { version = "3.0.0b1", allow-prereleases = true }
 
 [tool.poetry.extras]
 mamba = ["mamba-lens"]

{sae_lens-6.31.0 → sae_lens-6.34.0}/sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.31.0"
+__version__ = "6.34.0"
 
 import logging
 
@@ -125,6 +125,19 @@ __all__ = [
     "MatchingPursuitTrainingSAEConfig",
 ]
 
+# Conditional export for SAETransformerBridge (requires transformer-lens v3+)
+try:
+    from sae_lens.analysis.compat import has_transformer_bridge
+
+    if has_transformer_bridge():
+        from sae_lens.analysis.sae_transformer_bridge import (  # noqa: F401
+            SAETransformerBridge,
+        )
+
+        __all__.append("SAETransformerBridge")
+except ImportError:
+    pass
+
 
 register_sae_class("standard", StandardSAE, StandardSAEConfig)
 register_sae_training_class("standard", StandardTrainingSAE, StandardTrainingSAEConfig)

sae_lens-6.34.0/sae_lens/analysis/__init__.py

@@ -0,0 +1,15 @@
+from sae_lens.analysis.hooked_sae_transformer import HookedSAETransformer
+
+__all__ = ["HookedSAETransformer"]
+
+try:
+    from sae_lens.analysis.compat import has_transformer_bridge
+
+    if has_transformer_bridge():
+        from sae_lens.analysis.sae_transformer_bridge import (  # noqa: F401
+            SAETransformerBridge,
+        )
+
+        __all__.append("SAETransformerBridge")
+except ImportError:
+    pass

sae_lens-6.34.0/sae_lens/analysis/compat.py

@@ -0,0 +1,16 @@
+import importlib.metadata
+
+from packaging.version import parse as parse_version
+
+
+def get_transformer_lens_version() -> tuple[int, int, int]:
+    """Get transformer-lens version as (major, minor, patch)."""
+    version_str = importlib.metadata.version("transformer-lens")
+    version = parse_version(version_str)
+    return (version.major, version.minor, version.micro)
+
+
+def has_transformer_bridge() -> bool:
+    """Check if TransformerBridge is available (v3+)."""
+    major, _, _ = get_transformer_lens_version()
+    return major >= 3

{sae_lens-6.31.0 → sae_lens-6.34.0}/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,71 @@ 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
+        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 +144,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:
-            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 +377,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()