format

Author: Butanium

dictionary_learning/trainers/crosscoder.py

@@ -212,7 +212,9 @@ def loss(
                 "deads": deads if return_deads else None,
             }
             for layer in range(x.shape[1]):
-                log_dict[f"rms_norm_l{layer}"] = th.sqrt((x[:, layer, :].pow(2).sum(-1)).mean()).item()
+                log_dict[f"rms_norm_l{layer}"] = th.sqrt(
+                    (x[:, layer, :].pow(2).sum(-1)).mean()
+                ).item()
             return namedtuple("LossLog", ["x", "x_hat", "f", "losses"])(
                 x,
                 x_hat,

dictionary_learning/utils.py

@@ -122,4 +122,4 @@ def remove_gradient_parallel_to_decoder_directions(
         normed_W_dec_DF,
         "d_sae, d_in d_sae -> d_in d_sae",
     )
-    return W_dec_DF_grad
+    return W_dec_DF_grad

tests/test_cache.py

@@ -9,6 +9,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import numpy as np
 
+
 @pytest.fixture
 def temp_dir():
     """Create a temporary directory for test files."""
@@ -274,7 +275,7 @@ def test_activation_cache_with_normalizer(temp_dir):
 def test_sequence_ranges_no_bos_token(temp_dir):
     """Test that sequence ranges are stored when model has no BOS token."""
     # Set flag to handle meta tensors properly
-    if hasattr(th.fx, 'experimental'):
+    if hasattr(th.fx, "experimental"):
         th.fx.experimental._config.meta_nonzero_assume_all_nonzero = True
 
     # Skip test if CUDA not available
@@ -296,12 +297,18 @@ def test_sequence_ranges_no_bos_token(temp_dir):
     )
     model = LanguageModel(model, torch_dtype=th.float32, tokenizer=tokenizer)
     model.tokenizer.pad_token = model.tokenizer.eos_token
-    
+
     # Simulate model without BOS token
     original_bos_token_id = model.tokenizer.bos_token_id
     model.tokenizer.bos_token_id = None
 
-    tokens = model.tokenizer(test_strings, add_special_tokens=True, return_tensors="pt", padding=True, truncation=True)
+    tokens = model.tokenizer(
+        test_strings,
+        add_special_tokens=True,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+    )
     lengths = tokens["attention_mask"].sum(dim=1).tolist()
     ranges = np.cumsum([0] + lengths)
     try:
@@ -335,28 +342,40 @@ def test_sequence_ranges_no_bos_token(temp_dir):
 
         # Verify sequence ranges were stored
         sequence_ranges = cache.sequence_ranges
-        assert sequence_ranges is not None, "sequence ranges should be stored for model without BOS token"
-        
+        assert (
+            sequence_ranges is not None
+        ), "sequence ranges should be stored for model without BOS token"
+
         # Should have one sequence start per input string plus one for the last sequence
-        assert len(sequence_ranges) == len(test_strings) + 1, f"Expected {len(test_strings)} sequence ranges, got {len(sequence_ranges)}"
-        
+        assert (
+            len(sequence_ranges) == len(test_strings) + 1
+        ), f"Expected {len(test_strings)} sequence ranges, got {len(sequence_ranges)}"
+
         # First sequence should start at position 0
-        assert sequence_ranges[0].item() == 0, "First sequence should start at position 0"
+        assert (
+            sequence_ranges[0].item() == 0
+        ), "First sequence should start at position 0"
 
         # sequence ranges should be the same as the ranges computed from the tokens
-        assert np.allclose(sequence_ranges, ranges), "sequence ranges should be the same as the ranges computed from the tokens"
-        
+        assert np.allclose(
+            sequence_ranges, ranges
+        ), "sequence ranges should be the same as the ranges computed from the tokens"
+
         # sequence ranges should be in ascending order
         for i in range(1, len(sequence_ranges)):
-            assert sequence_ranges[i] > sequence_ranges[i-1], f"sequence ranges should be ascending: {sequence_ranges}"
+            assert (
+                sequence_ranges[i] > sequence_ranges[i - 1]
+            ), f"sequence ranges should be ascending: {sequence_ranges}"
 
         # Verify sequence ranges align with token boundaries
         tokens = cache.tokens
         total_tokens = len(tokens)
-        
+
         # All sequence ranges should be valid indices
         for start_idx in sequence_ranges:
