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

sae_lens/__init__.py

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

sae_lens/synthetic/__init__.py

@@ -17,11 +17,14 @@ from sae_lens.synthetic.activation_generator import (
     ActivationGenerator,
     ActivationsModifier,
     ActivationsModifierInput,
+    CorrelationMatrixInput,
 )
 from sae_lens.synthetic.correlation import (
+    LowRankCorrelationMatrix,
     create_correlation_matrix_from_correlations,
     generate_random_correlation_matrix,
     generate_random_correlations,
+    generate_random_low_rank_correlation_matrix,
 )
 from sae_lens.synthetic.evals import (
     SyntheticDataEvalResult,
@@ -66,6 +69,9 @@ __all__ = [
     "create_correlation_matrix_from_correlations",
     "generate_random_correlations",
     "generate_random_correlation_matrix",
+    "generate_random_low_rank_correlation_matrix",
+    "LowRankCorrelationMatrix",
+    "CorrelationMatrixInput",
     # Feature modifiers
     "ActivationsModifier",
     "ActivationsModifierInput",

sae_lens/synthetic/activation_generator.py

@@ -2,17 +2,21 @@
 Functions for generating synthetic feature activations.
 """
 
+import math
 from collections.abc import Callable, Sequence
 
 import torch
-from scipy.stats import norm
 from torch import nn
 from torch.distributions import MultivariateNormal
 
+from sae_lens.synthetic.correlation import LowRankCorrelationMatrix
 from sae_lens.util import str_to_dtype
 
 ActivationsModifier = Callable[[torch.Tensor], torch.Tensor]
 ActivationsModifierInput = ActivationsModifier | Sequence[ActivationsModifier] | None
+CorrelationMatrixInput = (
+    torch.Tensor | LowRankCorrelationMatrix | tuple[torch.Tensor, torch.Tensor]
+)
 
 
 class ActivationGenerator(nn.Module):
@@ -28,7 +32,9 @@ class ActivationGenerator(nn.Module):
     mean_firing_magnitudes: torch.Tensor
     modify_activations: ActivationsModifier | None
     correlation_matrix: torch.Tensor | None
+    low_rank_correlation: tuple[torch.Tensor, torch.Tensor] | None
     correlation_thresholds: torch.Tensor | None
+    use_sparse_tensors: bool
 
     def __init__(
         self,
@@ -37,10 +43,37 @@ class ActivationGenerator(nn.Module):
         std_firing_magnitudes: torch.Tensor | float = 0.0,
         mean_firing_magnitudes: torch.Tensor | float = 1.0,
         modify_activations: ActivationsModifierInput = None,
-        correlation_matrix: torch.Tensor | None = None,
+        correlation_matrix: CorrelationMatrixInput | None = None,
         device: torch.device | str = "cpu",
         dtype: torch.dtype | str = "float32",
+        use_sparse_tensors: bool = False,
     ):
+        """
+        Create a new ActivationGenerator.
+
+        Args:
+            num_features: Number of features to generate activations for.
+            firing_probabilities: Probability of each feature firing. Can be a single
+                float (applied to all features) or a tensor of shape (num_features,).
+            std_firing_magnitudes: Standard deviation of firing magnitudes. Can be a
+                single float or a tensor of shape (num_features,). Defaults to 0.0
+                (deterministic magnitudes).
+            mean_firing_magnitudes: Mean firing magnitude when a feature fires. Can be
+                a single float or a tensor of shape (num_features,). Defaults to 1.0.
+            modify_activations: Optional function(s) to modify activations after
+                generation. Can be a single callable, a sequence of callables (applied
+                in order), or None. Useful for applying hierarchy constraints.
+            correlation_matrix: Optional correlation structure between features. Can be:
+
+                - A full correlation matrix tensor of shape (num_features, num_features)
+                - A LowRankCorrelationMatrix for memory-efficient large-scale correlations
+                - A tuple of (factor, diag) tensors representing low-rank structure
+
+            device: Device to place tensors on. Defaults to "cpu".
+            dtype: Data type for tensors. Defaults to "float32".
+            use_sparse_tensors: If True, return sparse COO tensors from sample().
+                Only recommended when using massive numbers of features. Defaults to False.
+        """
         super().__init__()
         self.num_features = num_features
         self.firing_probabilities = _to_tensor(
@@ -54,14 +87,34 @@ class ActivationGenerator(nn.Module):
         )
         self.modify_activations = _normalize_modifiers(modify_activations)
         self.correlation_thresholds = None
+        self.correlation_matrix = None
+        self.low_rank_correlation = None
+        self.use_sparse_tensors = use_sparse_tensors
+
         if correlation_matrix is not None:
-            _validate_correlation_matrix(correlation_matrix, num_features)
-            self.correlation_thresholds = torch.tensor(
-                [norm.ppf(1 - p.item()) for p in self.firing_probabilities],
-                device=device,
-                dtype=self.firing_probabilities.dtype,
+            if isinstance(correlation_matrix, torch.Tensor):
+                # Full correlation matrix
+                _validate_correlation_matrix(correlation_matrix, num_features)
+                self.correlation_matrix = correlation_matrix
+            else:
+                # Low-rank correlation (tuple or LowRankCorrelationMatrix)
+                correlation_factor, correlation_diag = (
+                    correlation_matrix[0],
+                    correlation_matrix[1],
+                )
+                _validate_low_rank_correlation(
+                    correlation_factor, correlation_diag, num_features
+                )
+                # Pre-compute sqrt for efficiency (used every sample call)
+                self.low_rank_correlation = (
+                    correlation_factor,
+                    correlation_diag.sqrt(),
+                )
+
+            # Vectorized inverse normal CDF: norm.ppf(1-p) = sqrt(2) * erfinv(1 - 2*p)
+            self.correlation_thresholds = math.sqrt(2) * torch.erfinv(
+                1 - 2 * self.firing_probabilities
             )
-        self.correlation_matrix = correlation_matrix
 
     @torch.no_grad()
     def sample(self, batch_size: int) -> torch.Tensor:
@@ -84,30 +137,74 @@ class ActivationGenerator(nn.Module):
 
         if self.correlation_matrix is not None:
             assert self.correlation_thresholds is not None
-            firing_features = _generate_correlated_features(
+            firing_indices = _generate_correlated_features(
                 batch_size,
                 self.correlation_matrix,
                 self.correlation_thresholds,
                 device,
             )
+        elif self.low_rank_correlation is not None:
+            assert self.correlation_thresholds is not None
+            firing_indices = _generate_low_rank_correlated_features(
+                batch_size,
+                self.low_rank_correlation[0],
+                self.low_rank_correlation[1],
+                self.correlation_thresholds,
+                device,
+            )
         else:
-            firing_features = torch.bernoulli(
+            firing_indices = torch.bernoulli(
                 self.firing_probabilities.unsqueeze(0).expand(batch_size, -1)
+            ).nonzero(as_tuple=True)
+
+        # Compute activations only at firing positions (sparse optimization)
+        feature_indices = firing_indices[1]
+        num_firing = feature_indices.shape[0]
+        mean_at_firing = self.mean_firing_magnitudes[feature_indices]
+        std_at_firing = self.std_firing_magnitudes[feature_indices]
+        random_deltas = (
+            torch.randn(
+                num_firing, device=device, dtype=self.mean_firing_magnitudes.dtype
             )
-
-        firing_magnitude_delta = torch.normal(
-            torch.zeros_like(self.firing_probabilities)
-            .unsqueeze(0)
-            .expand(batch_size, -1),
-            self.std_firing_magnitudes.unsqueeze(0).expand(batch_size, -1),
+            * std_at_firing
         )
-        firing_magnitude_delta[firing_features == 0] = 0
-        feature_activations = (
-            firing_features * self.mean_firing_magnitudes + firing_magnitude_delta
-        ).relu()
+        activations_at_firing = (mean_at_firing + random_deltas).relu()
+
+        if self.use_sparse_tensors:
+            # Return sparse COO tensor
+            indices = torch.stack(firing_indices)  # [2, nnz]
+            feature_activations = torch.sparse_coo_tensor(
+                indices,
+                activations_at_firing,
+                size=(batch_size, self.num_features),
+                device=device,
+                dtype=self.mean_firing_magnitudes.dtype,
+            )
+        else:
+            # Dense tensor path
+            feature_activations = torch.zeros(
+                batch_size,
+                self.num_features,
+                device=device,
+                dtype=self.mean_firing_magnitudes.dtype,
+            )
+            feature_activations[firing_indices] = activations_at_firing
 
         if self.modify_activations is not None:
-            feature_activations = self.modify_activations(feature_activations).relu()
+            feature_activations = self.modify_activations(feature_activations)
+            if feature_activations.is_sparse:
+                # Apply relu to sparse values
+                feature_activations = feature_activations.coalesce()
+                feature_activations = torch.sparse_coo_tensor(
+                    feature_activations.indices(),
+                    feature_activations.values().relu(),
+                    feature_activations.shape,
+                    device=feature_activations.device,
+                    dtype=feature_activations.dtype,
+                )
+            else:
+                feature_activations = feature_activations.relu()
+
         return feature_activations
 
     def forward(self, batch_size: int) -> torch.Tensor:
@@ -119,7 +216,7 @@ def _generate_correlated_features(
     correlation_matrix: torch.Tensor,
     thresholds: torch.Tensor,
     device: torch.device,
-) -> torch.Tensor:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Generate correlated binary features using multivariate Gaussian sampling.
 
@@ -133,7 +230,7 @@ def _generate_correlated_features(
         device: Device to generate samples on
 
     Returns:
-        Binary feature matrix of shape [batch_size, num_features]
+        Tuple of (row_indices, col_indices) for firing features
     """
     num_features = correlation_matrix.shape[0]
 
@@ -143,7 +240,49 @@ def _generate_correlated_features(
     )
 
     gaussian_samples = mvn.sample((batch_size,))
-    return (gaussian_samples > thresholds.unsqueeze(0)).float()
+    indices = (gaussian_samples > thresholds.unsqueeze(0)).nonzero(as_tuple=True)
+    return indices[0], indices[1]
+
+
+def _generate_low_rank_correlated_features(
+    batch_size: int,
+    correlation_factor: torch.Tensor,
+    cov_diag_sqrt: torch.Tensor,
+    thresholds: torch.Tensor,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate correlated binary features using low-rank multivariate Gaussian sampling.
+
+    Uses the Gaussian copula approach with a low-rank covariance structure for scalability.
+    The covariance is represented as: cov = factor @ factor.T + diag(diag_term)
+
+    Args:
+        batch_size: Number of samples to generate
+        correlation_factor: Factor matrix of shape (num_features, rank)
+        cov_diag_sqrt: Pre-computed sqrt of diagonal term, shape (num_features,)
+        thresholds: Pre-computed thresholds for each feature (from inverse normal CDF)
+        device: Device to generate samples on
+
+    Returns:
+        Tuple of (row_indices, col_indices) for firing features
+    """
+    # Manual low-rank MVN sampling to enable autocast for the expensive matmul
+    # samples = eps @ cov_factor.T + eta * sqrt(cov_diag)
+    # where eps ~ N(0, I_rank) and eta ~ N(0, I_n)
+
+    num_features, rank = correlation_factor.shape
+
+    # Generate random samples in float32 for numerical stability
+    eps = torch.randn(batch_size, rank, device=device, dtype=correlation_factor.dtype)
+    eta = torch.randn(
+        batch_size, num_features, device=device, dtype=cov_diag_sqrt.dtype
+    )
+
+    gaussian_samples = eps @ correlation_factor.T + eta * cov_diag_sqrt
+
+    indices = (gaussian_samples > thresholds.unsqueeze(0)).nonzero(as_tuple=True)
+    return indices[0], indices[1]
 
 
 def _to_tensor(
@@ -194,7 +333,7 @@ def _validate_correlation_matrix(
 
     Args:
         correlation_matrix: The matrix to validate
-        num_features: Expected number of features (matrix should be [num_features, num_features])
+        num_features: Expected number of features (matrix should be (num_features, num_features))
 
     Raises:
         ValueError: If the matrix has incorrect shape, non-unit diagonal, or is not positive definite
@@ -214,3 +353,36 @@ def _validate_correlation_matrix(
         torch.linalg.cholesky(correlation_matrix)
     except RuntimeError as e:
         raise ValueError("Correlation matrix must be positive definite") from e
+
+
+def _validate_low_rank_correlation(
+    correlation_factor: torch.Tensor,
+    correlation_diag: torch.Tensor,
+    num_features: int,
+) -> None:
+    """Validate that low-rank correlation parameters have correct properties.
+
+    Args:
+        correlation_factor: Factor matrix of shape (num_features, rank)
+        correlation_diag: Diagonal term of shape (num_features,)
+        num_features: Expected number of features
+
+    Raises:
+        ValueError: If shapes are incorrect or diagonal terms are not positive
+    """
+    if correlation_factor.ndim != 2:
+        raise ValueError(
+            f"correlation_factor must be 2D, got {correlation_factor.ndim}D"
+        )
+    if correlation_factor.shape[0] != num_features:
+        raise ValueError(
+            f"correlation_factor must have shape ({num_features}, rank), "
+            f"got {tuple(correlation_factor.shape)}"
+        )
+    if correlation_diag.shape != (num_features,):
+        raise ValueError(
+            f"correlation_diag must have shape ({num_features},), "
+            f"got {tuple(correlation_diag.shape)}"
+        )
+    if torch.any(correlation_diag <= 0):
+        raise ValueError("correlation_diag must have all positive values")

sae_lens/synthetic/correlation.py

@@ -1,7 +1,32 @@
 import random
+from typing import NamedTuple
 
 import torch
 
+from sae_lens.util import str_to_dtype
+
+
+class LowRankCorrelationMatrix(NamedTuple):
+    """
+    Low-rank representation of a correlation matrix for scalable correlated sampling.
+
+    The correlation structure is represented as:
+        correlation = correlation_factor @ correlation_factor.T + diag(correlation_diag)
+
+    This requires O(num_features * rank) storage instead of O(num_features^2),
+    making it suitable for very large numbers of features (e.g., 1M+).
+
+    Attributes:
+        correlation_factor: Factor matrix of shape (num_features, rank) that captures
+            correlations through shared latent factors.
+        correlation_diag: Diagonal variance term of shape (num_features,). Should be
+            chosen such that the diagonal of the full correlation matrix equals 1.
+            Typically: correlation_diag[i] = 1 - sum(correlation_factor[i, :]^2)
+    """
+
+    correlation_factor: torch.Tensor
+    correlation_diag: torch.Tensor
+
 
 def create_correlation_matrix_from_correlations(
     num_features: int,
@@ -11,14 +36,18 @@ def create_correlation_matrix_from_correlations(
     """
     Create a correlation matrix with specified pairwise correlations.
 
+    Note: If the resulting matrix is not positive definite, it will be adjusted
+    to ensure validity. This adjustment may change the specified correlation
+    values. To minimize this effect, use smaller correlation magnitudes.
+
     Args:
         num_features: Number of features
         correlations: Dict mapping (i, j) pairs to correlation values.
-            Pairs should have i < j.
+            Pairs should have i < j. Pairs not specified will use default_correlation.
         default_correlation: Default correlation for unspecified pairs
 
     Returns:
-        Correlation matrix of shape [num_features, num_features]
+        Correlation matrix of shape (num_features, num_features)
     """
     matrix = torch.eye(num_features) + default_correlation * (
         1 - torch.eye(num_features)
@@ -50,6 +79,7 @@ def _fix_correlation_matrix(
     fixed_matrix = eigenvecs @ torch.diag(eigenvals) @ eigenvecs.T
 
     diag_vals = torch.diag(fixed_matrix)
+    diag_vals = torch.clamp(diag_vals, min=1e-8)  # Prevent division by zero
     fixed_matrix = fixed_matrix / torch.sqrt(
         diag_vals.unsqueeze(0) * diag_vals.unsqueeze(1)
     )
@@ -58,6 +88,25 @@ def _fix_correlation_matrix(
     return fixed_matrix
 
 
+def _validate_correlation_params(
+    positive_ratio: float,
+    uncorrelated_ratio: float,
+    min_correlation_strength: float,
+    max_correlation_strength: float,
+) -> None:
+    """Validate parameters for correlation generation."""
+    if not 0.0 <= positive_ratio <= 1.0:
+        raise ValueError("positive_ratio must be between 0.0 and 1.0")
+    if not 0.0 <= uncorrelated_ratio <= 1.0:
+        raise ValueError("uncorrelated_ratio must be between 0.0 and 1.0")
+    if min_correlation_strength < 0:
+        raise ValueError("min_correlation_strength must be non-negative")
+    if max_correlation_strength > 1.0:
+        raise ValueError("max_correlation_strength must be <= 1.0")
+    if min_correlation_strength > max_correlation_strength:
+        raise ValueError("min_correlation_strength must be <= max_correlation_strength")
+
+
 def generate_random_correlations(
     num_features: int,
     positive_ratio: float = 0.5,
@@ -71,29 +120,26 @@ def generate_random_correlations(
 
     Args:
         num_features: Number of features
-        positive_ratio: Fraction of correlations that should be positive (0.0 to 1.0)
-        uncorrelated_ratio: Fraction of feature pairs that should remain uncorrelated (0.0 to 1.0)
-        min_correlation_strength: Minimum absolute correlation strength
-        max_correlation_strength: Maximum absolute correlation strength
+        positive_ratio: Fraction of correlated pairs that should be positive (0.0 to 1.0)
+        uncorrelated_ratio: Fraction of feature pairs that should have zero correlation
+            (0.0 to 1.0). These pairs are omitted from the returned dictionary.
+        min_correlation_strength: Minimum absolute correlation strength for correlated pairs
+        max_correlation_strength: Maximum absolute correlation strength for correlated pairs
         seed: Random seed for reproducibility
 
     Returns:
-        Dictionary mapping (i, j) pairs to correlation values
+        Dictionary mapping (i, j) pairs to correlation values. Pairs with zero
+        correlation (determined by uncorrelated_ratio) are not included.
     """
     # Use local random number generator to avoid side effects on global state
     rng = random.Random(seed)
 
-    # Validate inputs
-    if not 0.0 <= positive_ratio <= 1.0:
-        raise ValueError("positive_ratio must be between 0.0 and 1.0")
-    if not 0.0 <= uncorrelated_ratio <= 1.0:
-        raise ValueError("uncorrelated_ratio must be between 0.0 and 1.0")
-    if min_correlation_strength < 0:
-        raise ValueError("min_correlation_strength must be non-negative")
-    if max_correlation_strength > 1.0:
-        raise ValueError("max_correlation_strength must be <= 1.0")
-    if min_correlation_strength > max_correlation_strength:
-        raise ValueError("min_correlation_strength must be <= max_correlation_strength")
+    _validate_correlation_params(
+        positive_ratio,
+        uncorrelated_ratio,
+        min_correlation_strength,
+        max_correlation_strength,
+    )
 
     # Generate all possible feature pairs (i, j) where i < j
     all_pairs = [
@@ -136,35 +182,170 @@ def generate_random_correlation_matrix(
     min_correlation_strength: float = 0.1,
     max_correlation_strength: float = 0.8,
     seed: int | None = None,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype | str = torch.float32,
 ) -> torch.Tensor:
     """
     Generate a random correlation matrix with specified constraints.
 
-    This is a convenience function that combines generate_random_correlations()
-    and create_correlation_matrix_from_correlations() into a single call.
+    Uses vectorized torch operations for efficiency with large numbers of features.
+
+    Note: If the randomly generated matrix is not positive definite, it will be
+    adjusted to ensure validity. This adjustment may change correlation values,
+    including turning some zero correlations into non-zero values. To minimize
+    this effect, use smaller correlation strengths (e.g., 0.01-0.1).
+
+    Args:
+        num_features: Number of features
+        positive_ratio: Fraction of correlated pairs that should be positive (0.0 to 1.0)
+        uncorrelated_ratio: Fraction of feature pairs that should have zero correlation
+            (0.0 to 1.0). Note that matrix fixing for positive definiteness may reduce
+            the actual number of zero correlations.
+        min_correlation_strength: Minimum absolute correlation strength for correlated pairs
+        max_correlation_strength: Maximum absolute correlation strength for correlated pairs
+        seed: Random seed for reproducibility
+        device: Device to create the matrix on
+        dtype: Data type for the matrix
+
+    Returns:
+        Random correlation matrix of shape (num_features, num_features)
+    """
+    dtype = str_to_dtype(dtype)
+    _validate_correlation_params(
+        positive_ratio,
+        uncorrelated_ratio,
+        min_correlation_strength,
+        max_correlation_strength,
+    )
+
+    if num_features <= 1:
+        return torch.eye(num_features, device=device, dtype=dtype)
+
+    # Set random seed if provided
+    generator = torch.Generator(device=device)
+    if seed is not None:
+        generator.manual_seed(seed)
+
+    # Get upper triangular indices (i < j)
+    row_idx, col_idx = torch.triu_indices(num_features, num_features, offset=1)
+    num_pairs = row_idx.shape[0]
+
+    # Generate random values for all pairs at once
+    # is_correlated: 1 if this pair should have a correlation, 0 otherwise
+    is_correlated = (
+        torch.rand(num_pairs, generator=generator, device=device) >= uncorrelated_ratio
+    )
+
+    # signs: +1 for positive correlation, -1 for negative
+    is_positive = (
+        torch.rand(num_pairs, generator=generator, device=device) < positive_ratio
+    )
+    signs = torch.where(is_positive, 1.0, -1.0)
+
+    # strengths: uniform in [min_strength, max_strength]
+    strengths = (
+        torch.rand(num_pairs, generator=generator, device=device, dtype=dtype)
+        * (max_correlation_strength - min_correlation_strength)
+        + min_correlation_strength
+    )
+
+    # Combine: correlation = is_correlated * sign * strength
+    correlations = is_correlated.to(dtype) * signs.to(dtype) * strengths
+
+    # Build the symmetric matrix
+    matrix = torch.eye(num_features, device=device, dtype=dtype)
+    matrix[row_idx, col_idx] = correlations
+    matrix[col_idx, row_idx] = correlations
+
+    # Check positive definiteness and fix if necessary
+    eigenvals = torch.linalg.eigvalsh(matrix)
+    if torch.any(eigenvals < -1e-6):
+        matrix = _fix_correlation_matrix(matrix)
+
+    return matrix
+
+
+def generate_random_low_rank_correlation_matrix(
+    num_features: int,
+    rank: int,
+    correlation_scale: float = 0.1,
+    seed: int | None = None,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype | str = torch.float32,
+) -> LowRankCorrelationMatrix:
+    """
+    Generate a random low-rank correlation structure for scalable correlated sampling.
+
+    The correlation structure is represented as:
+        correlation = factor @ factor.T + diag(diag_term)
+
+    This requires O(num_features * rank) storage instead of O(num_features^2),
+    making it suitable for very large numbers of features (e.g., 1M+).
+
+    The factor matrix is initialized with random values scaled by correlation_scale,
+    and the diagonal term is computed to ensure the implied correlation matrix has
+    unit diagonal.
 
     Args:
         num_features: Number of features
-        positive_ratio: Fraction of correlations that should be positive (0.0 to 1.0)
-        uncorrelated_ratio: Fraction of feature pairs that should remain uncorrelated (0.0 to 1.0)
-        min_correlation_strength: Minimum absolute correlation strength
-        max_correlation_strength: Maximum absolute correlation strength
+        rank: Rank of the low-rank approximation. Higher rank allows more complex
+            correlation structures but uses more memory. Typical values: 10-100.
+        correlation_scale: Scale factor for random correlations. Larger values produce
+            stronger correlations between features. Use 0 for no correlations (identity
+            matrix). Should be small enough that rank * correlation_scale^2 < 1 to
+            ensure valid diagonal terms.
         seed: Random seed for reproducibility
+        device: Device to create tensors on
+        dtype: Data type for tensors
 
     Returns:
-        Random correlation matrix of shape [num_features, num_features]
+        LowRankCorrelationMatrix containing the factor matrix and diagonal term
     """
-    # Generate random correlations
-    correlations = generate_random_correlations(
-        num_features=num_features,
-        positive_ratio=positive_ratio,
-        uncorrelated_ratio=uncorrelated_ratio,
-        min_correlation_strength=min_correlation_strength,
-        max_correlation_strength=max_correlation_strength,
-        seed=seed,
+    # Minimum diagonal value to ensure numerical stability in the covariance matrix.
+    # This limits how much variance can come from the low-rank factor.
+    _MIN_DIAG = 0.01
+
+    dtype = str_to_dtype(dtype)
+    device = torch.device(device)
+
+    if rank <= 0:
+        raise ValueError("rank must be positive")
+    if correlation_scale < 0:
+        raise ValueError("correlation_scale must be non-negative")
+
+    # Set random seed if provided
+    generator = torch.Generator(device=device)
+    if seed is not None:
+        generator.manual_seed(seed)
+
+    # Generate random factor matrix
+    # Each row has norm roughly sqrt(rank) * correlation_scale
+    factor = (
+        torch.randn(num_features, rank, generator=generator, device=device, dtype=dtype)
+        * correlation_scale
     )
 
-    # Create and return correlation matrix
-    return create_correlation_matrix_from_correlations(
-        num_features=num_features, correlations=correlations
+    # Compute diagonal term to ensure unit diagonal in implied correlation matrix
+    # diag(factor @ factor.T) + diag_term = 1
+    # diag_term = 1 - sum(factor[i, :]^2)
+    factor_sq_sum = (factor**2).sum(dim=1)
+    diag_term = 1 - factor_sq_sum
+
+    # Ensure diagonal terms are at least _MIN_DIAG for numerical stability
+    # If any diagonal term is too small, scale down the factor matrix
+    if torch.any(diag_term < _MIN_DIAG):
+        # Scale factor so max row norm squared is at most (1 - _MIN_DIAG)
+        # This ensures all diagonal terms are >= _MIN_DIAG
+        max_factor_contribution = 1 - _MIN_DIAG
+        max_sq_sum = factor_sq_sum.max()
+        scale = torch.sqrt(
+            torch.tensor(max_factor_contribution, device=device, dtype=dtype)
+            / max_sq_sum
+        )
+        factor = factor * scale
+        factor_sq_sum = (factor**2).sum(dim=1)
+        diag_term = 1 - factor_sq_sum
+
+    return LowRankCorrelationMatrix(
+        correlation_factor=factor, correlation_diag=diag_term
     )

sae_lens/synthetic/feature_dictionary.py

@@ -9,7 +9,7 @@ from typing import Callable
 
 import torch
 from torch import nn
-from tqdm import tqdm
+from tqdm.auto import tqdm
 
 FeatureDictionaryInitializer = Callable[["FeatureDictionary"], None]
 
@@ -168,9 +168,18 @@ class FeatureDictionary(nn.Module):
 
         Args:
             feature_activations: Tensor of shape [batch, num_features] containing
-                sparse feature activation values
+                sparse feature activation values. Can be dense or sparse COO.
 
         Returns:
             Tensor of shape [batch, hidden_dim] containing dense hidden activations
         """
+        if feature_activations.is_sparse:
+            # autocast is disabled here because sparse matmul is not supported with bfloat16
+            with torch.autocast(
+                device_type=feature_activations.device.type, enabled=False
+            ):
+                return (
+                    torch.sparse.mm(feature_activations, self.feature_vectors)
+                    + self.bias
+                )
         return feature_activations @ self.feature_vectors + self.bias

sae_lens/synthetic/hierarchy.py

@@ -147,6 +147,14 @@ class _SparseHierarchyData:
     # Total number of ME groups
     num_groups: int
 
+    # Sparse COO support: Feature-to-parent mapping
+    # feat_to_parent[f] = parent feature index, or -1 if root/no parent
+    feat_to_parent: torch.Tensor | None = None  # [num_features]
+
+    # Sparse COO support: Feature-to-ME-group mapping
+    # feat_to_me_group[f] = group index, or -1 if not in any ME group
+    feat_to_me_group: torch.Tensor | None = None  # [num_features]
+
 
 def _build_sparse_hierarchy(
     roots: Sequence[HierarchyNode],
@@ -232,7 +240,11 @@ def _build_sparse_hierarchy(
             me_indices = torch.empty(0, dtype=torch.long)
 
         level_data.append(
-            _LevelData(features=feats, parents=parents, me_group_indices=me_indices)
+            _LevelData(
+                features=feats,
+                parents=parents,
+                me_group_indices=me_indices,
+            )
         )
 
     # Build group siblings and parents tensors
@@ -254,12 +266,30 @@ def _build_sparse_hierarchy(
         me_group_parents = torch.empty(0, dtype=torch.long)
         num_groups = 0
 
+    # Build sparse COO support: feat_to_parent and feat_to_me_group mappings
+    # First determine num_features (max feature index + 1)
+    all_features = [f for f, _, _ in feature_info]
+    num_features = max(all_features) + 1 if all_features else 0
+
+    # Build feature-to-parent mapping
+    feat_to_parent = torch.full((num_features,), -1, dtype=torch.long)
+    for feat, parent, _ in feature_info:
+        feat_to_parent[feat] = parent
+
+    # Build feature-to-ME-group mapping
+    feat_to_me_group = torch.full((num_features,), -1, dtype=torch.long)
+    for g_idx, (_, _, siblings) in enumerate(me_groups):
+        for sib in siblings:
+            feat_to_me_group[sib] = g_idx
+
     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,
+        feat_to_parent=feat_to_parent,
+        feat_to_me_group=feat_to_me_group,
     )
 
 
@@ -396,8 +426,9 @@ def _apply_me_for_groups(
     # 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.
+    _INACTIVE_SCORE = -1e9
     random_scores = torch.rand(num_conflicts, max_siblings, device=device)
-    random_scores[~conflict_active] = -1e9
+    random_scores[~conflict_active] = _INACTIVE_SCORE
 
     winner_idx = random_scores.argmax(dim=1)
 
@@ -420,6 +451,275 @@ def _apply_me_for_groups(
     activations[deact_batch, deact_feat] = 0
 
 
+# ---------------------------------------------------------------------------
+# Sparse COO hierarchy implementation
+# ---------------------------------------------------------------------------
+
+
+def _apply_hierarchy_sparse_coo(
+    sparse_tensor: torch.Tensor,
+    sparse_data: _SparseHierarchyData,
+) -> torch.Tensor:
+    """
+    Apply hierarchy constraints to a sparse COO tensor.
+
+    This is the sparse analog of _apply_hierarchy_sparse. It processes
+    level-by-level, applying parent deactivation then mutual exclusion.
+    """
+    if sparse_tensor._nnz() == 0:
+        return sparse_tensor
+
+    sparse_tensor = sparse_tensor.coalesce()
+
+    for level_data in sparse_data.level_data:
+        # Step 1: Apply parent deactivation for features at this level
+        if level_data.features.numel() > 0:
+            sparse_tensor = _apply_parent_deactivation_coo(
+                sparse_tensor, level_data, sparse_data
+            )
+
+        # Step 2: Apply ME for groups whose parent is at this level
+        if level_data.me_group_indices.numel() > 0:
+            sparse_tensor = _apply_me_coo(
+                sparse_tensor, level_data.me_group_indices, sparse_data
+            )
+
+    return sparse_tensor
+
+
+def _apply_parent_deactivation_coo(
+    sparse_tensor: torch.Tensor,
+    level_data: _LevelData,
+    sparse_data: _SparseHierarchyData,
+) -> torch.Tensor:
+    """
+    Remove children from sparse COO tensor when their parent is inactive.
+
+    Uses searchsorted for efficient membership testing of parent activity.
+    """
+    if sparse_tensor._nnz() == 0 or level_data.features.numel() == 0:
+        return sparse_tensor
+
+    sparse_tensor = sparse_tensor.coalesce()
+    indices = sparse_tensor.indices()  # [2, nnz]
+    values = sparse_tensor.values()  # [nnz]
+    batch_indices = indices[0]
+    feat_indices = indices[1]
+
+    _, num_features = sparse_tensor.shape
+    device = sparse_tensor.device
+    nnz = indices.shape[1]
+
+    # Build set of active (batch, feature) pairs for efficient lookup
+    # Encode as: batch_idx * num_features + feat_idx
+    active_pairs = batch_indices * num_features + feat_indices
+    active_pairs_sorted, _ = active_pairs.sort()
+
+    # Use the precomputed feat_to_parent mapping
+    assert sparse_data.feat_to_parent is not None
+    hierarchy_num_features = sparse_data.feat_to_parent.numel()
+
+    # Handle features outside the hierarchy (they have no parent, pass through)
+    in_hierarchy = feat_indices < hierarchy_num_features
+    parent_of_feat = torch.full((nnz,), -1, dtype=torch.long, device=device)
+    parent_of_feat[in_hierarchy] = sparse_data.feat_to_parent[
+        feat_indices[in_hierarchy]
+    ]
+
+    # Find entries that have a parent (parent >= 0 means this feature has a parent)
+    has_parent = parent_of_feat >= 0
+
+    if not has_parent.any():
+        return sparse_tensor
+
+    # For entries with parents, check if parent is active
+    child_entry_indices = torch.where(has_parent)[0]
+    child_batch = batch_indices[has_parent]
+    child_parents = parent_of_feat[has_parent]
+
+    # Look up parent activity using searchsorted
+    parent_pairs = child_batch * num_features + child_parents
+    search_pos = torch.searchsorted(active_pairs_sorted, parent_pairs)
+    search_pos = search_pos.clamp(max=active_pairs_sorted.numel() - 1)
+    parent_active = active_pairs_sorted[search_pos] == parent_pairs
+
+    # Handle empty case
+    if active_pairs_sorted.numel() == 0:
+        parent_active = torch.zeros_like(parent_pairs, dtype=torch.bool)
+
+    # Build keep mask: keep entry if it's a root OR its parent is active
+    keep_mask = torch.ones(nnz, dtype=torch.bool, device=device)
+    keep_mask[child_entry_indices[~parent_active]] = False
+
+    if keep_mask.all():
+        return sparse_tensor
+
+    return torch.sparse_coo_tensor(
+        indices[:, keep_mask],
+        values[keep_mask],
+        sparse_tensor.shape,
+        device=device,
+        dtype=sparse_tensor.dtype,
+    )
+
+
+def _apply_me_coo(
+    sparse_tensor: torch.Tensor,
+    group_indices: torch.Tensor,
+    sparse_data: _SparseHierarchyData,
+) -> torch.Tensor:
+    """
+    Apply mutual exclusion to sparse COO tensor.
+
+    For each ME group with multiple active siblings in the same batch,
+    randomly selects one winner and removes the rest.
+    """
+    if sparse_tensor._nnz() == 0 or group_indices.numel() == 0:
+        return sparse_tensor
+
+    sparse_tensor = sparse_tensor.coalesce()
+    indices = sparse_tensor.indices()  # [2, nnz]
+    values = sparse_tensor.values()  # [nnz]
+    batch_indices = indices[0]
+    feat_indices = indices[1]
+
+    _, num_features = sparse_tensor.shape
+    device = sparse_tensor.device
+    nnz = indices.shape[1]
+
+    # Use precomputed feat_to_me_group mapping
+    assert sparse_data.feat_to_me_group is not None
+    hierarchy_num_features = sparse_data.feat_to_me_group.numel()
+
+    # Handle features outside the hierarchy (they are not in any ME group)
+    in_hierarchy = feat_indices < hierarchy_num_features
+    me_group_of_feat = torch.full((nnz,), -1, dtype=torch.long, device=device)
+    me_group_of_feat[in_hierarchy] = sparse_data.feat_to_me_group[
+        feat_indices[in_hierarchy]
+    ]
+
+    # Find entries that belong to ME groups we're processing (vectorized)
+    in_relevant_group = torch.isin(me_group_of_feat, group_indices)
+
+    if not in_relevant_group.any():
+        return sparse_tensor
+
+    # Get the ME entries
+    me_entry_indices = torch.where(in_relevant_group)[0]
+    me_batch = batch_indices[in_relevant_group]
+    me_group = me_group_of_feat[in_relevant_group]
+
+    # Check parent activity for ME groups (only apply ME if parent is active)
+    me_group_parents = sparse_data.me_group_parents[me_group]
+    has_parent = me_group_parents >= 0
+
+    if has_parent.any():
+        # Build active pairs for parent lookup
+        active_pairs = batch_indices * num_features + feat_indices
+        active_pairs_sorted, _ = active_pairs.sort()
+
+        parent_pairs = (
+            me_batch[has_parent] * num_features + me_group_parents[has_parent]
+        )
+        search_pos = torch.searchsorted(active_pairs_sorted, parent_pairs)
+        search_pos = search_pos.clamp(max=active_pairs_sorted.numel() - 1)
+        parent_active_for_has_parent = active_pairs_sorted[search_pos] == parent_pairs
+
+        # Build full parent_active mask
+        parent_active = torch.ones(
+            me_entry_indices.numel(), dtype=torch.bool, device=device
+        )
+        parent_active[has_parent] = parent_active_for_has_parent
+
+        # Filter to only ME entries where parent is active
+        valid_me = parent_active
+        me_entry_indices = me_entry_indices[valid_me]
+        me_batch = me_batch[valid_me]
+        me_group = me_group[valid_me]
+
+    if me_entry_indices.numel() == 0:
+        return sparse_tensor
+
+    # Encode (batch, group) pairs
+    num_groups = sparse_data.num_groups
+    batch_group_pairs = me_batch * num_groups + me_group
+
+    # Find unique (batch, group) pairs and count occurrences
+    unique_bg, inverse, counts = torch.unique(
+        batch_group_pairs, return_inverse=True, return_counts=True
+    )
+
+    # Only process pairs with count > 1 (conflicts)
+    has_conflict = counts > 1
+
+    if not has_conflict.any():
+        return sparse_tensor
+
+    # For efficiency, we process all conflicts together
+    # Assign random scores to each ME entry
+    random_scores = torch.rand(me_entry_indices.numel(), device=device)
+
+    # For each (batch, group) pair, we want the entry with highest score to be winner
+    # Use scatter_reduce to find max score per (batch, group)
+    bg_to_dense = torch.zeros(unique_bg.numel(), dtype=torch.long, device=device)
+    bg_to_dense[has_conflict.nonzero(as_tuple=True)[0]] = torch.arange(
+        has_conflict.sum(), device=device
+    )
+
+    # Map each ME entry to its dense conflict index
+    entry_has_conflict = has_conflict[inverse]
+
+    if not entry_has_conflict.any():
+        return sparse_tensor
+
+    conflict_entries_mask = entry_has_conflict
+    conflict_entry_indices = me_entry_indices[conflict_entries_mask]
+    conflict_random_scores = random_scores[conflict_entries_mask]
+    conflict_inverse = inverse[conflict_entries_mask]
+    conflict_dense_idx = bg_to_dense[conflict_inverse]
+
+    # Vectorized winner selection using sorting
+    # Sort entries by (group_idx, -random_score) so highest score comes first per group
+    # Use group * 2 - score to sort by group ascending, then score descending
+    sort_keys = conflict_dense_idx.float() * 2.0 - conflict_random_scores
+    sorted_order = sort_keys.argsort()
+    sorted_dense_idx = conflict_dense_idx[sorted_order]
+
+    # Find first entry of each group in sorted order (these are winners)
+    group_starts = torch.cat(
+        [
+            torch.tensor([True], device=device),
+            sorted_dense_idx[1:] != sorted_dense_idx[:-1],
+        ]
+    )
+
+    # Winners are entries at group starts in sorted order
+    winner_positions_in_sorted = torch.where(group_starts)[0]
+    winner_original_positions = sorted_order[winner_positions_in_sorted]
+
+    # Create winner mask (vectorized)
+    is_winner = torch.zeros(
+        conflict_entry_indices.numel(), dtype=torch.bool, device=device
+    )
+    is_winner[winner_original_positions] = True
+
+    # Build keep mask (vectorized)
+    keep_mask = torch.ones(nnz, dtype=torch.bool, device=device)
+    loser_entry_indices = conflict_entry_indices[~is_winner]
+    keep_mask[loser_entry_indices] = False
+
+    if keep_mask.all():
+        return sparse_tensor
+
+    return torch.sparse_coo_tensor(
+        indices[:, keep_mask],
+        values[keep_mask],
+        sparse_tensor.shape,
+        device=device,
+        dtype=sparse_tensor.dtype,
+    )
+
+
 @torch.no_grad()
 def hierarchy_modifier(
     roots: Sequence[HierarchyNode] | HierarchyNode,
@@ -475,12 +775,24 @@ def hierarchy_modifier(
                 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,
+                feat_to_parent=(
+                    sparse_data.feat_to_parent.to(device)
+                    if sparse_data.feat_to_parent is not None
+                    else None
+                ),
+                feat_to_me_group=(
+                    sparse_data.feat_to_me_group.to(device)
+                    if sparse_data.feat_to_me_group is not None
+                    else None
+                ),
             )
         return device_cache[device]
 
     def modifier(activations: torch.Tensor) -> torch.Tensor:
         device = activations.device
         cached = _get_sparse_for_device(device)
+        if activations.is_sparse:
+            return _apply_hierarchy_sparse_coo(activations, cached)
         return _apply_hierarchy_sparse(activations, cached)
 
     return modifier

sae_lens/synthetic/training.py

@@ -23,6 +23,8 @@ def train_toy_sae(
     device: str | torch.device = "cpu",
     n_snapshots: int = 0,
     snapshot_fn: Callable[[SAETrainer[Any, Any]], None] | None = None,
+    autocast_sae: bool = False,
+    autocast_data: bool = False,
 ) -> None:
     """
     Train an SAE on synthetic activations from a feature dictionary.
@@ -46,6 +48,8 @@ def train_toy_sae(
         snapshot_fn: Callback function called at each snapshot point. Receives
             the SAETrainer instance, allowing access to the SAE, training step,
             and other training state. Required if n_snapshots > 0.
+        autocast_sae: Whether to autocast the SAE to bfloat16. Only recommend for large SAEs on CUDA
+        autocast_data: Whether to autocast the activations generator and feature dictionary to bfloat16. Only recommend for large data on CUDA.
     """
 
     device_str = str(device) if isinstance(device, torch.device) else device
@@ -55,6 +59,7 @@ def train_toy_sae(
         feature_dict=feature_dict,
         activations_generator=activations_generator,
         batch_size=batch_size,
+        autocast=autocast_data,
     )
 
     # Create trainer config
@@ -64,7 +69,7 @@ def train_toy_sae(
         save_final_checkpoint=False,
         total_training_samples=training_samples,
         device=device_str,
-        autocast=False,
+        autocast=autocast_sae,
         lr=lr,
         lr_end=lr,
         lr_scheduler_name="constant",
@@ -119,6 +124,7 @@ class SyntheticActivationIterator(Iterator[torch.Tensor]):
         feature_dict: FeatureDictionary,
         activations_generator: ActivationGenerator,
         batch_size: int,
+        autocast: bool = False,
     ):
         """
         Create a new SyntheticActivationIterator.
@@ -127,16 +133,23 @@ class SyntheticActivationIterator(Iterator[torch.Tensor]):
             feature_dict: The feature dictionary to use for generating hidden activations
             activations_generator: Generator that produces feature activations
             batch_size: Number of samples per batch
+            autocast: Whether to autocast the activations generator and feature dictionary to bfloat16.
         """
         self.feature_dict = feature_dict
         self.activations_generator = activations_generator
         self.batch_size = batch_size
+        self.autocast = autocast
 
     @torch.no_grad()
     def next_batch(self) -> torch.Tensor:
         """Generate the next batch of hidden activations."""
-        features = self.activations_generator(self.batch_size)
-        return self.feature_dict(features)
+        with torch.autocast(
+            device_type=self.feature_dict.feature_vectors.device.type,
+            dtype=torch.bfloat16,
+            enabled=self.autocast,
+        ):
+            features = self.activations_generator(self.batch_size)
+            return self.feature_dict(features)
 
     def __iter__(self) -> "SyntheticActivationIterator":
         return self

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.28.2
+Version: 6.29.1
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE

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

@@ -1,4 +1,4 @@
-sae_lens/__init__.py,sha256=B9tY0Jt21pOHmSQrQLpMxQHyUAdLHIZpVP6pg3O0dfQ,4788
+sae_lens/__init__.py,sha256=emqKVNiJwD8YtYhtgHJyAT8YSX1QmruQYuG-J4CStC4,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
@@ -25,16 +25,16 @@ sae_lens/saes/standard_sae.py,sha256=_hldNZkFPAf9VGrxouR1-tN8T2OEk8IkWBcXoatrC1o
 sae_lens/saes/temporal_sae.py,sha256=83Ap4mYGfdN3sKdPF8nKjhdXph3-7E2QuLobqJ_YuoM,13273
 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=JEN7mEgdGDuXr0ArTwUsSdSVUAfvheT_1Eew2ojbA-g,7659
-sae_lens/synthetic/correlation.py,sha256=odr-S5h6c2U-bepwrAQeMfV1iBF_cnnQzqw7zapEXZ4,6056
+sae_lens/synthetic/__init__.py,sha256=MtTnGkTfHV2WjkIgs7zZyx10EK9U5fjOHXy69Aq3uKw,3095
+sae_lens/synthetic/activation_generator.py,sha256=8L9nwC4jFRv_wg3QN-n1sFwX8w1NqwJMysWaJ41lLlY,15197
+sae_lens/synthetic/correlation.py,sha256=tMTLo9fBfDpeXwqhyUgFqnTipj9x2W0t4oEtNxB7AG0,13256
 sae_lens/synthetic/evals.py,sha256=Nhi314ZnRgLfhBj-3tm_zzI-pGyFTcwllDXbIpPFXeU,4584
-sae_lens/synthetic/feature_dictionary.py,sha256=ysn0ihE3JgVlCLUZMb127WYZqbz4kMp9BGHfCZqERBg,6487
+sae_lens/synthetic/feature_dictionary.py,sha256=Nd4xjSTxKMnKilZ3uYi8Gv5SS5D4bv4wHiSL1uGB69E,6933
 sae_lens/synthetic/firing_probabilities.py,sha256=yclz1pWl5gE1r8LAxFvzQS88Lxwk5-3r8BCX9HLVejA,3370
-sae_lens/synthetic/hierarchy.py,sha256=j9-6K7xq6zQS9N8bB5nK_-EbuzAZsY5Z5AfUK-qlB5M,22138
+sae_lens/synthetic/hierarchy.py,sha256=nm7nwnTswktVJeKUsRZ0hLOdXcFWGbxnA1b6lefHm-4,33592
 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
+sae_lens/synthetic/training.py,sha256=fHcX2cZ6nDupr71GX0Gk17f1NvQ0SKIVXIA6IuAb2dw,5692
 sae_lens/tokenization_and_batching.py,sha256=uoHtAs9z3XqG0Fh-iQVYVlrbyB_E3kFFhrKU30BosCo,5438
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/training/activation_scaler.py,sha256=FzNfgBplLWmyiSlZ6TUvE-nur3lOiGTrlvC97ys8S24,1973
@@ -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.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,,
+sae_lens-6.29.1.dist-info/METADATA,sha256=0Pp1L3vNiUGzkMox_BdQR6B064tTHFgwAPGJz8FY8UM,6573
+sae_lens-6.29.1.dist-info/WHEEL,sha256=zp0Cn7JsFoX2ATtOhtaFYIiE2rmFAD4OcMhtUki8W3U,88
+sae_lens-6.29.1.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.29.1.dist-info/RECORD,,