test: expand test suite with fixtures, utils, integration, and model forward-pass tests

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)

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 <kilian.fatras@mila.quebec>
 

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

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])

tests/test_optimal_transport.py

@@ -1,4 +1,4 @@
-"""Tests for Conditional Flow Matcher classers."""
+"""Tests for Conditional Flow Matcher classes."""
 
 # Author: Kilian Fatras <kilian.fatras@mila.quebec>
 

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)