File size: 1,271 Bytes
eaeea8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """
Tests for Nested-Core Tensor-Train Layer.
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import pytest
from src.tensor_layers import TTLinear, TTFeedForward
def test_nested_tt_slicing_shapes():
in_dim, out_dim = 64, 128
layer = TTLinear(in_dim, out_dim, max_rank=8)
x = torch.randn(2, 4, in_dim)
# Test all candidate ranks
for r in [1, 2, 4, 8]:
layer.set_rank(r)
assert layer.rank == r
out = layer(x)
assert out.shape == (2, 4, out_dim), f"Expected (2, 4, {out_dim}), got {out.shape}"
assert not torch.isnan(out).any(), f"NaN in output at rank {r}"
traffic = layer.get_memory_traffic()
assert traffic["total_bytes"] > 0
print("✓ test_nested_tt_slicing_shapes passed")
def test_nested_tt_ffn():
ffn = TTFeedForward(hidden_dim=64, ff_multiplier=4, rank=8)
x = torch.randn(2, 64)
for r in [1, 2, 4, 8]:
ffn.set_rank(r)
out = ffn(x)
assert out.shape == (2, 64)
assert ffn.active_params > 0
print("✓ test_nested_tt_ffn passed")
if __name__ == "__main__":
test_nested_tt_slicing_shapes()
test_nested_tt_ffn()
print("All Nested TT tests passed!")
|