From e7b8cec56cd6a7f595598d1a83a1977bbcf0f5af Mon Sep 17 00:00:00 2001 From: Alex Tong Date: Mon, 20 Jul 2026 14:17:57 +0200 Subject: [PATCH] test: expand test suite with fixtures, utils, integration, and model forward-pass tests - Add tests/conftest.py with shared fixtures: shapes, rng, sample_data, and one fixture per CFM matcher class - Add tests/test_utils.py covering eight_normal_sample, sample_moons, sample_8gaussians, and torch_wrapper - Add tests/test_integration.py with a @pytest.mark.slow end-to-end CFM training smoke test (asserts loss decreases) - Expand tests/test_models.py with forward-pass tests for MLP (time-varying/non-time-varying/out_dim) and GradModel (instantiation, forward, gradient correctness) - Fix typo 'classers' -> 'classes' in test_conditional_flow_matcher.py and test_optimal_transport.py docstrings --- tests/conftest.py | 66 ++++++++++++++ tests/test_conditional_flow_matcher.py | 2 +- tests/test_integration.py | 53 +++++++++++ tests/test_models.py | 65 ++++++++++++++ tests/test_optimal_transport.py | 2 +- tests/test_utils.py | 118 +++++++++++++++++++++++++ 6 files changed, 304 insertions(+), 2 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/test_integration.py create mode 100644 tests/test_utils.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..bd1b28f9 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,66 @@ +"""Shared fixtures for the TorchCFM test suite.""" + +import pytest +import torch + +from torchcfm.conditional_flow_matching import ( # noqa: F401 (star import below) + ConditionalFlowMatcher, + ExactOptimalTransportConditionalFlowMatcher, + SchrodingerBridgeConditionalFlowMatcher, + TargetConditionalFlowMatcher, + VariancePreservingConditionalFlowMatcher, +) +from torchcfm.conditional_flow_matching import * # noqa: F401,F403 + +# A collection of (batch, *feature) shapes used to parametrize tests across the suite. +SHAPES = [(2, 3), (5, 2), (10, 1), (3, 3, 2)] + + +@pytest.fixture(params=SHAPES) +def shapes(request): + """Parametrized fixture yielding each (batch, *feature) shape in ``SHAPES``.""" + return request.param + + +@pytest.fixture +def rng(): + """A seeded ``torch.Generator`` (seed 42) for reproducible randomness.""" + return torch.Generator().manual_seed(42) + + +@pytest.fixture +def sample_data(rng): + """Small pair of standard-normal tensors ``(x0, x1)`` of shape ``(64, 2)``.""" + x0 = torch.randn(64, 2, generator=rng) + x1 = torch.randn(64, 2, generator=rng) + return x0, x1 + + +@pytest.fixture +def conditional_flow_matcher(): + """Base independent conditional flow matcher (sigma=0.0).""" + return ConditionalFlowMatcher(sigma=0.0) + + +@pytest.fixture +def exact_ot_flow_matcher(): + """Exact optimal-transport conditional flow matcher (sigma=0.0).""" + return ExactOptimalTransportConditionalFlowMatcher(sigma=0.0) + + +@pytest.fixture +def target_flow_matcher(): + """Target conditional flow matcher (sigma=0.0).""" + return TargetConditionalFlowMatcher(sigma=0.0) + + +@pytest.fixture +def schrodinger_bridge_flow_matcher(): + """Schrödinger bridge conditional flow matcher (sigma must be > 0).""" + return SchrodingerBridgeConditionalFlowMatcher(sigma=1.0) + + +@pytest.fixture +def variance_preserving_flow_matcher(): + """Variance-preserving (trigonometric interpolant) flow matcher (sigma=0.0).""" + return VariancePreservingConditionalFlowMatcher(sigma=0.0) diff --git a/tests/test_conditional_flow_matcher.py b/tests/test_conditional_flow_matcher.py index b080470f..edc2bd00 100644 --- a/tests/test_conditional_flow_matcher.py +++ b/tests/test_conditional_flow_matcher.py @@ -1,4 +1,4 @@ -"""Tests for Conditional Flow Matcher classers.""" +"""Tests for Conditional Flow Matcher classes.""" # Author: Kilian Fatras diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 00000000..8055b202 --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,53 @@ +"""End-to-end smoke tests for training a tiny CFM model.""" + +import numpy as np +import pytest +import torch + +from torchcfm.conditional_flow_matching import ConditionalFlowMatcher +from torchcfm.models.models import MLP +from torchcfm.utils import sample_moons + + +@pytest.mark.slow +def test_cfm_training_reduces_loss(): + """Train a small MLP with CFM on two moons and assert the loss decreases.""" + # Arrange — seed everything for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + n_samples = 256 + x0 = torch.as_tensor(sample_moons(n_samples)).float() + x1 = torch.as_tensor(sample_moons(n_samples)).float() + + model = MLP(dim=2, w=32, time_varying=True) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) + fm = ConditionalFlowMatcher(sigma=0.0) + + # Fixed evaluation batch (seeded so initial/final loss are comparable) + torch.manual_seed(0) + t_eval, xt_eval, ut_eval = fm.sample_location_and_conditional_flow(x0, x1) + t_eval = t_eval.reshape(-1, 1) + + def eval_loss(): + vt = model(torch.cat([xt_eval, t_eval], dim=1)) + return torch.nn.functional.mse_loss(vt, ut_eval).item() + + # Act — measure initial loss + initial_loss = eval_loss() + + # Train for a small number of steps + n_steps = 100 + for _ in range(n_steps): + t, xt, ut = fm.sample_location_and_conditional_flow(x0, x1) + t = t.reshape(-1, 1) + vt = model(torch.cat([xt, t], dim=1)) + loss = torch.nn.functional.mse_loss(vt, ut) + optimizer.zero_grad() + loss.backward() + optimizer.step() + + final_loss = eval_loss() + + # Assert — training should reduce the loss on the fixed evaluation batch + assert final_loss < initial_loss diff --git a/tests/test_models.py b/tests/test_models.py index 6d82f986..364ac785 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,4 +1,9 @@ +"""Tests for models in ``torchcfm.models``.""" + +import torch + from torchcfm.models import MLP +from torchcfm.models.models import GradModel from torchcfm.models.unet import UNetModel @@ -11,3 +16,63 @@ def test_initialize_models(): class_cond=True, ) MLP(dim=2, time_varying=True, w=64) + + +def test_mlp_forward_time_varying(): + # Arrange — time_varying MLP expects input of dim (dim + 1) + model = MLP(dim=2, w=32, time_varying=True) + x = torch.randn(10, 2) + t = torch.rand(10, 1) + xt = torch.cat([x, t], dim=1) + # Act + out = model(xt) + # Assert + assert out.shape == (10, 2) + assert torch.isfinite(out).all() + + +def test_mlp_forward_not_time_varying(): + # Arrange — non-time-varying MLP takes input of dim directly + model = MLP(dim=2, w=32, time_varying=False) + x = torch.randn(10, 2) + # Act + out = model(x) + # Assert + assert out.shape == (10, 2) + assert torch.isfinite(out).all() + + +def test_mlp_forward_out_dim(): + # Arrange — explicit out_dim different from dim + model = MLP(dim=2, out_dim=4, w=32, time_varying=False) + x = torch.randn(10, 2) + # Act + out = model(x) + # Assert + assert out.shape == (10, 4) + + +def test_grad_model_instantiation_and_forward(): + # Arrange — GradModel wraps an "action" and returns the gradient (minus last dim) + action = MLP(dim=3, w=32, time_varying=False) + grad_model = GradModel(action) + x = torch.randn(10, 3) + # Act + grad = grad_model(x) + # Assert — gradient w.r.t. x has shape (10, 3); [:, :-1] yields (10, 2) + assert grad.shape == (10, 2) + assert torch.isfinite(grad).all() + + +def test_grad_model_is_gradient(): + # Arrange — verify the output is actually the gradient of the action + action = MLP(dim=3, w=32, time_varying=False) + grad_model = GradModel(action) + x = torch.randn(5, 3).clone().requires_grad_(True) + # Act + grad = grad_model(x) + # Manually compute the expected gradient + loss = torch.sum(action(x)) + (expected_grad,) = torch.autograd.grad(loss, x, create_graph=True) + # Assert — GradModel returns all but the last column of the gradient + assert torch.allclose(grad, expected_grad[:, :-1]) diff --git a/tests/test_optimal_transport.py b/tests/test_optimal_transport.py index ad696e8b..157b5c6f 100644 --- a/tests/test_optimal_transport.py +++ b/tests/test_optimal_transport.py @@ -1,4 +1,4 @@ -"""Tests for Conditional Flow Matcher classers.""" +"""Tests for Conditional Flow Matcher classes.""" # Author: Kilian Fatras diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..5f4d5c13 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,118 @@ +"""Tests for utility functions in ``torchcfm.utils``.""" + +import numpy as np +import pytest +import torch + +from torchcfm.models.models import MLP +from torchcfm.utils import eight_normal_sample, sample_8gaussians, sample_moons, torch_wrapper + + +@pytest.mark.parametrize("n_samples", [1, 10, 100, 1000]) +@pytest.mark.parametrize("dim", [2, 3, 4]) +def test_eight_normal_sample_shape_and_finite(n_samples, dim): + # Arrange & Act + data = eight_normal_sample(n_samples, dim) + # Assert + assert isinstance(data, torch.Tensor) + assert data.shape == (n_samples, dim) + assert torch.isfinite(data).all() + + +@pytest.mark.parametrize("n_samples", [1, 10, 100, 1000]) +def test_eight_normal_sample_default_params(n_samples): + # Arrange & Act — default scale=1, var=1 + data = eight_normal_sample(n_samples, 2) + # Assert + assert isinstance(data, torch.Tensor) + assert data.shape == (n_samples, 2) + assert torch.isfinite(data).all() + + +def test_eight_normal_sample_edge_case_single_sample(): + # Arrange & Act — smallest batch size with non-default scale/var + data = eight_normal_sample(1, 2, scale=5, var=0.5) + # Assert + assert data.shape == (1, 2) + assert torch.isfinite(data).all() + + +@pytest.mark.parametrize("n_samples", [1, 10, 100, 1000]) +def test_sample_moons_shape_and_finite(n_samples): + # Arrange & Act + data = sample_moons(n_samples) + # Assert — sample_moons returns a numpy array of shape (n_samples, 2) + assert np.asarray(data).shape == (n_samples, 2) + assert torch.isfinite(torch.as_tensor(data)).all() + + +def test_sample_moons_edge_case_single_sample(): + # Arrange & Act — smallest batch size + data = sample_moons(1) + # Assert + assert np.asarray(data).shape == (1, 2) + assert torch.isfinite(torch.as_tensor(data)).all() + + +@pytest.mark.parametrize("n_samples", [1, 10, 100, 1000]) +def test_sample_8gaussians_shape_and_finite(n_samples): + # Arrange & Act + data = sample_8gaussians(n_samples) + # Assert + assert isinstance(data, torch.Tensor) + assert data.shape == (n_samples, 2) + assert torch.isfinite(data).all() + + +def test_sample_8gaussians_edge_case_single_sample(): + # Arrange & Act — smallest batch size + data = sample_8gaussians(1) + # Assert + assert data.shape == (1, 2) + assert torch.isfinite(data).all() + + +def test_torch_wrapper_returns_callable_with_correct_shape(): + # Arrange — MLP expects an input of dim (dim + 1) when time_varying + model = MLP(dim=2, w=16, time_varying=True) + wrapper = torch_wrapper(model) + t = torch.tensor(0.5) + x = torch.randn(10, 2) + # Act + out = wrapper(t, x) + # Assert + assert isinstance(out, torch.Tensor) + assert out.shape == (10, 2) + assert torch.isfinite(out).all() + + +def test_torch_wrapper_is_torch_module(): + # Arrange + model = MLP(dim=2, w=16, time_varying=True) + wrapper = torch_wrapper(model) + # Assert — torch_wrapper subclasses torch.nn.Module for torchdyn compatibility + assert isinstance(wrapper, torch.nn.Module) + assert wrapper.model is model + + +@pytest.mark.parametrize("batch_size", [1, 5, 32]) +def test_torch_wrapper_parametrized(batch_size): + # Arrange + model = MLP(dim=2, w=16, time_varying=True) + wrapper = torch_wrapper(model) + t = torch.tensor(0.5) + x = torch.randn(batch_size, 2) + # Act + out = wrapper(t, x) + # Assert + assert out.shape == (batch_size, 2) + assert torch.isfinite(out).all() + + +@pytest.mark.skip(reason="plot_trajectories requires a matplotlib display backend") +def test_plot_trajectories(): + """Skipped: plotting requires an interactive matplotlib display.""" + from torchcfm.utils import plot_trajectories + + traj = torch.randn(5, 100, 2) + plot_trajectories(traj)