diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 58b61c6..47d67d6 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -142,6 +142,20 @@ class TestWhitenedSVD: assert idx in results assert results[idx].directions.shape[0] == 2 + def test_extract_all_layers_skips_unpaired_harmful_layer(self): + extractor = WhitenedSVDExtractor() + harmful = { + 0: [torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])], + 1: [torch.tensor([1.0, 1.0]), torch.tensor([2.0, 2.0])], + } + harmless = { + 0: [torch.tensor([0.0, 0.0]), torch.tensor([0.0, 0.0])], + } + + results = extractor.extract_all_layers(harmful, harmless, n_directions=1) + + assert results.keys() == {0} + def test_compare_with_standard(self): """Comparison should return valid cosine similarities.""" torch.manual_seed(42) @@ -159,6 +173,13 @@ class TestWhitenedSVD: assert "subspace_principal_cosine" in comparison assert 0 <= comparison["primary_direction_cosine"] <= 1.0 + standard_subspace = torch.linalg.qr(torch.randn(16, 2)).Q.T + subspace_comparison = WhitenedSVDExtractor.compare_with_standard( + result, + standard_subspace, + ) + assert 0 <= subspace_comparison["subspace_principal_cosine"] <= 1.0 + def test_handles_3d_activations(self): """Should handle activations with an extra batch dimension.""" torch.manual_seed(42)