sae-lens 6.28.1__py3-none-any.whl → 6.28.2__py3-none-any.whl

sae_lens/__init__.py

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

sae_lens/pretrained_saes.yaml

@@ -40631,7 +40631,7 @@ gemma-3-1b-res-matryoshka-dc:
   conversion_func: null
   links:
     model: https://huggingface.co/google/gemma-3-1b-pt
-  model: gemma-3-1b
+  model: google/gemma-3-1b-pt
   repo_id: chanind/gemma-3-1b-batch-topk-matryoshka-saes-w-32k-l0-40
   saes:
   - id: blocks.0.hook_resid_post

sae_lens/synthetic/activation_generator.py

@@ -63,6 +63,7 @@ class ActivationGenerator(nn.Module):
             )
         self.correlation_matrix = correlation_matrix
 
+    @torch.no_grad()
     def sample(self, batch_size: int) -> torch.Tensor:
         """
         Generate a batch of feature activations with controlled properties.

sae_lens/synthetic/feature_dictionary.py

@@ -16,11 +16,28 @@ FeatureDictionaryInitializer = Callable[["FeatureDictionary"], None]
 
 def orthogonalize_embeddings(
     embeddings: torch.Tensor,
-    target_cos_sim: float = 0,
     num_steps: int = 200,
     lr: float = 0.01,
     show_progress: bool = False,
+    chunk_size: int = 1024,
 ) -> torch.Tensor:
+    """
+    Orthogonalize embeddings using gradient descent with chunked computation.
+
+    Uses chunked computation to avoid O(n²) memory usage when computing pairwise
+    dot products. Memory usage is O(chunk_size × n) instead of O(n²).
+
+    Args:
+        embeddings: Tensor of shape [num_vectors, hidden_dim]
+        num_steps: Number of optimization steps
+        lr: Learning rate for Adam optimizer
+        show_progress: Whether to show progress bar
+        chunk_size: Number of vectors to process at once. Smaller values use less
+            memory but may be slower.
+
+    Returns:
+        Orthogonalized embeddings of the same shape, normalized to unit length.
+    """
     num_vectors = embeddings.shape[0]
     # Create a detached copy and normalize, then enable gradients
     embeddings = embeddings.detach().clone()
@@ -29,24 +46,37 @@ def orthogonalize_embeddings(
 
     optimizer = torch.optim.Adam([embeddings], lr=lr)  # type: ignore[list-item]
 
-    # Create a mask to zero out diagonal elements (avoid in-place operations)
-    off_diagonal_mask = ~torch.eye(
-        num_vectors, dtype=torch.bool, device=embeddings.device
-    )
-
     pbar = tqdm(
         range(num_steps), desc="Orthogonalizing vectors", disable=not show_progress
     )
     for _ in pbar:
         optimizer.zero_grad()
 
-        dot_products = embeddings @ embeddings.T
-        diff = dot_products - target_cos_sim
-        # Use masking instead of in-place fill_diagonal_
-        off_diagonal_diff = diff * off_diagonal_mask.float()
-        loss = off_diagonal_diff.pow(2).sum()
-        loss = loss + num_vectors * (dot_products.diag() - 1).pow(2).sum()
+        off_diag_loss = torch.tensor(0.0, device=embeddings.device)
+        diag_loss = torch.tensor(0.0, device=embeddings.device)
+
+        for i in range(0, num_vectors, chunk_size):
+            end_i = min(i + chunk_size, num_vectors)
+            chunk = embeddings[i:end_i]
+            chunk_dots = chunk @ embeddings.T  # [chunk_size, num_vectors]
 
+            # Create mask to zero out diagonal elements for this chunk
+            # Diagonal of full matrix: position (i+k, i+k) → in chunk_dots: (k, i+k)
+            chunk_len = end_i - i
+            row_indices = torch.arange(chunk_len, device=embeddings.device)
+            col_indices = i + row_indices  # column indices in full matrix
+
+            # Boolean mask: True for off-diagonal elements we want to include
+            off_diag_mask = torch.ones_like(chunk_dots, dtype=torch.bool)
+            off_diag_mask[row_indices, col_indices] = False
+
+            off_diag_loss = off_diag_loss + chunk_dots[off_diag_mask].pow(2).sum()
+
+            # Diagonal loss: keep self-dot-products at 1
+            diag_vals = chunk_dots[row_indices, col_indices]
+            diag_loss = diag_loss + (diag_vals - 1).pow(2).sum()
+
+        loss = off_diag_loss + num_vectors * diag_loss
         loss.backward()
         optimizer.step()
         pbar.set_description(f"loss: {loss.item():.3f}")
@@ -59,7 +89,10 @@ def orthogonalize_embeddings(
 
 
 def orthogonal_initializer(
-    num_steps: int = 200, lr: float = 0.01, show_progress: bool = False
+    num_steps: int = 200,
+    lr: float = 0.01,
+    show_progress: bool = False,
+    chunk_size: int = 1024,
 ) -> FeatureDictionaryInitializer:
     def initializer(feature_dict: "FeatureDictionary") -> None:
         feature_dict.feature_vectors.data = orthogonalize_embeddings(
@@ -67,6 +100,7 @@ def orthogonal_initializer(
             num_steps=num_steps,
             lr=lr,
             show_progress=show_progress,
+            chunk_size=chunk_size,
         )
 
     return initializer
@@ -97,6 +131,7 @@ class FeatureDictionary(nn.Module):
         hidden_dim: int,
         bias: bool = False,
         initializer: FeatureDictionaryInitializer | None = orthogonal_initializer(),
+        device: str | torch.device = "cpu",
     ):
         """
         Create a new FeatureDictionary.
