Pad attribution batches before concatenation

evaluate_single_disease_mortality_attribution.py

@@ -47,6 +47,7 @@ from evaluate_event_free_survival import (
     mortality_hazard_from_risk,
 )
 from future_risk import death_risk_from_probabilities, probabilities_from_logits
+from targets import PAD_IDX
 
 
 OUTPUT_COLUMNS = [
@@ -196,6 +197,48 @@ def write_summary_csv(
     return len(rows)
 
 
+def concat_padded_tensor_batches(chunks: list[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    if not chunks:
+        raise ValueError("Cannot concatenate an empty chunk list")
+
+    fill_values = {
+        "event_seq": PAD_IDX,
+        "time_seq": 0.0,
+        "readout_mask": False,
+        "padding_mask": False,
+        "other_type": 0,
+        "other_value": 0.0,
+        "other_value_kind": 0,
+        "other_time": 0.0,
+    }
+    out: Dict[str, torch.Tensor] = {}
+    for key in chunks[0]:
+        tensors = [chunk[key] for chunk in chunks]
+        shapes = [tuple(t.shape) for t in tensors]
+        if len(set(shapes)) == 1:
+            out[key] = torch.cat(tensors, dim=0)
+            continue
+
+        if any(t.ndim == 0 for t in tensors):
+            raise ValueError(f"Cannot concatenate scalar tensor key={key!r} with mismatched shapes")
+        max_shape = list(shapes[0])
+        for shape in shapes[1:]:
+            if len(shape) != len(max_shape):
+                raise ValueError(f"Cannot concatenate key={key!r} with shapes {shapes}")
+            max_shape = [max(a, b) for a, b in zip(max_shape, shape)]
+
+        padded: list[torch.Tensor] = []
+        fill = fill_values.get(key, 0)
+        for tensor in tensors:
+            target_shape = [int(tensor.shape[0]), *max_shape[1:]]
+            padded_tensor = tensor.new_full(target_shape, fill)
+            slices = tuple(slice(0, int(size)) for size in tensor.shape)
+            padded_tensor[slices] = tensor
+            padded.append(padded_tensor)
+        out[key] = torch.cat(padded, dim=0)
+    return out
+
+
 def load_disease_metadata(
     mapping_path: Path,
     *,
@@ -503,10 +546,7 @@ def main() -> None:
         if pending_n == 0:
             return
 
-        ablated_batch = {
+        ablated_batch = concat_padded_tensor_batches(pending_batch_chunks)
-            key: torch.cat([chunk[key] for chunk in pending_batch_chunks], dim=0)
-            for key in pending_batch_chunks[0]
-        }
         meta_rows = [row for chunk in pending_meta_chunks for row in chunk]
         with torch.no_grad():
             ablated_risk = death_risk_for_batch(