sae-lens 6.22.3__py3-none-any.whl → 6.24.0__py3-none-any.whl

sae_lens/saes/__init__.py

@@ -33,6 +33,8 @@ from .topk_sae import (
     TopKTrainingSAEConfig,
 )
 from .transcoder import (
+    JumpReLUSkipTranscoder,
+    JumpReLUSkipTranscoderConfig,
     JumpReLUTranscoder,
     JumpReLUTranscoderConfig,
     SkipTranscoder,
@@ -70,6 +72,8 @@ __all__ = [
     "SkipTranscoderConfig",
     "JumpReLUTranscoder",
     "JumpReLUTranscoderConfig",
+    "JumpReLUSkipTranscoder",
+    "JumpReLUSkipTranscoderConfig",
     "MatryoshkaBatchTopKTrainingSAE",
     "MatryoshkaBatchTopKTrainingSAEConfig",
     "TemporalSAE",

sae_lens/saes/transcoder.py

@@ -368,3 +368,44 @@ class JumpReLUTranscoder(Transcoder):
     def from_dict(cls, config_dict: dict[str, Any]) -> "JumpReLUTranscoder":
         cfg = JumpReLUTranscoderConfig.from_dict(config_dict)
         return cls(cfg)
+
+
+@dataclass
+class JumpReLUSkipTranscoderConfig(JumpReLUTranscoderConfig):
+    """Configuration for JumpReLU transcoder."""
+
+    @classmethod
+    def architecture(cls) -> str:
+        """Return the architecture name for this config."""
+        return "jumprelu_skip_transcoder"
+
+    @classmethod
+    def from_dict(cls, config_dict: dict[str, Any]) -> "JumpReLUSkipTranscoderConfig":
+        """Create a JumpReLUSkipTranscoderConfig from a dictionary."""
+        # Filter to only include valid dataclass fields
+        filtered_config_dict = filter_valid_dataclass_fields(config_dict, cls)
+
+        # Create the config instance
+        res = cls(**filtered_config_dict)
+
+        # Handle metadata if present
+        if "metadata" in config_dict:
+            res.metadata = SAEMetadata(**config_dict["metadata"])
+
+        return res
+
+
+class JumpReLUSkipTranscoder(JumpReLUTranscoder, SkipTranscoder):
+    """
+    A transcoder with a learnable skip connection and JumpReLU activation function.
+    """
+
+    cfg: JumpReLUSkipTranscoderConfig  # type: ignore[assignment]
+
+    def __init__(self, cfg: JumpReLUSkipTranscoderConfig):
+        super().__init__(cfg)
+
+    @classmethod
+    def from_dict(cls, config_dict: dict[str, Any]) -> "JumpReLUSkipTranscoder":
+        cfg = JumpReLUSkipTranscoderConfig.from_dict(config_dict)
+        return cls(cfg)

sae_lens/training/activations_store.py

@@ -230,7 +230,7 @@ class ActivationsStore:
             load_dataset(
                 dataset,
                 split="train",
-                streaming=streaming,
+                streaming=streaming,  # type: ignore
                 trust_remote_code=dataset_trust_remote_code,  # type: ignore
             )
             if isinstance(dataset, str)

{sae_lens-6.22.3.dist-info → sae_lens-6.24.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.22.3
+Version: 6.24.0
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE
@@ -62,7 +62,7 @@ This library is maintained by [Joseph Bloom](https://www.decoderesearch.com/), [
 
 ## Loading Pre-trained SAEs.
 
-Pre-trained SAEs for various models can be imported via SAE Lens. See this [page](https://decoderesearch.github.io/SAELens/sae_table/) in the readme for a list of all SAEs.
+Pre-trained SAEs for various models can be imported via SAE Lens. See this [page](https://decoderesearch.github.io/SAELens/pretrained_saes/) for a list of all SAEs.
 
 ## Migrating to SAELens v6
 

{sae_lens-6.22.3.dist-info → sae_lens-6.24.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-sae_lens/__init__.py,sha256=-q6U-a-hKOgMuJj8jBEyV5JvAZolVh6ccirbA3zrpYg,4033
+sae_lens/__init__.py,sha256=nBhszyLHRXNVsbQ-Y3FfIZ-DB_0C8MOANVXe_1eTRtI,4268
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=dQRgGVwce8XwylL2AzJE7l9elhtMRFCs2hdUj-Qyy4g,14038
 sae_lens/analysis/neuronpedia_integration.py,sha256=Gx1W7hUBEuMoasNcnOnZ1wmqbXDd1pSZ1nqKEya1HQc,4962
@@ -9,12 +9,12 @@ sae_lens/evals.py,sha256=P0NUsJeGzYxFBiVKhbPzd72IFKY4gH40HHlEZ3jEAmg,39598
 sae_lens/llm_sae_training_runner.py,sha256=M7BK55gSFYu2qFQKABHX3c8i46P1LfODCeyHFzGGuqU,15196
 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=X-gVZ4A74E85lSMFMsZ_rEQhHlR9AYFwhxvoA_vt2CQ,56051
+sae_lens/loading/pretrained_sae_loaders.py,sha256=W2eIvUU1wAHrYxGiZs4s2D6DnGBQqqKjq0wvXzWbD5c,63561
 sae_lens/loading/pretrained_saes_directory.py,sha256=hejNfLUepYCSGPalRfQwxxCEUqMMUPsn1tufwvwct5k,3820
-sae_lens/pretokenize_runner.py,sha256=x-reJzVPFDS9iRFbZtrFYSzNguJYki9gd0pbHjYJ3r4,7085
-sae_lens/pretrained_saes.yaml,sha256=VzgJ_t-IEWpO2MabgQY6CAcg8FFsqZWiOVXjqvqfgeE,604973
+sae_lens/pretokenize_runner.py,sha256=amJwIz3CKi2s2wNQn-10E7eAV7VFhNqtFDNTeTkwEI8,7133
+sae_lens/pretrained_saes.yaml,sha256=si8wRWgFhsCmgaigzye8Nsmyg03nDo0DiczYMisVlgY,1509026
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
-sae_lens/saes/__init__.py,sha256=nTNPnJ7edyfedo1MX96xwn9WOG8504yHbT9LFw9od_0,1778
+sae_lens/saes/__init__.py,sha256=fYVujOzNnUgpzLL0MBLBt_DNX2CPcTaheukzCd2bEPo,1906
 sae_lens/saes/batchtopk_sae.py,sha256=x4EbgZl0GUickRPcCmtKNGS2Ra3Uy1Z1OtF2FnrSabQ,5422
 sae_lens/saes/gated_sae.py,sha256=mHnmw-RD7hqIbP9_EBj3p2SK0OqQIkZivdOKRygeRgw,8825
 sae_lens/saes/jumprelu_sae.py,sha256=udjGHp3WTABQSL2Qq57j-bINWX61GCmo68EmdjMOXoo,13310
@@ -23,11 +23,11 @@ sae_lens/saes/sae.py,sha256=Vb1aGSDPRv_0J2aL8-EICRSkIxsO6Q4lJaJE9NNmfdA,37749
 sae_lens/saes/standard_sae.py,sha256=nEVETwAmRD2tyX7ESIic1fij48gAq1Dh7s_GQ2fqCZ4,5747
 sae_lens/saes/temporal_sae.py,sha256=DsecivcHWId-MTuJpQbz8OhqtmGhZACxJauYZGHo0Ok,13272
 sae_lens/saes/topk_sae.py,sha256=vrMRPrCQR1o8G_kXqY_EAoGZARupkQNFB2dNZVLsusE,21073
-sae_lens/saes/transcoder.py,sha256=BfLSbTYVNZh-ruGxseZiZJ_acEL6_7QyTdfqUr0lDOg,12156
+sae_lens/saes/transcoder.py,sha256=CTpJs8ASOK06npih7gZHygZuxqTR7HICWlOYfTiKjI4,13501
 sae_lens/tokenization_and_batching.py,sha256=D_o7cXvRqhT89H3wNzoRymNALNE6eHojBWLdXOUwUGE,5438
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/training/activation_scaler.py,sha256=FzNfgBplLWmyiSlZ6TUvE-nur3lOiGTrlvC97ys8S24,1973
-sae_lens/training/activations_store.py,sha256=pG6YhPCLrf7bd7bZTZ3_aS5J92lfnhVtyK8e0bzYFnI,33752
+sae_lens/training/activations_store.py,sha256=yDWw7TZGPFM_O8_Oi78j8lLIHJJesxq9TKVP_TrMX-M,33768
 sae_lens/training/mixing_buffer.py,sha256=vDpYG5ZE70szDvBsRKcNHEES3h_WTKJ16qDYk5jPOVA,2015
 sae_lens/training/optim.py,sha256=bJpqqcK4enkcPvQAJkeH4Ci1LUOlfjIMTv6-IlaAbRA,5588
 sae_lens/training/sae_trainer.py,sha256=zhkabyIKxI_tZTV3_kwz6zMrHZ95Ecr97krmwc-9ffs,17600
@@ -35,7 +35,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=tCovQ-eZa1L7thPpNDL6PGOJrIMML2yLI5e0EHCOpS8,3309
-sae_lens-6.22.3.dist-info/METADATA,sha256=vsYZOnH0fqH5fsQHPFnwXRLd61-NzOWm1Hz98YRS52c,5369
-sae_lens-6.22.3.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-sae_lens-6.22.3.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.22.3.dist-info/RECORD,,
+sae_lens-6.24.0.dist-info/METADATA,sha256=it9DN0UITpje9mAJG5Gt7jsO3Bt_MgbnrE5HDb9lBmA,5361
+sae_lens-6.24.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+sae_lens-6.24.0.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.24.0.dist-info/RECORD,,