Refactor loss computation and model input handling for improved clarity and efficiency

models.py

@@ -260,12 +260,27 @@ class DeepHealth(nn.Module):
         other_time: torch.Tensor,
         other_mask: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        batch_size, _n_other, n_embd = h_other.shape
-        group_counts = [
-            int(torch.unique(other_time[b, other_mask[b]], sorted=True).numel())
-            for b in range(batch_size)
-        ]
-        max_groups = max(group_counts, default=0)
+        batch_size, n_other, n_embd = h_other.shape
+        if n_other == 0:
+            empty_h = h_other.new_zeros(batch_size, 0, n_embd)
+            empty_t = other_time.new_zeros(batch_size, 0)
+            empty_m = torch.zeros(batch_size, 0, dtype=torch.bool, device=h_other.device)
+            return empty_h, empty_t, empty_m
+
+        masked_time = other_time.masked_fill(~other_mask, float("inf"))
+        _sorted_time_with_pad, order = masked_time.sort(dim=1)
+        sorted_time = other_time.gather(1, order)
+        sorted_mask = other_mask.gather(1, order)
+        sorted_h = h_other.gather(1, order.unsqueeze(-1).expand(-1, -1, n_embd))
+
+        group_start = torch.zeros_like(sorted_mask)
+        group_start[:, 0] = sorted_mask[:, 0]
+        group_start[:, 1:] = sorted_mask[:, 1:] & (
+            sorted_time[:, 1:] != sorted_time[:, :-1]
+        )
+        group_id = group_start.long().cumsum(dim=1) - 1
+        max_groups = int(group_start.sum(dim=1).max().item())
+
         pooled_h = h_other.new_zeros(batch_size, max_groups, n_embd)
         pooled_time = other_time.new_zeros(batch_size, max_groups)
         pooled_mask = torch.zeros(
@@ -277,35 +292,29 @@ class DeepHealth(nn.Module):
         if max_groups == 0:
             return pooled_h, pooled_time, pooled_mask
 
-        for b in range(batch_size):
-            valid_time = other_time[b, other_mask[b]]
-            if valid_time.numel() == 0:
-                continue
-            valid_h = h_other[b, other_mask[b]]
-            unique_time, inverse = torch.unique(
-                valid_time,
-                sorted=True,
-                return_inverse=True,
+        safe_group_id = group_id.clamp_min(0)
+        pooled_h.scatter_add_(
+            1,
+            safe_group_id.unsqueeze(-1).expand_as(sorted_h),
+            sorted_h * sorted_mask.unsqueeze(-1).to(sorted_h.dtype),
+        )
+        if self.extra_pool_reduce == "mean":
+            counts = h_other.new_zeros(batch_size, max_groups, 1)
+            counts.scatter_add_(
+                1,
+                safe_group_id.unsqueeze(-1),
+                sorted_mask.unsqueeze(-1).to(h_other.dtype),
             )
-            n_groups = unique_time.numel()
-            group_h = valid_h.new_zeros(n_groups, n_embd)
-            group_h.scatter_add_(
-                0,
-                inverse[:, None].expand(-1, n_embd),
-                valid_h,
-            )
-            if self.extra_pool_reduce == "mean":
-                counts = valid_h.new_zeros(n_groups, 1)
-                counts.scatter_add_(
-                    0,
-                    inverse[:, None],
-                    torch.ones_like(valid_h[:, :1]),
-                )
-                group_h = group_h / counts.clamp_min(1.0)
+            pooled_h = pooled_h / counts.clamp_min(1.0)
 
-            pooled_h[b, :n_groups] = group_h
-            pooled_time[b, :n_groups] = unique_time
-            pooled_mask[b, :n_groups] = True
+        pooled_time.scatter_add_(
+            1,
+            safe_group_id,
+            sorted_time * group_start.to(sorted_time.dtype),
+        )
+        group_count = group_start.sum(dim=1)
+        arange_groups = torch.arange(max_groups, device=h_other.device)
+        pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1)
         return pooled_h, pooled_time, pooled_mask
 
     def _forward_shared(