sae-lens 6.3.0__tar.gz → 6.3.1__tar.gz

{sae_lens-6.3.0 → sae_lens-6.3.1}/PKG-INFO

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

{sae_lens-6.3.0 → sae_lens-6.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.3.0"
+version = "6.3.1"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"

{sae_lens-6.3.0 → sae_lens-6.3.1}/sae_lens/__init__.py

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

{sae_lens-6.3.0 → sae_lens-6.3.1}/sae_lens/llm_sae_training_runner.py

@@ -17,6 +17,7 @@ from sae_lens.config import HfDataset, LanguageModelSAERunnerConfig
 from sae_lens.constants import ACTIVATIONS_STORE_STATE_FILENAME, RUNNER_CFG_FILENAME
 from sae_lens.evals import EvalConfig, run_evals
 from sae_lens.load_model import load_model
+from sae_lens.saes.batchtopk_sae import BatchTopKTrainingSAEConfig
 from sae_lens.saes.gated_sae import GatedTrainingSAEConfig
 from sae_lens.saes.jumprelu_sae import JumpReLUTrainingSAEConfig
 from sae_lens.saes.sae import (
@@ -291,7 +292,7 @@ def _parse_cfg_args(
     architecture_parser.add_argument(
         "--architecture",
         type=str,
-        choices=["standard", "gated", "jumprelu", "topk"],
+        choices=["standard", "gated", "jumprelu", "topk", "batchtopk"],
         default="standard",
         help="SAE architecture to use",
     )
@@ -352,6 +353,7 @@ def _parse_cfg_args(
         "gated": GatedTrainingSAEConfig,
         "jumprelu": JumpReLUTrainingSAEConfig,
         "topk": TopKTrainingSAEConfig,
+        "batchtopk": BatchTopKTrainingSAEConfig,
     }
 
     sae_config_type = sae_config_map[architecture]

{sae_lens-6.3.0 → sae_lens-6.3.1}/sae_lens/training/sae_trainer.py

@@ -1,5 +1,4 @@
 import contextlib
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable, Generic, Protocol
 
@@ -38,13 +37,6 @@ def _update_sae_lens_training_version(sae: TrainingSAE[Any]) -> None:
     sae.cfg.sae_lens_training_version = str(__version__)
 
 
-@dataclass
-class TrainSAEOutput:
-    sae: TrainingSAE[Any]
-    checkpoint_path: str
-    log_feature_sparsities: torch.Tensor
-
-
 class SaveCheckpointFn(Protocol):
     def __call__(
         self,