-            assert 0 <= start_idx <= total_tokens, f"Invalid sequence start index: {start_idx}"
+            assert (
+                0 <= start_idx <= total_tokens
+            ), f"Invalid sequence start index: {start_idx}"
 
     finally:
         # Restore original BOS token
@@ -366,7 +385,7 @@ def test_sequence_ranges_no_bos_token(temp_dir):
 def test_sequence_ranges_with_bos_token(temp_dir):
     """Test that sequence ranges are NOT stored when model has BOS token."""
     # Set flag to handle meta tensors properly
-    if hasattr(th.fx, 'experimental'):
+    if hasattr(th.fx, "experimental"):
         th.fx.experimental._config.meta_nonzero_assume_all_nonzero = True
 
     # Skip test if CUDA not available
@@ -382,7 +401,7 @@ def test_sequence_ranges_with_bos_token(temp_dir):
     )
     model = LanguageModel(model, torch_dtype=th.float32, tokenizer=tokenizer)
     model.tokenizer.pad_token = model.tokenizer.eos_token
-    
+
     # Ensure model has BOS token (set it explicitly)
     model.tokenizer.bos_token_id = model.tokenizer.eos_token_id
 
@@ -411,7 +430,9 @@ def test_sequence_ranges_with_bos_token(temp_dir):
 
     # Verify sequence ranges were NOT stored
     sequence_ranges = cache.sequence_ranges
-    assert sequence_ranges is None, "sequence ranges should not be stored for model with BOS token"
+    assert (
+        sequence_ranges is None
+    ), "sequence ranges should not be stored for model with BOS token"
 
 
 def test_activation_cache_slice_indexing_cross_shard(temp_dir):
@@ -469,39 +490,45 @@ def test_activation_cache_slice_indexing_cross_shard(temp_dir):
 
     # Load the cached activations
     cache = ActivationCache(temp_dir, submodule_name + "_out")
-    
+
     # Verify we have multiple shards
-    assert len(cache.shards) >= 2, f"Expected at least 2 shards, got {len(cache.shards)}"
-    
+    assert (
+        len(cache.shards) >= 2
+    ), f"Expected at least 2 shards, got {len(cache.shards)}"
+
     total_size = len(cache)
     print(f"Cache has {len(cache.shards)} shards with total size {total_size}")
-    
+
     # Print shard boundaries for debugging
     shard_boundaries = cache._range_to_shard_idx
     print(f"Shard boundaries: {shard_boundaries}")
-    
+
     # Test 1: Slice that crosses exactly one shard boundary
     if len(cache.shards) >= 2:
         # Find a slice that starts in first shard and ends in second shard
         first_shard_end = shard_boundaries[1]
         start_idx = max(0, first_shard_end - 10)
         end_idx = min(total_size, first_shard_end + 10)
-        
+
         # Get slice result
         slice_result = cache[start_idx:end_idx]
-        
+
         # Get individual results for comparison
-        individual_results = th.stack([cache[i] for i in range(start_idx, end_idx)], dim=0)
-        
+        individual_results = th.stack(
+            [cache[i] for i in range(start_idx, end_idx)], dim=0
+        )
+
         # Verify they match
-        assert th.allclose(slice_result, individual_results, atol=1e-5, rtol=1e-5), \
-            f"Slice result doesn't match individual indexing for indices {start_idx}:{end_idx}"
-        
+        assert th.allclose(
+            slice_result, individual_results, atol=1e-5, rtol=1e-5
+        ), f"Slice result doesn't match individual indexing for indices {start_idx}:{end_idx}"
+
         # Verify correct shape
         expected_length = end_idx - start_idx
-        assert slice_result.shape[0] == expected_length, \
-            f"Expected slice length {expected_length}, got {slice_result.shape[0]}"
-        
+        assert (
+            slice_result.shape[0] == expected_length
+        ), f"Expected slice length {expected_length}, got {slice_result.shape[0]}"
+
         print(f"✓ Cross-shard slice test 1 passed: indices {start_idx}:{end_idx}")
 
     # Test 2: Slice that spans multiple shards
@@ -510,54 +537,70 @@ def test_activation_cache_slice_indexing_cross_shard(temp_dir):
         second_shard_end = shard_boundaries[2]
         start_idx = max(0, shard_boundaries[1] - 5)  # Start near end of first shard
         end_idx = min(total_size, second_shard_end + 5)  # End in third shard
