import math import unittest import torch from backbones import ( SharedEventTrajectoryCore, SharedTrajectoryMixer, TrajectoryCrossAttention, ) from models import ( EVENT_TRAJECTORY_ARCHITECTURE, MODEL_SIZE_PRESETS, DeepHealth, resolve_model_size, validate_event_trajectory_config, validate_event_trajectory_state_dict, ) from train_util import get_model_parameter_counts def build_test_model( *, target_mode: str = "next_token", time_mode: str = "absolute", n_reasoning_rounds: int = 3, ) -> DeepHealth: return DeepHealth( vocab_size=32, model_size="nano", n_reasoning_rounds=n_reasoning_rounds, n_types=2, n_cont_types=0, n_categories=2, cont_type_ids=[], target_mode=target_mode, time_mode=time_mode, ) def model_inputs() -> dict[str, torch.Tensor]: return { "event_seq": torch.tensor([[1, 2, 3, 4], [5, 6, 0, 0]]), "time_seq": torch.tensor( [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 0.0, 0.0]] ), "sex": torch.tensor([0, 1]), "padding_mask": torch.tensor( [[True, True, True, True], [True, True, False, False]] ), "other_type": torch.zeros(2, 1, dtype=torch.long), "other_value": torch.zeros(2, 1), "other_value_kind": torch.zeros(2, 1, dtype=torch.long), "other_time": torch.zeros(2, 1), } class EventTrajectoryBackboneTest(unittest.TestCase): def test_model_size_presets(self) -> None: expected = { "nano": (120, 6, 20, 24), "tiny": (256, 8, 32, 32), "small": (512, 8, 64, 32), "medium": (768, 12, 64, 48), "huge": (1024, 16, 64, 64), } self.assertEqual(set(MODEL_SIZE_PRESETS), set(expected)) for name, values in expected.items(): preset = resolve_model_size(name) self.assertEqual( ( preset.d_model, preset.n_trajectory, preset.trajectory_dim, preset.traj_hidden, ), values, ) def test_default_mixer_shapes_and_parameter_count(self) -> None: mixer = SharedTrajectoryMixer( n_trajectory=8, trajectory_dim=32, ) state = torch.randn(2, 5, 8, 32) self.assertEqual(mixer(state).shape, state.shape) self.assertEqual(mixer.traj_hidden, 32) self.assertEqual(tuple(mixer.gate_proj.shape), (32, 8, 32)) self.assertEqual(tuple(mixer.value_proj.shape), (32, 8, 32)) self.assertEqual(tuple(mixer.output_proj.shape), (32, 32, 8)) self.assertEqual( get_model_parameter_counts(mixer), { "model_parameter_count": 24_576, "trainable_parameter_count": 24_576, }, ) def test_reasoning_rounds_share_one_core_parameter_set(self) -> None: core_one = SharedEventTrajectoryCore( d_model=256, n_trajectory=8, n_reasoning_rounds=1, ) core_twelve = SharedEventTrajectoryCore( d_model=256, n_trajectory=8, n_reasoning_rounds=12, ) self.assertEqual( sum(p.numel() for p in core_one.parameters()), sum(p.numel() for p in core_twelve.parameters()), ) self.assertAlmostEqual(core_one.attn_scale.item(), 1.0) self.assertAlmostEqual( core_twelve.attn_scale.item(), 1.0 / math.sqrt(12), places=6, ) self.assertAlmostEqual( core_twelve.mixer_scale.item(), 1.0 / math.sqrt(12), places=6, ) def test_all_masked_attention_is_finite_and_zero(self) -> None: attention = TrajectoryCrossAttention( d_model=32, n_trajectory=4, ) memory = torch.randn(2, 3, 32) key_value = attention.project_event_memory(memory) state = torch.randn(2, 2, 4, 8) invalid_mask = torch.ones(2, 2, 3, dtype=torch.bool) output = attention( trajectory_state=state, event_key_value=key_value, event_invalid_mask=invalid_mask, ) self.assertTrue(torch.isfinite(output).all()) torch.testing.assert_close(output, torch.zeros_like(output)) def test_next_token_future_events_do_not_change_earlier_query(self) -> None: torch.manual_seed(0) model = build_test_model(n_reasoning_rounds=2) model.eval() inputs = model_inputs() original = model(**inputs) changed_inputs = dict(inputs) changed_inputs["event_seq"] = inputs["event_seq"].clone() changed_inputs["event_seq"][0, 3] = 9 changed = model(**changed_inputs) torch.testing.assert_close(original[0, 1], changed[0, 1]) def test_next_token_later_equal_time_event_is_not_visible(self) -> None: torch.manual_seed(0) model = build_test_model(n_reasoning_rounds=2) model.eval() inputs = model_inputs() inputs["time_seq"] = inputs["time_seq"].clone() inputs["time_seq"][0] = torch.tensor([1.0, 1.0, 2.0, 3.0]) original = model(**inputs) changed_inputs = dict(inputs) changed_inputs["event_seq"] = inputs["event_seq"].clone() changed_inputs["event_seq"][0, 1] = 9 changed = model(**changed_inputs) torch.testing.assert_close(original[0, 0], changed[0, 0]) def test_padding_content_does_not_change_valid_queries(self) -> None: torch.manual_seed(0) model = build_test_model( time_mode="relative", n_reasoning_rounds=2, ) model.eval() inputs = model_inputs() original = model(**inputs) changed_inputs = dict(inputs) changed_inputs["event_seq"] = inputs["event_seq"].clone() changed_inputs["time_seq"] = inputs["time_seq"].clone() changed_inputs["event_seq"][1, 2:] = torch.tensor([9, 10]) changed_inputs["time_seq"][1, 2:] = torch.tensor([30.0, 40.0]) changed = model(**changed_inputs) torch.testing.assert_close(original[1, :2], changed[1, :2]) def test_next_token_and_all_future_output_contracts(self) -> None: inputs = model_inputs() next_model = build_test_model(target_mode="next_token") next_hidden = next_model(**inputs) self.assertEqual(tuple(next_hidden.shape), (2, 4, 120)) next_output = next_model(**inputs, return_output=True) self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120)) self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4)) future_model = build_test_model(target_mode="all_future") future_hidden = future_model( **inputs, t_query=torch.tensor([5.0, 3.0]), ) self.assertEqual(tuple(future_hidden.shape), (2, 120)) def test_model_contains_one_shared_core_and_no_block_stack(self) -> None: model = build_test_model(n_reasoning_rounds=12) self.assertFalse(hasattr(model, "blocks")) reasoning_keys = [ key for key in model.state_dict() if key.startswith("reasoning_core.") ] self.assertTrue(reasoning_keys) self.assertFalse(any("blocks." in key for key in model.state_dict())) self.assertFalse(any("out_proj" in key for key in reasoning_keys)) self.assertFalse(any("group_align" in key for key in reasoning_keys)) def test_event_key_and_value_are_projected_once_per_forward(self) -> None: model = build_test_model(n_reasoning_rounds=12) call_counts = {"key": 0, "value": 0} def count_key(*_args) -> None: call_counts["key"] += 1 def count_value(*_args) -> None: call_counts["value"] += 1 key_handle = ( model.reasoning_core.cross_attention.k_proj .register_forward_hook(count_key) ) value_handle = ( model.reasoning_core.cross_attention.v_proj .register_forward_hook(count_value) ) try: model(**model_inputs()) finally: key_handle.remove() value_handle.remove() self.assertEqual(call_counts, {"key": 1, "value": 1}) def test_relative_time_forward_and_backward_are_finite(self) -> None: torch.manual_seed(0) model = build_test_model( target_mode="all_future", time_mode="relative", n_reasoning_rounds=2, ) hidden = model( **model_inputs(), t_query=torch.tensor([5.0, 3.0]), ) (hidden * torch.randn_like(hidden)).sum().backward() self.assertTrue(torch.isfinite(hidden).all()) self.assertIsNotNone(model.event_projection.weight.grad) self.assertTrue(torch.isfinite(model.event_projection.weight.grad).all()) time_scale = model.reasoning_core.cross_attention.time_bias_scale self.assertIsNotNone(time_scale) self.assertIsNotNone(time_scale.grad) self.assertGreater(abs(float(time_scale.grad)), 0.0) def test_architecture_marker_and_checkpoint_are_required(self) -> None: validate_event_trajectory_config( { "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE, "model_size": "nano", "d_model": 120, "n_trajectory": 6, "trajectory_dim": 20, "traj_hidden": 24, "n_reasoning_rounds": 3, } ) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_event_trajectory_config( {"model_architecture": "traj_mixer_v2"} ) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_event_trajectory_config( {"model_architecture": "event_trajectory_shared_v1"} ) with self.assertRaisesRegex(ValueError, "trajectory_dim"): validate_event_trajectory_config( { "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE, "model_size": "nano", "d_model": 120, "n_trajectory": 6, "trajectory_dim": 10, "traj_hidden": 24, "n_reasoning_rounds": 3, } ) model = build_test_model() state_dict = model.state_dict() validate_event_trajectory_state_dict( state_dict, expected_d_model=120, expected_n_trajectory=6, expected_n_reasoning_rounds=3, ) with self.assertRaisesRegex( ValueError, "Checkpoint architecture does not match", ): validate_event_trajectory_state_dict( state_dict, expected_n_reasoning_rounds=12, ) state_dict.pop("reasoning_core.attn_scale") with self.assertRaisesRegex( ValueError, "not a shared event-trajectory checkpoint", ): validate_event_trajectory_state_dict(state_dict) def test_unknown_model_size_is_rejected(self) -> None: with self.assertRaisesRegex(ValueError, "Unknown model_size"): DeepHealth( vocab_size=32, model_size="giant", n_reasoning_rounds=2, n_types=2, n_cont_types=0, n_categories=2, cont_type_ids=[], ) if __name__ == "__main__": unittest.main()