sae-lens 6.6.4__py3-none-any.whl → 6.6.5__py3-none-any.whl

sae_lens/__init__.py

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

sae_lens/evals.py

@@ -718,17 +718,9 @@ def get_recons_loss(
         **model_kwargs,
     )
 
-    def kl(original_logits: torch.Tensor, new_logits: torch.Tensor):
-        original_probs = torch.nn.functional.softmax(original_logits, dim=-1)
-        log_original_probs = torch.log(original_probs)
-        new_probs = torch.nn.functional.softmax(new_logits, dim=-1)
-        log_new_probs = torch.log(new_probs)
-        kl_div = original_probs * (log_original_probs - log_new_probs)
-        return kl_div.sum(dim=-1)
-
     if compute_kl:
-        recons_kl_div = kl(original_logits, recons_logits)
-        zero_abl_kl_div = kl(original_logits, zero_abl_logits)
+        recons_kl_div = _kl(original_logits, recons_logits)
+        zero_abl_kl_div = _kl(original_logits, zero_abl_logits)
         metrics["kl_div_with_sae"] = recons_kl_div
         metrics["kl_div_with_ablation"] = zero_abl_kl_div
 
@@ -740,6 +732,18 @@ def get_recons_loss(
     return metrics
 
 
+def _kl(original_logits: torch.Tensor, new_logits: torch.Tensor):
+    # Computes the log-probabilities of the new logits (approximation).
+    log_probs_new = torch.nn.functional.log_softmax(new_logits, dim=-1)
+    # Computes the probabilities of the original logits (true distribution).
+    probs_orig = torch.nn.functional.softmax(original_logits, dim=-1)
+    # Compute the KL divergence. torch.nn.functional.kl_div expects the first argument to be the log
+    # probabilities of the approximation (new), and the second argument to be the true distribution
+    # (original) as probabilities. This computes KL(original || new).
+    kl = torch.nn.functional.kl_div(log_probs_new, probs_orig, reduction="none")
+    return kl.sum(dim=-1)
+
+
 def all_loadable_saes() -> list[tuple[str, str, float, float]]:
     all_loadable_saes = []
     saes_directory = get_pretrained_saes_directory()

sae_lens/loading/pretrained_sae_loaders.py

@@ -1001,10 +1001,14 @@ def get_sparsify_config_from_disk(
         layer = int(match.group(1))
         hook_name = f"blocks.{layer}.hook_resid_post"
 
+    d_sae = old_cfg_dict.get("num_latents")
+    if d_sae is None:
+        d_sae = old_cfg_dict["d_in"] * old_cfg_dict["expansion_factor"]
+
     cfg_dict: dict[str, Any] = {
         "architecture": "standard",
         "d_in": old_cfg_dict["d_in"],
-        "d_sae": old_cfg_dict["d_in"] * old_cfg_dict["expansion_factor"],
+        "d_sae": d_sae,
         "dtype": "bfloat16",
         "device": device or "cpu",
         "model_name": config_dict.get("model", path.parts[-2]),

{sae_lens-6.6.4.dist-info → sae_lens-6.6.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sae-lens
-Version: 6.6.4
+Version: 6.6.5
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 Keywords: deep-learning,sparse-autoencoders,mechanistic-interpretability,PyTorch

{sae_lens-6.6.4.dist-info → sae_lens-6.6.5.dist-info}/RECORD

@@ -1,15 +1,15 @@
-sae_lens/__init__.py,sha256=yJm-dL16dtJ9Eo9p4cOcV3Y4Zrhim6FMuzysx0MnYas,3588
+sae_lens/__init__.py,sha256=gvg9photJRtatuXa9YF-uDv1tYiHwHTMh29X1GNQd6Y,3588
 sae_lens/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/analysis/hooked_sae_transformer.py,sha256=vRu6JseH1lZaEeILD5bEkQEQ1wYHHDcxD-f2olKmE9Y,14275
 sae_lens/analysis/neuronpedia_integration.py,sha256=Fj4gVyaXMGBUxoK0vPeTwGVFr4n40fmfPrRENo4WzPs,19324
 sae_lens/cache_activations_runner.py,sha256=cNeAtp2JQ_vKbeddZVM-tcPLYyyfTWL8NDna5KQpkLI,12583
 sae_lens/config.py,sha256=IrjbsKBbaZoFXYrsPJ5xBwIqi9uZJIIFXjV_uoErJaE,28176
 sae_lens/constants.py,sha256=CSjmiZ-bhjQeVLyRvWxAjBokCgkfM8mnvd7-vxLIWTY,639
-sae_lens/evals.py,sha256=2YHR_IBhXdjktpmoVtvvNrqUZIx5ok7yERuiFY40HHY,39186
+sae_lens/evals.py,sha256=4hanbyG8qZLItWqft94F4ZjUoytPVB7fw5s0P4Oi0VE,39504
 sae_lens/llm_sae_training_runner.py,sha256=exxNX_OEhdiUrlgmBP9bjX9DOf0HUcNQGO4unKeDjKM,13713
 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=iIcHM24qfb45JOGEmUn7jr5E9vl8L2FYSlArsobCwlI,44388
+sae_lens/loading/pretrained_sae_loaders.py,sha256=tLeHArWFpu8CI6vXH1ZxFkhmsrhO2UsZyi7DzVzqAUs,44477
 sae_lens/loading/pretrained_saes_directory.py,sha256=4Vn-Jex6SveD7EbxcSOBv8cx1gkPfUMLU1QOP-ww1ZE,3752
 sae_lens/pretokenize_runner.py,sha256=w0f6SfZLAxbp5eAAKnet8RqUB_DKofZ9RGsoJwFnYbA,7058
 sae_lens/pretrained_saes.yaml,sha256=O_FwoOe7fU9_WLEOnMk1IWXRxD4nwzf1tCfbof1r0D0,598578
@@ -33,7 +33,7 @@ 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.6.4.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.6.4.dist-info/METADATA,sha256=2fVRS6CU6AgaiZZBDgwLcFPClSs_KAMUXWP_uJB27TE,5356
-sae_lens-6.6.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-sae_lens-6.6.4.dist-info/RECORD,,
+sae_lens-6.6.5.dist-info/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.6.5.dist-info/METADATA,sha256=U5oP3RYgIE2EnHA2mwRImUcoyVBhYYwiRU199LM_R7c,5356
+sae_lens-6.6.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+sae_lens-6.6.5.dist-info/RECORD,,