Premchan369/Q-TensorFormer

1"""2Tests for Nested-Core Tensor-Train Layer.3"""4 5import sys6import os7sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))8 9import torch10import pytest11from src.tensor_layers import TTLinear, TTFeedForward12 13 14def test_nested_tt_slicing_shapes():15    in_dim, out_dim = 64, 12816    layer = TTLinear(in_dim, out_dim, max_rank=8)17 18    x = torch.randn(2, 4, in_dim)19 20    # Test all candidate ranks21    for r in [1, 2, 4, 8]:22        layer.set_rank(r)23        assert layer.rank == r24        out = layer(x)25        assert out.shape == (2, 4, out_dim), f"Expected (2, 4, {out_dim}), got {out.shape}"26        assert not torch.isnan(out).any(), f"NaN in output at rank {r}"27        traffic = layer.get_memory_traffic()28        assert traffic["total_bytes"] > 029    print("✓ test_nested_tt_slicing_shapes passed")30 31 32def test_nested_tt_ffn():33    ffn = TTFeedForward(hidden_dim=64, ff_multiplier=4, rank=8)34    x = torch.randn(2, 64)35 36    for r in [1, 2, 4, 8]:37        ffn.set_rank(r)38        out = ffn(x)39        assert out.shape == (2, 64)40        assert ffn.active_params > 041    print("✓ test_nested_tt_ffn passed")42 43 44if __name__ == "__main__":45    test_nested_tt_slicing_shapes()46    test_nested_tt_ffn()47    print("All Nested TT tests passed!")48