From ec99df8da72b40604b9953a96dee9644dc457ceb Mon Sep 17 00:00:00 2001 From: Tauheed Elahee Date: Thu, 12 Mar 2026 21:13:06 -0400 Subject: [PATCH] Add compatibility for newer sae_lens forward() return type Newer versions of sae_lens changed forward() to return a reconstructed tensor instead of an object with feature_acts. Add hasattr check and fall back to encode() to get feature activations. --- .../generate_feature_occurence_data.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/sae_multid_feature_discovery/generate_feature_occurence_data.py b/sae_multid_feature_discovery/generate_feature_occurence_data.py index 63733aa..7301f76 100644 --- a/sae_multid_feature_discovery/generate_feature_occurence_data.py +++ b/sae_multid_feature_discovery/generate_feature_occurence_data.py @@ -139,8 +139,12 @@ def next_batch_activations(): forward_pass = ae.forward(activations) if isinstance(forward_pass, tuple): hidden_sae = forward_pass[1] - else: + elif hasattr(forward_pass, "feature_acts"): hidden_sae = forward_pass.feature_acts + else: + # Newer sae_lens returns reconstructed tensor from forward(); + # use encode() to get feature activations instead + hidden_sae = ae.encode(activations) nonzero_sae = hidden_sae.abs() > 1e-6 nonzero_sae_values = hidden_sae[nonzero_sae]