Fix ANHN construction and sampled-state shapes

ANHN was unusable: it subscripted ModelConfig (which has no
__getitem__), so construction raised TypeError, and
compute_states_at_sample_times computed the same tensor seq_len times
in a loop and stacked the copies, returning a 5-D tensor that crashed
loglike_loss. Use attribute access on the config and return the
single (already correctly broadcast) state tensor. Add construction,
loss/backward, and shape tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: iLampard

easy_tpp/model/torch_model/torch_anhn.py

@@ -18,12 +18,12 @@ def __init__(self, model_config):
         """
         super(ANHN, self).__init__(model_config)
 
-        self.d_time = model_config['time_emb_size']
-        self.use_norm = model_config['use_ln']
+        self.d_time = model_config.time_emb_size
+        self.use_norm = model_config.use_ln
 
-        self.n_layers = model_config['num_layers']
-        self.n_head = model_config['num_heads']
-        self.dropout = model_config['dropout']
+        self.n_layers = model_config.num_layers
+        self.n_head = model_config.num_heads
+        self.dropout = model_config.dropout_rate
 
         self.layer_rnn = nn.LSTM(input_size=self.hidden_size, hidden_size=self.hidden_size, batch_first=True)
 
@@ -230,14 +230,8 @@ def compute_states_at_sample_times(self, intensity_base, intensity_alpha, intens
         # [batch_size, seq_len, num_samples, 1, 1]
         sample_dtimes_ = sample_dtimes[:, :, :, None, None]
 
-        states_samples = []
-        seq_len = intensity_base.size()[1]
-        for _ in range(seq_len):
-            states_samples_ = self.compute_states_at_event_times(mu, alpha, delta, base_elapses + sample_dtimes_)
-            states_samples.append(states_samples_)
-
         # [batch_size, seq_len, num_sample, hidden_size]
-        states_samples = torch.stack(states_samples, dim=1)
+        states_samples = self.compute_states_at_event_times(mu, alpha, delta, base_elapses + sample_dtimes_)
         return states_samples
 
     def compute_intensities_at_sample_times(self, time_seqs, time_delta_seqs, type_seqs, sample_dtimes, **kwargs):

tests/test_anhn.py

@@ -0,0 +1,74 @@
+import torch
+
+from easy_tpp.config_factory import ModelConfig
+from easy_tpp.model import TorchANHN
+
+
+def make_model_config():
+    return ModelConfig.parse_from_yaml_config({
+        'model_id': 'ANHN',
+        'hidden_size': 8,
+        'time_emb_size': 8,
+        'num_layers': 1,
+        'num_heads': 1,
+        'num_event_types': 3,
+        'num_event_types_pad': 4,
+        'event_pad_index': 3,
+        'gpu': -1,
+        'loss_integral_num_sample_per_step': 5,
+    })
+
+
+def make_batch():
+    time_delta_seqs = torch.tensor([
+        [0.0, 0.4, 0.6, 0.5, 0.7],
+        [0.0, 0.3, 0.2, 0.6, 0.4],
+    ])
+    time_seqs = torch.cumsum(time_delta_seqs, dim=-1)
+    type_seqs = torch.tensor([
+        [0, 1, 2, 0, 1],
+        [2, 1, 0, 2, 1],
+    ], dtype=torch.long)
+    seq_non_pad_mask = torch.ones_like(type_seqs, dtype=torch.bool)
+    attention_mask = torch.triu(torch.ones(2, 5, 5, dtype=torch.bool), diagonal=1)
+    return {
+        'time_seqs': time_seqs,
+        'time_delta_seqs': time_delta_seqs,
+        'type_seqs': type_seqs,
+        'seq_non_pad_mask': seq_non_pad_mask,
+        'attention_mask': attention_mask,
+    }
+
+
+def test_construction():
+    TorchANHN(make_model_config())
+
+
+def test_loglike_loss_runs():
+    model = TorchANHN(make_model_config())
+
+    loss, num_events = model.loglike_loss(**make_batch())
+
+    assert torch.isfinite(loss)
+    assert num_events > 0
+    assert loss.requires_grad
+    loss.backward()
+
+
+def test_sample_states_shape():
+    model = TorchANHN(make_model_config())
+    batch = make_batch()
+    dtime_seqs = batch['time_delta_seqs'][:, 1:]
+    type_seqs = batch['type_seqs'][:, :-1]
+    attention_mask = batch['attention_mask'][:, 1:, :-1]
+
+    _, (intensity_base, intensity_alpha, intensity_delta), (base_dtime, _) = model.forward(
+        dtime_seqs, type_seqs, attention_mask
+    )
+    sample_dtimes = model.make_dtime_loss_samples(dtime_seqs)
+    states = model.compute_states_at_sample_times(
+        intensity_base, intensity_alpha, intensity_delta, base_dtime, sample_dtimes
+    )
+
+    assert states.ndim == 4
+    assert states.shape == (2, 4, 5, 8)