import unittest import torch from backbones import GPTBlock, TemporalAttention, TrajMixer from models import ( TRAJ_MIXER_ARCHITECTURE, validate_traj_mixer_config, validate_traj_mixer_state_dict, ) from train_util import get_model_parameter_counts class TrajMixerTest(unittest.TestCase): def test_zero_rbf_bias_has_live_projection_gradient(self) -> None: attention = TemporalAttention( n_embd=12, n_head=3, use_time_rope=False, use_rbf_bias=True, ) features = torch.randn(2, 4, 4, 16) target = torch.randn(2, 4, 4, 3) initial_bias = ( attention.time_bias_scale.tanh() * attention.rbf_proj(features) ) torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias)) loss = (initial_bias * target).sum() loss.backward() projection_grad = attention.rbf_proj.weight.grad self.assertIsNotNone(projection_grad) self.assertGreater(projection_grad.abs().sum().item(), 0.0) with torch.no_grad(): attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3) attention.zero_grad(set_to_none=True) updated_bias = ( attention.time_bias_scale.tanh() * attention.rbf_proj(features) ) (updated_bias * target).sum().backward() scale_grad = attention.time_bias_scale.grad self.assertIsNotNone(scale_grad) self.assertGreater(scale_grad.abs().item(), 0.0) def test_default_shape_parameters_and_initialization(self) -> None: mixer = TrajMixer( n_embd=120, n_head=10, dropout=0.0, ) x = torch.randn(2, 7, 120) self.assertEqual(mixer(x).shape, x.shape) self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040) self.assertFalse(hasattr(mixer, "group_align")) self.assertFalse(hasattr(mixer, "intra_norm")) self.assertFalse(hasattr(mixer, "cross_norm")) self.assertEqual(tuple(mixer.norm.normalized_shape), (120,)) self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12)) torch.testing.assert_close( torch.sigmoid(mixer.intra_gate_logits.detach()), torch.full((10, 12), 0.1), ) self.assertEqual(mixer.intra_hidden, 48) self.assertEqual( tuple(mixer.intra_gate_proj.shape), (10, 12, 48), ) self.assertEqual( tuple(mixer.intra_value_proj.shape), (10, 12, 48), ) self.assertEqual( tuple(mixer.intra_output_proj.shape), (10, 48, 12), ) self.assertEqual(mixer.hidden_group, 40) self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40)) self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40)) self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10)) def test_zero_final_output_projection_makes_mixer_identity(self) -> None: torch.manual_seed(0) mixer = TrajMixer(120, n_head=10, dropout=0.0) with torch.no_grad(): mixer.output_proj.zero_() x = torch.randn(2, 5, 120) torch.testing.assert_close(mixer(x), x) def test_forward_matches_single_outer_residual_formula(self) -> None: torch.manual_seed(0) mixer = TrajMixer(120, n_head=10, dropout=0.0) mixer.eval() x = torch.randn(2, 5, 120) grouped = mixer.norm(x).reshape(2, 5, 10, 12) intra_output = mixer._intra_mix(grouped) static_gate = torch.sigmoid(mixer.intra_gate_logits).view( 1, 1, 10, 12 ) mixed_input = grouped + static_gate * intra_output update = mixer._cross_mix(mixed_input).reshape(2, 5, 120) torch.testing.assert_close(mixer(x), x + update) def test_intra_stage_is_independent_across_groups(self) -> None: torch.manual_seed(0) mixer = TrajMixer(120, n_head=10, dropout=0.0) mixer.eval() grouped = torch.randn(2, 4, 10, 12) changed = grouped.clone() changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :]) original_out = mixer._intra_mix(grouped) changed_out = mixer._intra_mix(changed) unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9]) torch.testing.assert_close( original_out.index_select(2, unchanged_groups), changed_out.index_select(2, unchanged_groups), ) def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None: mixer = TrajMixer(6, n_head=3, dropout=0.0) mixer.eval() with torch.no_grad(): mixer.gate_proj.zero_() mixer.value_proj.zero_() mixer.output_proj.zero_() # For coordinate 0 only, read group 0 through hidden unit 0 and # write the resulting gated value into group 1. mixer.gate_proj[0, 0, 0] = 1.0 mixer.value_proj[0, 0, 0] = 1.0 mixer.output_proj[0, 0, 1] = 1.0 grouped = torch.tensor( [[[ [-1.0, 4.0], [0.0, 5.0], [1.0, 6.0], ]]] ) changed = grouped.clone() changed[0, 0, 0, 0] = 2.0 original_out = mixer._cross_mix(grouped) changed_out = mixer._cross_mix(changed) self.assertNotEqual( original_out[0, 0, 1, 0].item(), changed_out[0, 0, 1, 0].item(), ) torch.testing.assert_close( original_out[..., 1], changed_out[..., 1], ) def test_mixer_does_not_mix_sequence_positions(self) -> None: torch.manual_seed(0) mixer = TrajMixer(120, n_head=10, dropout=0.0) mixer.eval() x = torch.randn(2, 5, 120) changed = x.clone() changed[:, 3, :] += torch.randn_like(changed[:, 3, :]) original_out = mixer(x) changed_out = mixer(changed) unchanged_positions = torch.tensor([0, 1, 2, 4]) torch.testing.assert_close( original_out.index_select(1, unchanged_positions), changed_out.index_select(1, unchanged_positions), ) def test_gradients_reach_all_projection_families(self) -> None: torch.manual_seed(1) mixer = TrajMixer(120, n_head=10, dropout=0.0) x = torch.randn(2, 4, 120, requires_grad=True) mixer(x).square().mean().backward() self.assertIsNotNone(x.grad) for name, parameter in mixer.named_parameters(): self.assertIsNotNone(parameter.grad, name) self.assertTrue(torch.isfinite(parameter.grad).all(), name) def test_gpt_block_delegates_single_mixer_residual_to_traj_mixer(self) -> None: block = GPTBlock(n_embd=120, n_head=10) self.assertIsInstance(block.mlp, TrajMixer) self.assertFalse(hasattr(block, "ln2")) self.assertIsInstance(block.mlp.norm, torch.nn.LayerNorm) self.assertFalse(hasattr(block.mlp, "intra_norm")) self.assertFalse(hasattr(block.mlp, "cross_norm")) x = torch.randn(2, 6, 120) self.assertEqual(block(x).shape, x.shape) def test_architecture_marker_is_required(self) -> None: validate_traj_mixer_config( {"model_architecture": TRAJ_MIXER_ARCHITECTURE} ) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_traj_mixer_config({}) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_traj_mixer_config({"model_architecture": "delphi_swiglu"}) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_traj_mixer_config( {"model_architecture": "traj_mixer_v2"} ) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_traj_mixer_config( {"model_architecture": "traj_mixer_v3"} ) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_traj_mixer_config( {"model_architecture": "traj_mixer_v4"} ) def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None: block = GPTBlock(n_embd=120, n_head=10) state_dict = { f"blocks.0.{key}": value for key, value in block.state_dict().items() } validate_traj_mixer_state_dict(state_dict) state_dict.pop("blocks.0.mlp.intra_gate_proj") with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"): validate_traj_mixer_state_dict(state_dict) def test_invalid_group_partition_is_rejected(self) -> None: with self.assertRaisesRegex(ValueError, "divisible"): TrajMixer(n_embd=121, n_head=10) def test_parameter_counts_match_traj_mixer_parameters(self) -> None: mixer = TrajMixer(n_embd=120, n_head=10) self.assertEqual( get_model_parameter_counts(mixer), { "model_parameter_count": 32_040, "trainable_parameter_count": 32_040, }, ) if __name__ == "__main__": unittest.main()