Initial pi0.5 minimal validation

tests/test_flow.py

@@ -0,0 +1,85 @@
+import torch
+
+from pi05mini.flow import FlowMatching
+
+
+def _expand_batch_time(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """Expand a per-example time tensor over non-batch dimensions."""
+    return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))
+
+
+def test_training_batch_obeys_linear_flow_equations() -> None:
+    flow = FlowMatching()
+    target = torch.tensor(
+        [
+            [[1.0, -2.0], [0.5, 3.0]],
+            [[-1.0, 0.0], [2.5, -0.5]],
+        ]
+    )
+    generator = torch.Generator().manual_seed(123)
+
+    x_t, t, target_velocity = flow.training_batch(target, generator=generator)
+
+    assert x_t.shape == target.shape
+    assert target_velocity.shape == target.shape
+    assert t.shape[0] == target.shape[0]
+    assert torch.all((t >= 0.0) & (t <= 1.0))
+
+    t_view = _expand_batch_time(t, target)
+    # v = target - noise, so the sampled noise is exactly recoverable.
+    noise = target - target_velocity
+    expected_x_t = (1.0 - t_view) * noise + t_view * target
+    torch.testing.assert_close(x_t, expected_x_t)
+
+
+def test_training_batch_is_reproducible_with_generator() -> None:
+    flow = FlowMatching()
+    target = torch.randn(3, 5, 2)
+
+    batch_a = flow.training_batch(
+        target, generator=torch.Generator().manual_seed(99)
+    )
+    batch_b = flow.training_batch(
+        target, generator=torch.Generator().manual_seed(99)
+    )
+
+    for tensor_a, tensor_b in zip(batch_a, batch_b):
+        torch.testing.assert_close(tensor_a, tensor_b)
+
+
+def test_loss_is_mean_squared_velocity_error() -> None:
+    flow = FlowMatching()
+    target_velocity = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
+    prediction = target_velocity + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
+
+    # Squared errors are [1, 1, 4, 0], whose mean is 1.5.
+    torch.testing.assert_close(flow.loss(prediction, target_velocity), torch.tensor(1.5))
+    torch.testing.assert_close(
+        flow.loss(target_velocity, target_velocity), torch.tensor(0.0)
+    )
+
+
+def test_euler_sampler_reaches_fixed_target_for_exact_vector_field() -> None:
+    flow = FlowMatching()
+    target = torch.tensor(
+        [
+            [[0.25, -0.5], [1.0, 2.0]],
+            [[-2.0, 0.75], [0.0, 0.5]],
+        ]
+    )
+
+    def exact_model(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        t_view = _expand_batch_time(t, x_t)
+        # Straight conditional path to target. The sampler evaluates t < 1.
+        return (target.to(x_t) - x_t) / (1.0 - t_view)
+
+    torch.manual_seed(5)
+    sampled = flow.sample(
+        exact_model,
+        shape=tuple(target.shape),
+        device=target.device,
+        steps=8,
+    )
+
+    torch.testing.assert_close(sampled, target, atol=1e-5, rtol=1e-5)
+