fix(topk activation): add keepdim=True to enable broadcasting; make dtype consistent without hardcode

Author: Hzfinfdu

src/lm_saes/evaluator.py

@@ -4,11 +4,11 @@
 from torch import Tensor
 from tqdm import tqdm
 from transformer_lens import HookedTransformer
+from wandb.sdk.wandb_run import Run
 
 from lm_saes.config import EvalConfig, MixCoderConfig, SAEConfig
 from lm_saes.mixcoder import MixCoder
 from lm_saes.sae import SparseAutoEncoder
-from wandb.sdk.wandb_run import Run
 
 
 class Evaluator:

src/lm_saes/runner.py

@@ -2,11 +2,11 @@
 from pathlib import Path
 from typing import Literal, Optional, TypeVar, overload
 
+import wandb
 from pydantic import model_validator
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from torch.distributed.device_mesh import init_device_mesh
 
-import wandb
 from lm_saes.activation.factory import ActivationFactory
 from lm_saes.activation.writer import ActivationWriter
 from lm_saes.analysis.feature_analyzer import FeatureAnalyzer

src/lm_saes/trainer.py

@@ -7,11 +7,11 @@
 from torch import Tensor
 from torch.optim import Adam, Optimizer
 from tqdm import tqdm
+from wandb.sdk.wandb_run import Run
 
 from lm_saes.config import TrainerConfig
 from lm_saes.optim import get_scheduler
 from lm_saes.sae import SparseAutoEncoder
-from wandb.sdk.wandb_run import Run
 
 
 class Trainer: