Initial pi0.6 minimal validation

tests/test_recap.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import torch
+
+from pi06mini.recap import (
+    apply_condition_dropout,
+    classifier_free_guidance,
+    condition_dropout_mask,
+)
+
+
+def test_default_condition_dropout_is_statistically_thirty_percent() -> None:
+    generator = torch.Generator().manual_seed(11)
+    mask = condition_dropout_mask(20_000, generator=generator)
+
+    assert mask.dtype == torch.bool
+    assert abs(float(mask.float().mean()) - 0.3) < 0.015
+
+
+def test_condition_dropout_replaces_whole_examples() -> None:
+    conditions = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+    dropped, mask = apply_condition_dropout(
+        conditions,
+        null_condition=-1.0,
+        mask=torch.tensor([False, True, False]),
+    )
+
+    assert mask.tolist() == [False, True, False]
+    torch.testing.assert_close(dropped[0], conditions[0])
+    torch.testing.assert_close(dropped[1], torch.tensor([-1.0, -1.0]))
+    torch.testing.assert_close(dropped[2], conditions[2])
+
+
+def test_cfg_interpolates_and_extrapolates_flow_predictions() -> None:
+    conditional = torch.tensor([[2.0, 4.0]])
+    unconditional = torch.tensor([[0.0, 2.0]])
+
+    torch.testing.assert_close(
+        classifier_free_guidance(conditional, unconditional, guidance_scale=0.0),
+        unconditional,
+    )
+    torch.testing.assert_close(
+        classifier_free_guidance(conditional, unconditional, guidance_scale=1.0),
+        conditional,
+    )
+    torch.testing.assert_close(
+        classifier_free_guidance(conditional, unconditional, guidance_scale=1.5),
+        torch.tensor([[3.0, 5.0]]),
+    )