Validate DeMemWM DiT stream boundaries

.exp_artifact/dememwm_remaining_fix_plan.md

@@ -232,7 +232,7 @@ Reject stale DeMemWM experiment flags
 
 ## Substep 4: Add Minimal DeMemWM DiT Boundary Checks
 
-Status: `[ ]`
+Status: `[x]`
 
 Bug:
 

algorithms/dememwm/models/dit.py

@@ -356,18 +356,42 @@ class SpatioTemporalDiTBlock(nn.Module):
         anchor = int(frame_memory_segments.get("anchor", 0))
         dynamic = int(frame_memory_segments.get("dynamic", 0))
         revisit = int(frame_memory_segments.get("revisit", 0))
+        if min(target, anchor, dynamic, revisit) < 0:
+            raise ValueError(
+                "frame_memory_segments lengths must be nonnegative; "
+                f"got target={target}, anchor={anchor}, dynamic={dynamic}, revisit={revisit}"
+            )
+        total = target + anchor + dynamic + revisit
+        if total != int(x.shape[1]):
+            raise ValueError(
+                f"frame_memory_segments lengths sum to {total}, expected x.shape[1]={int(x.shape[1])}"
+            )
         a0, a1 = target, target + anchor
         d0, d1 = a1, a1 + dynamic
         r0, r1 = d1, d1 + revisit
         return x[:, :target], x[:, a0:a1], x[:, d0:d1], x[:, r0:r1]
 
     def _frame_memory_stream_mask(self, frame_memory_masks, stream_name, stream_hidden):
+        self._check_frame_memory_stream_mask_shape(frame_memory_masks, stream_name, stream_hidden)
         if stream_hidden is None or int(stream_hidden.shape[1]) == 0:
             return None
         if frame_memory_masks is None or frame_memory_masks.get(stream_name) is None:
             return None
         return frame_memory_masks[stream_name].to(device=stream_hidden.device, dtype=torch.bool)
 
+    def _check_frame_memory_stream_mask_shape(self, frame_memory_masks, stream_name, stream_hidden):
+        if stream_hidden is None:
+            return
+        if frame_memory_masks is None or frame_memory_masks.get(stream_name) is None:
+            return
+        stream_mask = frame_memory_masks[stream_name]
+        expected_shape = (int(stream_hidden.shape[0]), int(stream_hidden.shape[1]))
+        if tuple(stream_mask.shape) != expected_shape:
+            raise ValueError(
+                f"frame_memory_masks[{stream_name!r}] shape {tuple(stream_mask.shape)} "
+                f"must match {expected_shape}"
+            )
+
     def _apply_frame_memory_reference_attention(self, x, c, frame_memory_segments, frame_memory_masks, frame_memory_geometry=None):
         x_target, x_anchor, x_dynamic, x_revisit = self._split_frame_memory(x, frame_memory_segments)
         if int(x_target.shape[1]) == 0:
@@ -412,6 +436,12 @@ class SpatioTemporalDiTBlock(nn.Module):
         frame_memory_segments=None, frame_memory_masks=None, frame_memory_pose=None,
         image_hw=None, frame_memory_geometry=None):
         B, T, H, W, D = x.shape
+        if frame_memory_segments is not None:
+            for stream_name, stream_hidden in zip(
+                ("target", "anchor", "dynamic", "revisit"),
+                self._split_frame_memory(x, frame_memory_segments),
+            ):
+                self._check_frame_memory_stream_mask_shape(frame_memory_masks, stream_name, stream_hidden)
 
         # spatial block
         

tests/test_dememwm_temporal_attention.py

@@ -108,6 +108,70 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         self.assertIs(spy.calls[0][0], segments)
         self.assertIs(spy.calls[0][1], masks)
 
+    def test_block_rejects_bad_frame_memory_segment_lengths_before_temporal_attention(self):
+        class ZeroAttention(nn.Module):
+            def forward(self, x):
+                return torch.zeros_like(x)
+
+        class UnexpectedTemporalAttention(nn.Module):
+            def forward(self, x, frame_memory_segments=None, frame_memory_masks=None):
+                raise AssertionError("temporal attention should not run before frame-memory validation")
+
+        block = SpatioTemporalDiTBlock(
+            hidden_size=4,
+            num_heads=1,
+            reference_length=0,
+            spatial_rotary_emb=None,
+            temporal_rotary_emb=None,
+        )
+        block.s_attn = ZeroAttention()
+        block.t_attn = UnexpectedTemporalAttention()
+        x = torch.zeros((1, 5, 1, 1, 4))
+        c = torch.zeros((1, 5, 4))
+
+        with self.assertRaisesRegex(ValueError, "lengths must be nonnegative"):
+            block(x, c, frame_memory_segments={"target": 2, "anchor": -1, "dynamic": 2, "revisit": 2})
+
+        with self.assertRaisesRegex(ValueError, r"sum to 6, expected x\.shape\[1\]=5"):
+            block(x, c, frame_memory_segments={"target": 2, "anchor": 1, "dynamic": 2, "revisit": 1})
+
+    def test_block_rejects_bad_frame_memory_stream_mask_shape_before_temporal_attention(self):
+        class ZeroAttention(nn.Module):
+            def forward(self, x):
+                return torch.zeros_like(x)
+
+        class UnexpectedTemporalAttention(nn.Module):
+            def forward(self, x, frame_memory_segments=None, frame_memory_masks=None):
+                raise AssertionError("temporal attention should not run before frame-memory validation")
+
+        block = SpatioTemporalDiTBlock(
+            hidden_size=4,
+            num_heads=1,
+            reference_length=0,
+            spatial_rotary_emb=None,
+            temporal_rotary_emb=None,
+        )
+        block.s_attn = ZeroAttention()
+        block.t_attn = UnexpectedTemporalAttention()
+        x = torch.zeros((1, 5, 1, 1, 4))
+        c = torch.zeros((1, 5, 4))
+        segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1}
+        masks = {"dynamic": torch.ones((1, 2), dtype=torch.bool)}
+
+        with self.assertRaisesRegex(
+            ValueError,
+            r"frame_memory_masks\[\'dynamic\'\] shape \(1, 2\) must match \(1, 1\)",
+        ):
+            block(x, c, frame_memory_segments=segments, frame_memory_masks=masks)
+
+        zero_dynamic_segments = {"target": 2, "anchor": 1, "dynamic": 0, "revisit": 2}
+        zero_dynamic_masks = {"dynamic": torch.ones((1, 1), dtype=torch.bool)}
+        with self.assertRaisesRegex(
+            ValueError,
+            r"frame_memory_masks\[\'dynamic\'\] shape \(1, 1\) must match \(1, 0\)",
+        ):
+            block(x, c, frame_memory_segments=zero_dynamic_segments, frame_memory_masks=zero_dynamic_masks)
+
     def test_dit_threads_frame_memory_metadata_to_blocks(self):
         class SpyBlock(nn.Module):
             def __init__(self):