@@ -106,20 +141,23 @@ class FeatureDictionary(nn.Module):
             hidden_dim: Dimensionality of the hidden space
             bias: Whether to include a bias term in the embedding
             initializer: Initializer function to use. If None, the embeddings are initialized to random unit vectors. By default will orthogonalize embeddings.
+            device: Device to use for the feature dictionary.
         """
         super().__init__()
         self.num_features = num_features
         self.hidden_dim = hidden_dim
 
         # Initialize feature vectors as unit vectors
-        embeddings = torch.randn(num_features, hidden_dim)
+        embeddings = torch.randn(num_features, hidden_dim, device=device)
         embeddings = embeddings / embeddings.norm(p=2, dim=1, keepdim=True).clamp(
             min=1e-8
         )
         self.feature_vectors = nn.Parameter(embeddings)
 
         # Initialize bias (zeros if not using bias, but still a parameter for consistent API)
-        self.bias = nn.Parameter(torch.zeros(hidden_dim), requires_grad=bias)
+        self.bias = nn.Parameter(
+            torch.zeros(hidden_dim, device=device), requires_grad=bias
+        )
 
         if initializer is not None:
             initializer(self)

sae_lens/synthetic/hierarchy.py

@@ -11,7 +11,9 @@ https://github.com/noanabeshima/matryoshka-saes/blob/main/toy_model.py
 
 from __future__ import annotations
 
+from collections import deque
 from collections.abc import Callable, Sequence
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -19,6 +21,7 @@ import torch
 ActivationsModifier = Callable[[torch.Tensor], torch.Tensor]
 
 
+@torch.no_grad()
 def _validate_hierarchy(roots: Sequence[HierarchyNode]) -> None:
     """
     Validate a forest of hierarchy trees.
@@ -104,6 +107,320 @@ def _node_description(node: HierarchyNode) -> str:
     return "unnamed node"
 
 
