Initial pi0 minimal validation

tests/test_model.py

@@ -0,0 +1,65 @@
+import torch
+
+from pi0mini.model import Pi0Config, Pi0Mini
+
+
+def test_pi0_three_block_attention_mask() -> None:
+    config = Pi0Config(
+        image_size=32,
+        patch_size=16,
+        max_text_tokens=4,
+        max_action_tokens=3,
+        d_model=32,
+        n_heads=4,
+        n_layers=1,
+        time_feature_dim=8,
+    )
+    model = Pi0Mini(config)
+
+    # Two image patches per axis -> 4 image tokens, plus 2 language tokens.
+    prefix_edges, suffix_edges = model._masks(
+        batch=1,
+        prefix_length=6,
+        suffix_length=4,  # one state + three action tokens
+        image_tokens=4,
+        attention_mask=torch.tensor([[1, 1]]),
+        device=torch.device("cpu"),
+    )
+
+    assert prefix_edges.shape == (1, 6, 6)
+    assert torch.all(prefix_edges)
+
+    # Suffix key order is [six prefix keys, state, action0, action1, action2].
+    assert suffix_edges.shape == (1, 4, 10)
+    assert torch.all(suffix_edges[0, 0, :7])
+    assert not torch.any(suffix_edges[0, 0, 7:])
+    assert torch.all(suffix_edges[0, 1:, :])
+
+
+def test_pi0_mask_respects_padded_language_keys() -> None:
+    model = Pi0Mini(
+        Pi0Config(
+            image_size=16,
+            patch_size=16,
+            max_text_tokens=3,
+            max_action_tokens=1,
+            d_model=16,
+            n_heads=4,
+            n_layers=1,
+            time_feature_dim=8,
+        )
+    )
+    prefix_edges, suffix_edges = model._masks(
+        batch=1,
+        prefix_length=4,
+        suffix_length=2,
+        image_tokens=1,
+        attention_mask=torch.tensor([[1, 0, 0]]),
+        device=torch.device("cpu"),
+    )
+
+    # Padded prefix keys are hidden from valid prefix, state, and action rows.
+    assert not prefix_edges[0, 0, 2]
+    assert not prefix_edges[0, 0, 3]
+    assert not suffix_edges[0, :, 2].any()
+    assert not suffix_edges[0, :, 3].any()