Make DeMemWM memory streams self-only

.exp_artifact/dememwm_future_memory_selection_plan.md

@@ -339,7 +339,7 @@ Allow noncausal DeMemWM dynamic events
 
 ## Substep 4: Make All Packed Memory Streams Diagonal/Self-Only
 
-Status: `[ ]`
+Status: `[x]`
 
 Goal:
 

algorithms/dememwm/models/attention.py

@@ -68,11 +68,8 @@ class TemporalAxialAttention(nn.Module):
             segment_slice = slice(cursor, cursor + length)
             segment_slices[segment] = segment_slice
             if length > 0:
-                if segment == "dynamic":
-                    allow[:, segment_slice, segment_slice] = True
-                else:
-                    idx = torch.arange(cursor, cursor + length, device=device)
-                    allow[:, idx, idx] = True
+                idx = torch.arange(cursor, cursor + length, device=device)
+                allow[:, idx, idx] = True
             cursor += length
 
         if frame_memory_masks is not None:

tests/test_dememwm_temporal_attention.py

@@ -40,6 +40,65 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
         expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 200.0])
         self.assertTrue(torch.allclose(out.flatten(), expected))
 
+    def test_frame_memory_attention_bias_keeps_target_causal_and_streams_self_only(self):
+        attn = _averaging_temporal_attention()
+        segments = {"target": 3, "anchor": 2, "dynamic": 2, "revisit": 2}
+        total_frames = sum(segments.values())
+
+        bias = attn._frame_memory_attn_bias(
+            B=1,
+            T=total_frames,
+            H=1,
+            W=1,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            frame_memory_segments=segments,
+            frame_memory_masks=None,
+        )[0, 0]
+        allow = torch.isfinite(bias)
+
+        target_expected = torch.tril(torch.ones((3, 3), dtype=torch.bool))
+        self.assertTrue(torch.equal(allow[:3, :3], target_expected))
+        self.assertFalse(allow[:3, 3:].any().item())
+
+        cursor = segments["target"]
+        for stream in ("anchor", "dynamic", "revisit"):
+            length = segments[stream]
+            rows = slice(cursor, cursor + length)
+            self.assertTrue(torch.equal(allow[rows, rows], torch.eye(length, dtype=torch.bool)))
+            self.assertFalse(allow[rows, :cursor].any().item())
+            self.assertFalse(allow[rows, cursor + length :].any().item())
+            cursor += length
+
+    def test_frame_memory_attention_bias_restores_invalid_memory_row_diagonal(self):
+        attn = _averaging_temporal_attention()
+        segments = {"target": 2, "anchor": 2, "dynamic": 2, "revisit": 2}
+        total_frames = sum(segments.values())
+        masks = {
+            "target": torch.ones((1, 2), dtype=torch.bool),
+            "anchor": torch.tensor([[False, True]]),
+            "dynamic": torch.tensor([[True, False]]),
+            "revisit": torch.tensor([[False, False]]),
+        }
+
+        bias = attn._frame_memory_attn_bias(
+            B=1,
+            T=total_frames,
+            H=1,
+            W=1,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            frame_memory_segments=segments,
+            frame_memory_masks=masks,
+        )[0, 0]
+        allow = torch.isfinite(bias)
+
+        self.assertTrue(allow.any(dim=-1).all().item())
+        for row in (2, 5, 6, 7):
+            row_expected = torch.zeros(total_frames, dtype=torch.bool)
+            row_expected[row] = True
+            self.assertTrue(torch.equal(allow[row], row_expected))
+
     def test_frame_memory_segments_mask_temporal_streams(self):
         attn = _averaging_temporal_attention()
         x = torch.tensor([10.0, 20.0, 30.0, 100.0, 1.0, 3.0, 200.0]).view(1, 7, 1, 1, 1)
@@ -47,7 +106,7 @@ class DeMemWMTemporalAttentionTests(unittest.TestCase):
 
         out = attn(x, frame_memory_segments=segments)
 
-        expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 2.0, 2.0, 200.0])
+        expected = torch.tensor([10.0, 15.0, 20.0, 100.0, 1.0, 3.0, 200.0])
         self.assertTrue(torch.allclose(out.flatten(), expected))
 
         masks = {