+# ---------------------------------------------------------------------------
+# Vectorized hierarchy implementation
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _LevelData:
+    """Data for a single level in the hierarchy."""
+
+    # Features at this level and their parents (for parent deactivation)
+    features: torch.Tensor  # [num_features_at_level]
+    parents: torch.Tensor  # [num_features_at_level]
+
+    # ME group indices to process AFTER this level's parent deactivation
+    # These are groups whose parent node is at this level
+    # ME must be applied here before processing next level's parent deactivation
+    me_group_indices: torch.Tensor  # [num_groups_at_level], may be empty
+
+
+@dataclass
+class _SparseHierarchyData:
+    """Precomputed data for sparse hierarchy processing.
+
+    This structure enables O(active_features) processing instead of O(all_groups).
+    ME is applied at each level after parent deactivation to ensure cascading works.
+    """
+
+    # Per-level data for parent deactivation and ME (processed in order)
+    level_data: list[_LevelData]
+
+    # ME group data (shared across levels, indexed by me_group_indices)
+    me_group_siblings: torch.Tensor  # [num_groups, max_siblings]
+    me_group_sizes: torch.Tensor  # [num_groups]
+    me_group_parents: (
+        torch.Tensor
+    )  # [num_groups] - parent feature index (-1 if no parent)
+
+    # Total number of ME groups
+    num_groups: int
+
+
+def _build_sparse_hierarchy(
+    roots: Sequence[HierarchyNode],
+) -> _SparseHierarchyData:
+    """
+    Build sparse hierarchy data structure for O(active_features) processing.
+
+    The key insight is that ME groups must be applied at the level of their parent node,
+    AFTER parent deactivation at that level, but BEFORE processing the next level.
+    This ensures that when a child is deactivated by ME, its grandchildren are also
+    deactivated during the next level's parent deactivation.
+    """
+    # Collect feature info by level using BFS
+    # Each entry: (feature_index, effective_parent, level)
+    feature_info: list[tuple[int, int, int]] = []
+
+    # ME groups: list of (parent_level, parent_feature, child_feature_indices)
+    me_groups: list[tuple[int, int, list[int]]] = []
+
+    # BFS queue: (node, effective_parent, level)
+    queue: deque[tuple[HierarchyNode, int, int]] = deque()
+    for root in roots:
+        queue.append((root, -1, 0))
+
+    while queue:
+        node, effective_parent, level = queue.popleft()
+
+        if node.feature_index is not None:
+            feature_info.append((node.feature_index, effective_parent, level))
+            new_effective_parent = node.feature_index
+        else:
+            new_effective_parent = effective_parent
+
+        # Handle mutual exclusion children - record the parent's level and feature
+        if node.mutually_exclusive_children and len(node.children) >= 2:
+            child_feats = [
+                c.feature_index for c in node.children if c.feature_index is not None
+            ]
+            if len(child_feats) >= 2:
+                # ME group belongs to the parent's level (current level)
+                # Parent feature is the node's feature_index (-1 if organizational node)
+                parent_feat = (
+                    node.feature_index if node.feature_index is not None else -1
+                )
+                me_groups.append((level, parent_feat, child_feats))
+
+        for child in node.children:
+            queue.append((child, new_effective_parent, level + 1))
+
+    # Determine max level for both features and ME groups
+    max_feature_level = max((info[2] for info in feature_info), default=-1)
+    max_me_level = max((lvl for lvl, _, _ in me_groups), default=-1)
+    max_level = max(max_feature_level, max_me_level)
+
+    # Build level data with ME group indices per level
+    level_data: list[_LevelData] = []
+
+    # Group ME groups by their parent level
+    me_groups_by_level: dict[int, list[int]] = {}
+    for g_idx, (parent_level, _, _) in enumerate(me_groups):
+        if parent_level not in me_groups_by_level:
+            me_groups_by_level[parent_level] = []
+        me_groups_by_level[parent_level].append(g_idx)
+
+    for level in range(max_level + 1):
+        # Get features at this level that have parents
+        features_at_level = [
+            (feat, parent) for feat, parent, lv in feature_info if lv == level
+        ]
+        with_parents = [(f, p) for f, p in features_at_level if p >= 0]
+
+        if with_parents:
+            feats = torch.tensor([f for f, _ in with_parents], dtype=torch.long)
+            parents = torch.tensor([p for _, p in with_parents], dtype=torch.long)
+        else:
+            feats = torch.empty(0, dtype=torch.long)
+            parents = torch.empty(0, dtype=torch.long)
+
+        # Get ME group indices for this level
+        if level in me_groups_by_level:
+            me_indices = torch.tensor(me_groups_by_level[level], dtype=torch.long)
+        else:
+            me_indices = torch.empty(0, dtype=torch.long)
+
+        level_data.append(
+            _LevelData(features=feats, parents=parents, me_group_indices=me_indices)
+        )
+
+    # Build group siblings and parents tensors
+    if me_groups:
+        max_siblings = max(len(children) for _, _, children in me_groups)
+        num_groups = len(me_groups)
+        me_group_siblings = torch.full((num_groups, max_siblings), -1, dtype=torch.long)
+        me_group_sizes = torch.zeros(num_groups, dtype=torch.long)
+        me_group_parents = torch.full((num_groups,), -1, dtype=torch.long)
+        for g_idx, (_, parent_feat, siblings) in enumerate(me_groups):
+            me_group_sizes[g_idx] = len(siblings)
+            me_group_parents[g_idx] = parent_feat
+            me_group_siblings[g_idx, : len(siblings)] = torch.tensor(
+                siblings, dtype=torch.long
+            )
+    else:
+        me_group_siblings = torch.empty((0, 0), dtype=torch.long)
+        me_group_sizes = torch.empty(0, dtype=torch.long)
+        me_group_parents = torch.empty(0, dtype=torch.long)
+        num_groups = 0
+
+    return _SparseHierarchyData(
+        level_data=level_data,
+        me_group_siblings=me_group_siblings,
+        me_group_sizes=me_group_sizes,
+        me_group_parents=me_group_parents,
+        num_groups=num_groups,
+    )
+
+
+def _apply_hierarchy_sparse(
+    activations: torch.Tensor,
+    sparse_data: _SparseHierarchyData,
+) -> torch.Tensor:
+    """
+    Apply hierarchy constraints using precomputed sparse indices.
+
+    Processes level by level:
+    1. Apply parent deactivation for features at this level
+    2. Apply mutual exclusion for groups whose parent is at this level
+    3. Move to next level
+
+    This ensures that ME at level L affects parent deactivation at level L+1.
+    """
+    result = activations.clone()
+
+    # Data is already on correct device from cache
+    me_group_siblings = sparse_data.me_group_siblings
+    me_group_sizes = sparse_data.me_group_sizes
+    me_group_parents = sparse_data.me_group_parents
+
+    for level_data in sparse_data.level_data:
+        # Step 1: Deactivate children where parent is inactive
+        if level_data.features.numel() > 0:
+            parent_vals = result[:, level_data.parents]
+            child_vals = result[:, level_data.features]
+            result[:, level_data.features] = child_vals * (parent_vals > 0)
+
+        # Step 2: Apply ME for groups whose parent is at this level
+        if level_data.me_group_indices.numel() > 0:
+            _apply_me_for_groups(
+                result,
+                level_data.me_group_indices,
+                me_group_siblings,
+                me_group_sizes,
+                me_group_parents,
+            )
+
+    return result
+
+
+def _apply_me_for_groups(
+    activations: torch.Tensor,
+    group_indices: torch.Tensor,
+    me_group_siblings: torch.Tensor,
+    me_group_sizes: torch.Tensor,
+    me_group_parents: torch.Tensor,
+) -> None:
+    """
+    Apply mutual exclusion for the specified groups.
+
+    Only processes groups where the parent is active (or has no parent).
+    This is a key optimization since most groups are skipped when parent is inactive.
+
+    Args:
+        activations: [batch_size, num_features] - modified in place
+        group_indices: [num_groups_to_process] - which groups to apply ME for
+        me_group_siblings: [total_groups, max_siblings] - sibling indices per group
+        me_group_sizes: [total_groups] - number of valid siblings per group
+        me_group_parents: [total_groups] - parent feature index (-1 if no parent)
+    """
+    batch_size = activations.shape[0]
+    device = activations.device
+    num_groups = group_indices.numel()
+
+    if num_groups == 0:
+        return
+
+    # Get parent indices for these groups
+    parents = me_group_parents[group_indices]  # [num_groups]
+
+    # Check which parents are active: [batch_size, num_groups]
+    # Groups with parent=-1 are always active (root-level ME)
+    has_parent = parents >= 0
+    if has_parent.all():
+        # All groups have parents - check their activation directly
+        parent_active = activations[:, parents] > 0  # [batch, num_groups]
+        if not parent_active.any():
+            return
+    elif has_parent.any():
+        # Mixed case: some groups have parents, some don't
+        # Use clamp to avoid indexing with -1 (reads feature 0, but result is masked out)
+        safe_parents = parents.clamp(min=0)
+        parent_active = activations[:, safe_parents] > 0  # [batch, num_groups]
+        # Groups without parent are always "active"
+        parent_active = parent_active | ~has_parent
+    else:
+        # No groups have parents - all are always active, skip parent check
+        parent_active = None
+
+    # Get siblings for the groups we're processing
+    siblings = me_group_siblings[group_indices]  # [num_groups, max_siblings]
+    sizes = me_group_sizes[group_indices]  # [num_groups]
+    max_siblings = siblings.shape[1]
+
+    # Get activations for all siblings: [batch_size, num_groups, max_siblings]
+    safe_siblings = siblings.clamp(min=0)
+    sibling_activations = activations[:, safe_siblings.view(-1)].view(
+        batch_size, num_groups, max_siblings
+    )
+
+    # Create validity mask for padding: [num_groups, max_siblings]
+    sibling_range = torch.arange(max_siblings, device=device)
+    valid_mask = sibling_range < sizes.unsqueeze(1)
+
+    # Find active valid siblings, but only where parent is active: [batch, groups, siblings]
+    sibling_active = (sibling_activations > 0) & valid_mask
+    if parent_active is not None:
+        sibling_active = sibling_active & parent_active.unsqueeze(2)
+
+    # Count active per group and check for conflicts: [batch_size, num_groups]
+    active_counts = sibling_active.sum(dim=2)
+    needs_exclusion = active_counts > 1
+
+    if not needs_exclusion.any():
+        return
+
+    # Get (batch, group) pairs needing exclusion
+    batch_with_conflict, groups_with_conflict = torch.where(needs_exclusion)
+    num_conflicts = batch_with_conflict.numel()
+
+    if num_conflicts == 0:
+        return
+
+    # Get siblings and activations for conflicts
+    conflict_siblings = siblings[groups_with_conflict]  # [num_conflicts, max_siblings]
+    conflict_active = sibling_active[
+        batch_with_conflict, groups_with_conflict
+    ]  # [num_conflicts, max_siblings]
+
+    # Random selection for winner
+    # Use -1e9 instead of -inf to avoid creating a tensor (torch.tensor(-float("inf")))
+    # on every call. Since random scores are in [0,1], -1e9 is effectively -inf for argmax.
+    random_scores = torch.rand(num_conflicts, max_siblings, device=device)
+    random_scores[~conflict_active] = -1e9
+
+    winner_idx = random_scores.argmax(dim=1)
+
+    # Determine losers using scatter for efficiency
+    is_winner = torch.zeros(
+        num_conflicts, max_siblings, dtype=torch.bool, device=device
+    )
+    is_winner.scatter_(1, winner_idx.unsqueeze(1), True)
+    should_deactivate = conflict_active & ~is_winner
+
+    # Get (conflict, sibling) pairs to deactivate
+    conflict_idx, sib_idx = torch.where(should_deactivate)
+
+    if conflict_idx.numel() == 0:
+        return
+
+    # Map back to (batch, feature) and deactivate
+    deact_batch = batch_with_conflict[conflict_idx]
+    deact_feat = conflict_siblings[conflict_idx, sib_idx]
+    activations[deact_batch, deact_feat] = 0
+
+
+@torch.no_grad()
 def hierarchy_modifier(
     roots: Sequence[HierarchyNode] | HierarchyNode,
 ) -> ActivationsModifier:
