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

sae_lens/__init__.py

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

sae_lens/loading/pretrained_sae_loaders.py

@@ -753,10 +753,14 @@ def get_dictionary_learning_config_1_from_hf(
     activation_fn = "topk" if trainer["dict_class"] == "AutoEncoderTopK" else "relu"
     activation_fn_kwargs = {"k": trainer["k"]} if activation_fn == "topk" else {}
 
+    architecture = "standard"
+    if trainer["dict_class"] == "GatedAutoEncoder":
+        architecture = "gated"
+    elif trainer["dict_class"] == "MatryoshkaBatchTopKSAE":
+        architecture = "jumprelu"
+
     return {
-        "architecture": (
-            "gated" if trainer["dict_class"] == "GatedAutoEncoder" else "standard"
-        ),
+        "architecture": architecture,
         "d_in": trainer["activation_dim"],
         "d_sae": trainer["dict_size"],
         "dtype": "float32",
@@ -905,9 +909,12 @@ def dictionary_learning_sae_huggingface_loader_1(
     )
     encoder = torch.load(encoder_path, map_location="cpu")
 
+    W_enc = encoder["W_enc"] if "W_enc" in encoder else encoder["encoder.weight"].T
+    W_dec = encoder["W_dec"] if "W_dec" in encoder else encoder["decoder.weight"].T
+
     state_dict = {
-        "W_enc": encoder["encoder.weight"].T,
-        "W_dec": encoder["decoder.weight"].T,
+        "W_enc": W_enc,
+        "W_dec": W_dec,
         "b_dec": encoder.get(
             "b_dec", encoder.get("bias", encoder.get("decoder_bias", None))
         ),
@@ -915,6 +922,8 @@ def dictionary_learning_sae_huggingface_loader_1(
 
     if "encoder.bias" in encoder:
         state_dict["b_enc"] = encoder["encoder.bias"]
+    if "b_enc" in encoder:
+        state_dict["b_enc"] = encoder["b_enc"]
 
     if "mag_bias" in encoder:
         state_dict["b_mag"] = encoder["mag_bias"]
@@ -923,6 +932,12 @@ def dictionary_learning_sae_huggingface_loader_1(
     if "r_mag" in encoder:
         state_dict["r_mag"] = encoder["r_mag"]
 
+    if "threshold" in encoder:
+        threshold = encoder["threshold"]
+        if threshold.ndim == 0:
+            threshold = torch.full((W_enc.size(1),), threshold)
+        state_dict["threshold"] = threshold
+
     return cfg_dict, state_dict, None
 
 

sae_lens/pretrained_saes.yaml

@@ -14959,3 +14959,81 @@ goodfire-llama-3.1-8b-instruct:
     path: Llama-3.1-8B-Instruct-SAE-l19.pth
     l0: 91
     neuronpedia: llama3.1-8b-it/19-resid-post-gf
+
+saebench_gemma-2-2b_width-2pow12_date-0108:
+  conversion_func: dictionary_learning_1
+  links:
+    model: https://huggingface.co/google/gemma-2-2b
+  model: gemma-2-2b
+  repo_id: adamkarvonen/saebench_gemma-2-2b_width-2pow12_date-0108
+  saes:
+  - id: blocks.12.hook_resid_post__trainer_0
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_0_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_0
+  - id: blocks.12.hook_resid_post__trainer_1
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_1_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_1
+  - id: blocks.12.hook_resid_post__trainer_2
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_2_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_2
+  - id: blocks.12.hook_resid_post__trainer_3
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_3_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_3
+  - id: blocks.12.hook_resid_post__trainer_4
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_4_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_4
+  - id: blocks.12.hook_resid_post__trainer_5
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-4k__trainer_5_step_final
+    path: MatryoshkaBatchTopK_gemma-2-2b__0108/resid_post_layer_12/trainer_5
+
+saebench_gemma-2-2b_width-2pow14_date-0107:
+  conversion_func: dictionary_learning_1
+  links:
+    model: https://huggingface.co/google/gemma-2-2b
+  model: gemma-2-2b
+  repo_id: canrager/saebench_gemma-2-2b_width-2pow14_date-0107
+  saes:
+  - id: blocks.12.hook_resid_post__trainer_0
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_0_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_0
+  - id: blocks.12.hook_resid_post__trainer_1
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_1_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_1
+  - id: blocks.12.hook_resid_post__trainer_2
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_2_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_2
+  - id: blocks.12.hook_resid_post__trainer_3
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_3_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_3
+  - id: blocks.12.hook_resid_post__trainer_4
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_4_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_4
+  - id: blocks.12.hook_resid_post__trainer_5
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-16k__trainer_5_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow14_date-0107/resid_post_layer_12/trainer_5
+
+saebench_gemma-2-2b_width-2pow16_date-0107:
+  conversion_func: dictionary_learning_1
+  links:
+    model: https://huggingface.co/google/gemma-2-2b
+  model: gemma-2-2b
+  repo_id: canrager/saebench_gemma-2-2b_width-2pow16_date-0107
+  saes:
+  - id: blocks.12.hook_resid_post__trainer_0
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_0_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_0
+  - id: blocks.12.hook_resid_post__trainer_1
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_1_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_1
+  - id: blocks.12.hook_resid_post__trainer_2
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_2_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_2
+  - id: blocks.12.hook_resid_post__trainer_3
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_3_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_3
+  - id: blocks.12.hook_resid_post__trainer_4
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_4_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_4
+  - id: blocks.12.hook_resid_post__trainer_5
+    neuronpedia: gemma-2-2b/12-sae_bench-matryoshka-res-65k__trainer_5_step_final
+    path: gemma-2-2b_matryoshka_batch_top_k_width-2pow16_date-0107/resid_post_layer_12/trainer_5

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

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

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

@@ -1,4 +1,4 @@
-sae_lens/__init__.py,sha256=-q6U-a-hKOgMuJj8jBEyV5JvAZolVh6ccirbA3zrpYg,4033
+sae_lens/__init__.py,sha256=XYVTI_d_xGq0CTozbkj8QZA3Mrf4U4lNZyCcE8uv3x8,4033
 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,10 +9,10 @@ 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=_YmZAiFd2Av-mgwEBzvQaylgwlwfY-Aqj62s2-C8SIs,56589
 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/pretrained_saes.yaml,sha256=pKfO-KfBU7HlhdukEP9Wdxxoka6LRrDDcsVAI0HYqjg,609596
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
 sae_lens/saes/__init__.py,sha256=nTNPnJ7edyfedo1MX96xwn9WOG8504yHbT9LFw9od_0,1778
 sae_lens/saes/batchtopk_sae.py,sha256=x4EbgZl0GUickRPcCmtKNGS2Ra3Uy1Z1OtF2FnrSabQ,5422
@@ -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.23.0.dist-info/METADATA,sha256=BJ0kFOaNRgcbMx-bLpI17h42zVugibTBCZPG0JaCDLU,5369
+sae_lens-6.23.0.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+sae_lens-6.23.0.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.23.0.dist-info/RECORD,,