""" Unit tests for decomposition module. Tests tensor decomposition functionality including SVD, QR, and error handling. """ import pytest import torch # Imported as `src.*` so the package-relative imports inside src/ resolve. # conftest.py puts the project root on sys.path and mocks the ComfyUI modules. from src.decomposition import ( SVDDecomposer, RandomizedSVDDecomposer, EnergyBasedRandomizedSVDDecomposer, SingularValueDistribution, ) class TestSVDDecomposer: """Tests for standard SVD decomposer.""" def test_basic_2d_decomposition(self): """Test basic 2D tensor decomposition.""" # return_statistics is a constructor option, not a decompose() argument. decomposer = SVDDecomposer(return_statistics=True) weight = torch.randn(100, 50) target_rank = 10 up, down, alpha, stats = decomposer.decompose(weight, target_rank) # Check shapes assert up.shape == (100, 10) assert down.shape == (10, 50) assert isinstance(alpha, float) # Reconstruction quality is bounded by the discarded singular values, not # by the decomposer: for a full-rank random Gaussian, a rank-10 truncation # of 50 singular values necessarily loses ~75% of the Frobenius norm. # Assert we land at that optimum rather than at an arbitrary threshold. reconstructed = up @ down reconstruction_error = torch.norm(weight - reconstructed) / torch.norm(weight) S = torch.linalg.svdvals(weight) optimal_error = (S[target_rank:].pow(2).sum() / S.pow(2).sum()).sqrt() assert reconstruction_error < optimal_error + 0.01 def test_low_rank_input_reconstructs_accurately(self): """A genuinely low-rank matrix is recovered with little error.""" decomposer = SVDDecomposer() weight = torch.randn(100, 10) @ torch.randn(10, 50) weight = weight + 0.01 * torch.randn(100, 50) up, down, alpha, _ = decomposer.decompose(weight, target_rank=10) reconstruction_error = torch.norm(weight - up @ down) / torch.norm(weight) assert reconstruction_error < 0.05 def test_4d_conv_decomposition(self): """Test 4D convolutional tensor decomposition.""" decomposer = SVDDecomposer() weight = torch.randn(64, 32, 3, 3) # Conv layer target_rank = 16 up, down, alpha, _ = decomposer.decompose(weight, target_rank) # Check shapes assert up.shape == (64, 16, 1, 1) assert down.shape == (16, 32, 3, 3) def test_symmetric_vs_asymmetric_distribution(self): """Test different singular value distributions.""" weight = torch.randn(50, 30) target_rank = 10 # Symmetric distribution decomposer_sym = SVDDecomposer( distribution=SingularValueDistribution.SYMMETRIC ) up_sym, down_sym, _, _ = decomposer_sym.decompose(weight, target_rank) # Asymmetric distribution decomposer_asym = SVDDecomposer( distribution=SingularValueDistribution.ASYMMETRIC ) up_asym, down_asym, _, _ = decomposer_asym.decompose(weight, target_rank) # Both should reconstruct similarly but with different scaling recon_sym = up_sym @ down_sym recon_asym = up_asym @ down_asym # The product is mathematically identical either way -- the distributions # only decide whether S goes into up, into down, or is split as sqrt(S) # across both. atol is needed because the default (1e-8) is below float32 # rounding for entries near zero; observed difference is ~5e-7. assert torch.allclose(recon_sym, recon_asym, rtol=1e-4, atol=1e-5) def test_statistics_calculation(self): """Test that statistics are calculated correctly.""" decomposer = SVDDecomposer(return_statistics=True) weight = torch.randn(80, 40) target_rank = 20 _, _, _, stats = decomposer.decompose(weight, target_rank) assert stats is not None assert 'new_rank' in stats assert 'new_alpha' in stats assert 'sum_retained' in stats assert 'fro_retained' in stats assert 'max_ratio' in stats # Check statistics are reasonable assert stats['new_rank'] == 20 assert 0.0 <= stats['sum_retained'] <= 1.0 assert 0.0 <= stats['fro_retained'] <= 1.0 def test_dynamic_rank_selection_ratio(self): """Test dynamic rank selection by singular value ratio.""" decomposer = SVDDecomposer(return_statistics=True) weight = torch.randn(100, 50) _, _, _, stats = decomposer.decompose( weight, target_rank=50, dynamic_method="sv_ratio", dynamic_param=100.0 # Ratio threshold ) # Rank should be selected based on ratio assert stats['new_rank'] <= 50 assert stats['new_rank'] >= 1 def test_dynamic_rank_selection_cumulative(self): """Test dynamic rank selection by cumulative singular values.""" decomposer = SVDDecomposer(return_statistics=True) weight = torch.randn(100, 50) _, _, _, stats = decomposer.decompose( weight, target_rank=50, dynamic_method="sv_cumulative", dynamic_param=0.95 # 95% of cumulative sum ) # Rank should be selected to capture 95% of singular values assert stats['sum_retained'] >= 0.90 # Allow some tolerance def test_dynamic_rank_selection_frobenius(self): """Test dynamic rank selection by Frobenius norm.""" decomposer = SVDDecomposer(return_statistics=True) weight = torch.randn(100, 50) _, _, _, stats = decomposer.decompose( weight, target_rank=50, dynamic_method="sv_fro", dynamic_param=0.99 # 99% of Frobenius norm ) # Rank should be selected to retain 99% of Frobenius norm assert stats['fro_retained'] >= 0.95 # Allow some tolerance class TestRandomizedSVDDecomposer: """Tests for randomized SVD decomposer.""" def test_randomized_svd_approximation(self): """Test that randomized SVD produces good approximation.""" weight = torch.randn(200, 100) target_rank = 20 # Standard SVD decomposer_std = SVDDecomposer() up_std, down_std, _, _ = decomposer_std.decompose(weight, target_rank) recon_std = up_std @ down_std # Randomized SVD decomposer_rand = RandomizedSVDDecomposer(n_oversamples=10, n_iter=2) up_rand, down_rand, _, _ = decomposer_rand.decompose(weight, target_rank) recon_rand = up_rand @ down_rand # Reconstructions should be similar error_std = torch.norm(weight - recon_std) error_rand = torch.norm(weight - recon_rand) # Randomized should be close to standard (within 50% relative error) assert abs(error_rand - error_std) / error_std < 0.5 def test_randomized_svd_small_matrix(self): """Test that small matrices fall back to standard SVD.""" decomposer = RandomizedSVDDecomposer() weight = torch.randn(50, 30) # Small matrix target_rank = 10 # Should not raise error up, down, alpha, _ = decomposer.decompose(weight, target_rank) assert up.shape == (50, 10) assert down.shape == (10, 30) class TestEnergyBasedRandomizedSVDDecomposer: """Tests for energy-based randomized SVD.""" def test_energy_based_rank_selection(self): """Test that energy threshold affects rank selection.""" weight = torch.randn(100, 50) # Low energy threshold (fewer components) decomposer_low = EnergyBasedRandomizedSVDDecomposer( energy_threshold=0.8, return_statistics=True ) _, _, _, stats_low = decomposer_low.decompose(weight, target_rank=50) # High energy threshold (more components) decomposer_high = EnergyBasedRandomizedSVDDecomposer( energy_threshold=0.99, return_statistics=True ) _, _, _, stats_high = decomposer_high.decompose(weight, target_rank=50) # Higher threshold should generally use more components # (though not guaranteed due to randomness) assert stats_low is not None assert stats_high is not None class TestErrorHandling: """Tests for error handling in decomposition.""" def test_invalid_tensor_dimensions(self): """Test that invalid tensor dimensions raise errors.""" decomposer = SVDDecomposer() weight = torch.randn(10) # 1D tensor (invalid) with pytest.raises(ValueError, match="must be 2D, 3D, or 4D"): decomposer.decompose(weight, target_rank=5) def test_zero_matrix_handling(self): """Test handling of numerically zero matrices.""" decomposer = SVDDecomposer(return_statistics=True) weight = torch.zeros(50, 30) # The zero-matrix guard lives in dynamic rank selection; without a # dynamic_method the caller has pinned the rank and target_rank is honoured. up, down, alpha, stats = decomposer.decompose( weight, target_rank=10, dynamic_method="sv_ratio", dynamic_param=2.0 ) # Rank should collapse to the minimum for a numerically zero matrix assert stats['new_rank'] == 1 def test_zero_matrix_fixed_rank_is_honoured(self): """A pinned rank (no dynamic_method) is kept even for a zero matrix.""" decomposer = SVDDecomposer(return_statistics=True) weight = torch.zeros(50, 30) up, down, alpha, stats = decomposer.decompose(weight, target_rank=10) assert stats['new_rank'] == 10 assert torch.allclose(up @ down, torch.zeros(50, 30)) def test_invalid_dynamic_method(self): """Test that invalid dynamic method raises error.""" decomposer = SVDDecomposer() weight = torch.randn(50, 30) with pytest.raises(ValueError, match="Unknown dynamic method"): decomposer.decompose( weight, target_rank=10, dynamic_method="invalid_method" ) # Fixtures @pytest.fixture def sample_2d_weight(): """Fixture providing sample 2D weight tensor.""" return torch.randn(100, 50) @pytest.fixture def sample_4d_weight(): """Fixture providing sample 4D convolutional weight.""" return torch.randn(64, 32, 3, 3) class TestIntegrationWithFixtures: """Integration tests using fixtures.""" def test_all_decomposers_with_2d(self, sample_2d_weight): """Test all decomposers work with 2D tensors.""" decomposers = [ SVDDecomposer(), RandomizedSVDDecomposer(), EnergyBasedRandomizedSVDDecomposer(), ] for decomposer in decomposers: up, down, alpha, _ = decomposer.decompose( sample_2d_weight, target_rank=20 ) assert up.shape[0] == 100 assert up.shape[1] == 20 assert down.shape[0] == 20 assert down.shape[1] == 50 def test_all_decomposers_with_4d(self, sample_4d_weight): """Test all decomposers work with 4D conv tensors.""" decomposers = [ SVDDecomposer(), RandomizedSVDDecomposer(), ] for decomposer in decomposers: up, down, alpha, _ = decomposer.decompose( sample_4d_weight, target_rank=16 ) assert up.shape == (64, 16, 1, 1) assert down.shape == (16, 32, 3, 3) if __name__ == "__main__": pytest.main([__file__, "-v"])