@@ -136,12 +453,35 @@ def hierarchy_modifier(
         roots = [roots]
     _validate_hierarchy(roots)
 
-    # Create modifier function that applies all hierarchies
+    # Build sparse hierarchy data
+    sparse_data = _build_sparse_hierarchy(roots)
+
+    # Cache for device-specific tensors
+    device_cache: dict[torch.device, _SparseHierarchyData] = {}
+
+    def _get_sparse_for_device(device: torch.device) -> _SparseHierarchyData:
+        """Get or create device-specific sparse hierarchy data."""
+        if device not in device_cache:
+            device_cache[device] = _SparseHierarchyData(
+                level_data=[
+                    _LevelData(
+                        features=ld.features.to(device),
+                        parents=ld.parents.to(device),
+                        me_group_indices=ld.me_group_indices.to(device),
+                    )
+                    for ld in sparse_data.level_data
+                ],
+                me_group_siblings=sparse_data.me_group_siblings.to(device),
+                me_group_sizes=sparse_data.me_group_sizes.to(device),
+                me_group_parents=sparse_data.me_group_parents.to(device),
+                num_groups=sparse_data.num_groups,
+            )
+        return device_cache[device]
+
     def modifier(activations: torch.Tensor) -> torch.Tensor:
-        result = activations.clone()
-        for root in roots:
-            root._apply_hierarchy(result, parent_active_mask=None)
-        return result
+        device = activations.device
+        cached = _get_sparse_for_device(device)
+        return _apply_hierarchy_sparse(activations, cached)
 
     return modifier
 
@@ -222,85 +562,6 @@ class HierarchyNode:
         if self.mutually_exclusive_children and len(self.children) < 2:
             raise ValueError("Need at least 2 children for mutual exclusion")
 
-    def _apply_hierarchy(
-        self,
-        activations: torch.Tensor,
-        parent_active_mask: torch.Tensor | None,
-    ) -> None:
-        """Recursively apply hierarchical constraints."""
-        batch_size = activations.shape[0]
-
-        # Determine which samples have this node active
-        if self.feature_index is not None:
-            is_active = activations[:, self.feature_index] > 0
-        else:
-            # Non-readout node: active if parent is active (or always if root)
-            is_active = (
-                parent_active_mask
-                if parent_active_mask is not None
-                else torch.ones(batch_size, dtype=torch.bool, device=activations.device)
-            )
-
-        # Deactivate this node if parent is inactive
-        if parent_active_mask is not None and self.feature_index is not None:
-            activations[~parent_active_mask, self.feature_index] = 0
-            # Update is_active after deactivation
-            is_active = activations[:, self.feature_index] > 0
-
-        # Handle mutually exclusive children
-        if self.mutually_exclusive_children and len(self.children) >= 2:
-            self._enforce_mutual_exclusion(activations, is_active)
-
-        # Recursively process children
-        for child in self.children:
-            child._apply_hierarchy(activations, parent_active_mask=is_active)
-
-    def _enforce_mutual_exclusion(
-        self,
-        activations: torch.Tensor,
-        parent_active_mask: torch.Tensor,
-    ) -> None:
-        """Ensure at most one child is active per sample."""
-        batch_size = activations.shape[0]
-
-        # Get indices of children that have feature indices
-        child_indices = [
-            child.feature_index
-            for child in self.children
-            if child.feature_index is not None
-        ]
-
-        if len(child_indices) < 2:
-            return
-
-        # For each sample where parent is active, enforce mutual exclusion.
-        # Note: This loop is not vectorized because we need to randomly select
-        # which child to keep active per sample. Vectorizing would require either
-        # a deterministic selection (losing randomness) or complex gather/scatter
-        # operations that aren't more efficient for typical batch sizes.
-        for batch_idx in range(batch_size):
-            if not parent_active_mask[batch_idx]:
-                continue
-
-            # Find which children are active
-            active_children = [
-                i
-                for i, feat_idx in enumerate(child_indices)
-                if activations[batch_idx, feat_idx] > 0
-            ]
-
-            if len(active_children) <= 1:
-                continue
-
-            # Randomly select one to keep active
-            random_idx = int(torch.randint(len(active_children), (1,)).item())
-            keep_idx = active_children[random_idx]
-
-            # Deactivate all others
-            for i, feat_idx in enumerate(child_indices):
-                if i != keep_idx and i in active_children:
-                    activations[batch_idx, feat_idx] = 0
-
     def get_all_feature_indices(self) -> list[int]:
         """Get all feature indices in this subtree."""
         indices = []

{sae_lens-6.28.1.dist-info → sae_lens-6.28.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.28.1
+Version: 6.28.2
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE
@@ -50,6 +50,8 @@ SAELens exists to help researchers:
 - Analyse sparse autoencoders / research mechanistic interpretability.
 - Generate insights which make it easier to create safe and aligned AI systems.
 
+SAELens inference works with any PyTorch-based model, not just TransformerLens. While we provide deep integration with TransformerLens via `HookedSAETransformer`, SAEs can be used with Hugging Face Transformers, NNsight, or any other framework by extracting activations and passing them to the SAE's `encode()` and `decode()` methods.
+
 Please refer to the [documentation](https://decoderesearch.github.io/SAELens/) for information on how to:
 
 - Download and Analyse pre-trained sparse autoencoders.
@@ -84,6 +86,14 @@ The new v6 update is a major refactor to SAELens and changes the way training co
 
 Feel free to join the [Open Source Mechanistic Interpretability Slack](https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-375zalm04-GFd5tdBU1yLKlu_T_JSqZQ) for support!
 
+## Other SAE Projects
+
+- [dictionary-learning](https://github.com/saprmarks/dictionary_learning): An SAE training library that focuses on having hackable code.
+- [Sparsify](https://github.com/EleutherAI/sparsify): A lean SAE training library focused on TopK SAEs.
+- [Overcomplete](https://github.com/KempnerInstitute/overcomplete): SAE training library focused on vision models.
+- [SAE-Vis](https://github.com/callummcdougall/sae_vis): A library for visualizing SAE features, works with SAELens.
+- [SAEBench](https://github.com/adamkarvonen/SAEBench): A suite of LLM SAE benchmarks, works with SAELens.
+
 ## Citation
 
 Please cite the package as follows:

{sae_lens-6.28.1.dist-info → sae_lens-6.28.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-sae_lens/__init__.py,sha256=S-AS72IxkvKO-wItRQjuyczikDxmfDaUgXRSfu5PU-o,4788
+sae_lens/__init__.py,sha256=B9tY0Jt21pOHmSQrQLpMxQHyUAdLHIZpVP6pg3O0dfQ,4788
 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
@@ -12,7 +12,7 @@ sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 sae_lens/loading/pretrained_sae_loaders.py,sha256=hHMlew1u6zVlbzvS9S_SfUPnAG0_OAjjIcjoUTIUZrU,63657
 sae_lens/loading/pretrained_saes_directory.py,sha256=1at_aQbD8WFywchQCKuwfP-yvCq_Z2aUYrpKDnSN5Nc,4283
 sae_lens/pretokenize_runner.py,sha256=amJwIz3CKi2s2wNQn-10E7eAV7VFhNqtFDNTeTkwEI8,7133
-sae_lens/pretrained_saes.yaml,sha256=Hn8jXwZ7V6QQxzgu41LFEP-LAzuDxwYL5vhoar-pPX8,1509922
+sae_lens/pretrained_saes.yaml,sha256=Nq43dTcFvDDONTuJ9Me_HQ5nHqr9BdbP5-ZJGXj0TAQ,1509932
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
 sae_lens/saes/__init__.py,sha256=SBqPaP6Gl5uPFwHlumAZATC4Wd26xKIYLAAAo4MSa5Q,2200
 sae_lens/saes/batchtopk_sae.py,sha256=x4EbgZl0GUickRPcCmtKNGS2Ra3Uy1Z1OtF2FnrSabQ,5422
@@ -26,12 +26,12 @@ sae_lens/saes/temporal_sae.py,sha256=83Ap4mYGfdN3sKdPF8nKjhdXph3-7E2QuLobqJ_YuoM
 sae_lens/saes/topk_sae.py,sha256=vrMRPrCQR1o8G_kXqY_EAoGZARupkQNFB2dNZVLsusE,21073
 sae_lens/saes/transcoder.py,sha256=CTpJs8ASOK06npih7gZHygZuxqTR7HICWlOYfTiKjI4,13501
 sae_lens/synthetic/__init__.py,sha256=FGUasB6fLPXRFCcrtKfL7vCKDOWebZ5Rx5F9QNJZklI,2875
-sae_lens/synthetic/activation_generator.py,sha256=thWGTwRmhu0K8m66WfJUajHmuIPHkwV4_HjmG0dL3G8,7638
+sae_lens/synthetic/activation_generator.py,sha256=JEN7mEgdGDuXr0ArTwUsSdSVUAfvheT_1Eew2ojbA-g,7659
 sae_lens/synthetic/correlation.py,sha256=odr-S5h6c2U-bepwrAQeMfV1iBF_cnnQzqw7zapEXZ4,6056
 sae_lens/synthetic/evals.py,sha256=Nhi314ZnRgLfhBj-3tm_zzI-pGyFTcwllDXbIpPFXeU,4584
-sae_lens/synthetic/feature_dictionary.py,sha256=2A9wqdT1KejRLuIoFWdoiWdDtaHHgIluaKsHGizsVxI,4864
+sae_lens/synthetic/feature_dictionary.py,sha256=ysn0ihE3JgVlCLUZMb127WYZqbz4kMp9BGHfCZqERBg,6487
 sae_lens/synthetic/firing_probabilities.py,sha256=yclz1pWl5gE1r8LAxFvzQS88Lxwk5-3r8BCX9HLVejA,3370
-sae_lens/synthetic/hierarchy.py,sha256=dlQdPnnG3VzQDB3QOaqSXwoH8Ij2ioxmTlZg1lXHaRQ,11754
+sae_lens/synthetic/hierarchy.py,sha256=j9-6K7xq6zQS9N8bB5nK_-EbuzAZsY5Z5AfUK-qlB5M,22138
 sae_lens/synthetic/initialization.py,sha256=orMGW-786wRDHIS2W7bEH0HmlVFQ4g2z4bnnwdv5w4s,1386
 sae_lens/synthetic/plotting.py,sha256=5lFrej1QOkGAcImFNo5-o-8mI_rUVqvEI57KzUQPPtQ,8208
 sae_lens/synthetic/training.py,sha256=Bg6NYxdzifq_8g-dJQSZ_z_TXDdGRtEi7tqNDb-gCVc,4986
@@ -46,7 +46,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=oIMoeyEP2IzcPFmRbKUzOAycgEyMcOasGeO_BGVZbc4,4846
-sae_lens-6.28.1.dist-info/METADATA,sha256=OdPVG1dwWoLGqiutKkAJGazfBLLbYQLBUbs_3h58BKg,5633
-sae_lens-6.28.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
-sae_lens-6.28.1.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.28.1.dist-info/RECORD,,
+sae_lens-6.28.2.dist-info/METADATA,sha256=i_kbAa64It0NRDrnSlmwNa8qgqOyEMntT_Ifxdx4Q90,6573
+sae_lens-6.28.2.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+sae_lens-6.28.2.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.28.2.dist-info/RECORD,,