Premchan369/Q-TensorFormer

1"""2Tests for tensor decomposition layers.3 4Verifies:5  - Correct output shapes6  - Rank truncation preserves structure7  - Compression ratio computation8  - Gradient flow9"""10 11import sys12import os13sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))14 15import torch16import pytest17from src.tensor_layers import TTLinear, TTFeedForward, factorize_dim18 19 20class TestFactorizeDim:21    def test_power_of_two(self):22        factors = factorize_dim(64)23        assert all(f >= 2 for f in factors), f"Got dead factor in {factors}"24 25    def test_prime(self):26        factors = factorize_dim(7)27        # Some factors may be 1 for primes (unavoidable)28        product = 129        for f in factors:30            product *= f31        assert product == 7, f"Prime 7 product: {factors} = {product}"32 33    def test_one(self):34        factors = factorize_dim(1)35        assert factors == (1,)36 37    def test_large(self):38        for dim in [128, 256, 512, 1024]:39            factors = factorize_dim(dim)40            product = 141            for f in factors:42                product *= f43            assert product == dim, f"Product mismatch: {factors} = {product} != {dim}"44 45 46class TestTTLinear:47    def test_output_shape(self):48        layer = TTLinear(64, 128, rank=8)49        x = torch.randn(4, 64)50        y = layer(x)51        assert y.shape == (4, 128)52 53    def test_batched(self):54        layer = TTLinear(64, 128, rank=8)55        x = torch.randn(3, 5, 64)56        y = layer(x)57        assert y.shape == (3, 5, 128)58 59    def test_gradient_flow(self):60        layer = TTLinear(64, 128, rank=8)61        x = torch.randn(4, 64, requires_grad=False)62        y = layer(x)63        loss = y.sum()64        loss.backward()65        for core in layer.cores:66            assert core.grad is not None67            assert not torch.isnan(core.grad).any()68 69    def test_set_rank_smaller(self):70        layer = TTLinear(64, 128, rank=8)71        x = torch.randn(4, 64)72        y_before = layer(x)73 74        layer.set_rank(4)75        y_after = layer(x)76 77        assert y_after.shape == y_before.shape78        assert layer.rank == 479 80    def test_set_rank_larger(self):81        layer = TTLinear(64, 128, rank=4)82        layer.set_rank(8)83        assert layer.rank == 884 85    def test_compression_ratio(self):86        layer = TTLinear(128, 256, rank=8)87        assert layer.compression_ratio > 1.088 89    def test_bias(self):90        layer = TTLinear(64, 128, rank=8, bias=True)91        assert layer.bias is not None92 93        layer_nb = TTLinear(64, 128, rank=8, bias=False)94        assert layer_nb.bias is None95 96 97class TestTTFeedForward:98    def test_output_shape(self):99        ffn = TTFeedForward(128, ff_multiplier=4, rank=8)100        x = torch.randn(4, 128)101        y = ffn(x)102        assert y.shape == (4, 128)103 104    def test_set_rank(self):105        ffn = TTFeedForward(128, ff_multiplier=4, rank=8)106        x = torch.randn(4, 128)107        y_before = ffn(x)108 109        ffn.set_rank(4)110        y_after = ffn(x)111 112        assert y_after.shape == y_before.shape113 114    def test_total_params(self):115        ffn = TTFeedForward(128, ff_multiplier=4, rank=8)116        params = ffn.total_params117        assert params > 0118        # Should be fewer than dense equivalent119        dense = 128 * 512 + 512 * 128  # up + down120        assert params < dense121