Premchan369/Q-TensorFormer

1"""2Tests for Token Information State Module.3"""4 5import sys6import os7sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))8 9import torch10import pytest11from src.information_state import TokenInformationState12 13 14def test_information_state_shape_and_range():15    d_model = 6416    n_heads = 417    engine = TokenInformationState(d_model=d_model, n_heads=n_heads)18 19    batch_size = 220    seq_len = 1621    x = torch.randn(batch_size, seq_len, d_model)22    attn = torch.softmax(torch.randn(batch_size, n_heads, seq_len, seq_len), dim=-1)23 24    z_t = engine(x, attn_weights=attn)25 26    assert z_t.shape == (batch_size, seq_len, 8), f"Expected (2, 16, 8), got {z_t.shape}"27    # Values should be normalized to [0, 1]28    assert z_t.min() >= 0.0, f"Min value {z_t.min()} < 0.0"29    assert z_t.max() <= 1.0, f"Max value {z_t.max()} > 1.0"30    print("✓ test_information_state_shape_and_range passed")31 32 33def test_information_state_ablation():34    engine = TokenInformationState(d_model=32, n_heads=2)35    x = torch.randn(1, 8, 32)36 37    # Disable predictive_entropy38    engine.set_ablation("predictive_entropy", False)39    z_t = engine(x)40 41    pred_idx = engine.FEATURE_NAMES.index("predictive_entropy")42    assert (z_t[:, :, pred_idx] == 0.0).all(), "Ablated feature should be zero"43 44    # Reset45    engine.reset_ablation()46    z_t_reset = engine(x)47    assert (z_t_reset[:, :, pred_idx] > 0.0).any(), "Reset feature should be restored"48    print("✓ test_information_state_ablation passed")49 50 51if __name__ == "__main__":52    test_information_state_shape_and_range()53    test_information_state_ablation()54    print("All Information State tests passed!")55