-        
+
         slice_result = cache[start_idx:end_idx]
-        individual_results = th.stack([cache[i] for i in range(start_idx, end_idx)], dim=0)
-        
-        assert th.allclose(slice_result, individual_results, atol=1e-5, rtol=1e-5), \
-            f"Multi-shard slice result doesn't match individual indexing for indices {start_idx}:{end_idx}"
-        
+        individual_results = th.stack(
+            [cache[i] for i in range(start_idx, end_idx)], dim=0
+        )
+
+        assert th.allclose(
+            slice_result, individual_results, atol=1e-5, rtol=1e-5
+        ), f"Multi-shard slice result doesn't match individual indexing for indices {start_idx}:{end_idx}"
+
         expected_length = end_idx - start_idx
-        assert slice_result.shape[0] == expected_length, \
-            f"Expected multi-shard slice length {expected_length}, got {slice_result.shape[0]}"
-        
+        assert (
+            slice_result.shape[0] == expected_length
+        ), f"Expected multi-shard slice length {expected_length}, got {slice_result.shape[0]}"
+
         print(f"✓ Multi-shard slice test passed: indices {start_idx}:{end_idx}")
 
     # Test 3: Slice with step parameter across shards
     if total_size >= 50:
         start_idx = 5
         end_idx = min(total_size, 45)
         step = 3
-        
+
         slice_result = cache[start_idx:end_idx:step]
-        individual_results = th.stack([cache[i] for i in range(start_idx, end_idx, step)], dim=0)
-        
-        assert th.allclose(slice_result, individual_results, atol=1e-5, rtol=1e-5), \
-            f"Stepped slice result doesn't match individual indexing for indices {start_idx}:{end_idx}:{step}"
-        
+        individual_results = th.stack(
+            [cache[i] for i in range(start_idx, end_idx, step)], dim=0
+        )
+
+        assert th.allclose(
+            slice_result, individual_results, atol=1e-5, rtol=1e-5
+        ), f"Stepped slice result doesn't match individual indexing for indices {start_idx}:{end_idx}:{step}"
+
         expected_length = len(range(start_idx, end_idx, step))
-        assert slice_result.shape[0] == expected_length, \
-            f"Expected stepped slice length {expected_length}, got {slice_result.shape[0]}"
-        
+        assert (
+            slice_result.shape[0] == expected_length
+        ), f"Expected stepped slice length {expected_length}, got {slice_result.shape[0]}"
+
         print(f"✓ Stepped slice test passed: indices {start_idx}:{end_idx}:{step}")
 
     # Test 4: Edge cases - slice at boundaries
     if len(cache.shards) >= 2:
         # Test slice starting exactly at shard boundary
         boundary_idx = shard_boundaries[1]
         if boundary_idx < total_size - 5:
-            slice_result = cache[boundary_idx:boundary_idx + 5]
-            individual_results = th.stack([cache[i] for i in range(boundary_idx, boundary_idx + 5)], dim=0)
-            
-            assert th.allclose(slice_result, individual_results, atol=1e-5, rtol=1e-5), \
-                f"Boundary slice result doesn't match individual indexing"
-            
-            print(f"✓ Boundary slice test passed: starting at shard boundary {boundary_idx}")
+            slice_result = cache[boundary_idx : boundary_idx + 5]
+            individual_results = th.stack(
+                [cache[i] for i in range(boundary_idx, boundary_idx + 5)], dim=0
+            )
+
+            assert th.allclose(
+                slice_result, individual_results, atol=1e-5, rtol=1e-5
+            ), f"Boundary slice result doesn't match individual indexing"
+
+            print(
+                f"✓ Boundary slice test passed: starting at shard boundary {boundary_idx}"
+            )
 
     # Test 5: Empty slice
     empty_slice = cache[10:10]
-    assert empty_slice.shape[0] == 0, f"Expected empty slice, got shape {empty_slice.shape}"
+    assert (
+        empty_slice.shape[0] == 0
+    ), f"Expected empty slice, got shape {empty_slice.shape}"
     print("✓ Empty slice test passed")
-    
 
-    print(f"✓ All slice indexing tests passed for cache with {len(cache.shards)} shards")
+    print(
+        f"✓ All slice indexing tests passed for cache with {len(cache.shards)} shards"
+    )