Add time attention mask handling and baseline class time computation to DeepHealth model

backbones.py

@@ -395,11 +395,28 @@ class BaselineEncoder(nn.Module):
             dtype=dtype,
         ).masked_fill(~mask[:, None, None, :], -1e4)
 
+    def _make_time_attn_mask(
+        self,
+        mask: torch.Tensor,
+        time: torch.Tensor,
+        dtype: torch.dtype,
+    ):
+        valid_key = mask[:, None, :]
+        visible_by_time = time[:, None, :] <= time[:, :, None]
+        valid = valid_key & visible_by_time
+        return torch.zeros(
+            valid.shape,
+            device=valid.device,
+            dtype=dtype,
+        ).masked_fill(~valid, -1e4)[:, None, :, :]
+
     def forward(
         self,
         other_type: torch.LongTensor,        # (B, K), 0 = padding
         other_value: torch.Tensor,           # (B, K), cate stores global id
         other_value_kind: torch.LongTensor,  # (B, K), 0=PAD, 1=CONT, 2=CATE
+        other_time: torch.Tensor | None = None,  # (B, K)
+        cls_time: torch.Tensor | None = None,    # (B,)
     ):
         if other_type.shape != other_value.shape:
             raise ValueError(
@@ -451,7 +468,24 @@ class BaselineEncoder(nn.Module):
         )
         full_valid = torch.cat([cls_valid, other_valid], dim=1)
 
-        attn_mask = self._make_attn_mask(full_valid, f.dtype)
+        if other_time is None:
+            attn_mask = self._make_attn_mask(full_valid, f.dtype)
+        else:
+            if other_time.shape != other_type.shape:
+                raise ValueError(
+                    "other_time must have the same shape as other_type, got "
+                    f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
+                )
+            if cls_time is None:
+                raise ValueError("cls_time is required when other_time is provided")
+            full_time = torch.cat(
+                [
+                    cls_time.to(device=other_time.device, dtype=other_time.dtype)[:, None],
+                    other_time,
+                ],
+                dim=1,
+            )
+            attn_mask = self._make_time_attn_mask(full_valid, full_time, f.dtype)
         for block in self.blocks:
             f = block(f, attn_mask=attn_mask)
             f = f * full_valid.unsqueeze(-1).to(f.dtype)