Skip to content

fix: numerical guard, OT warnings, lazy matplotlib, and GradModel export #182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
atong01 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: numerical guard, OT warnings, lazy matplotlib, and GradModel export #182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions torchcfm/conditional_flow_matching.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,15 @@ def compute_lambda(self, t):
[4] Simulation-free Schrodinger bridges via score and flow matching, Preprint, Tong et al.
"""
sigma_t = self.compute_sigma_t(t)
return 2 * sigma_t / (self.sigma**2 + 1e-8)
if self.sigma == 0:
# Avoid division by zero when sigma=0 (the default).
# When sigma=0 the path is deterministic and lambda is not well-defined;
# return ones so downstream loss weighting is a no-op.
# Use `t` (always a tensor per the docstring) rather than `sigma_t`,
# because the base class's `compute_sigma_t` returns the scalar
# `self.sigma` when sigma=0, and `torch.ones_like` requires a tensor.
return torch.ones_like(t)
return 2 * sigma_t / (self.sigma**2)


class ExactOptimalTransportConditionalFlowMatcher(ConditionalFlowMatcher):
Expand All @@ -226,17 +234,19 @@ class ExactOptimalTransportConditionalFlowMatcher(ConditionalFlowMatcher):
It overrides the sample_location_and_conditional_flow.
"""

def __init__(self, sigma: Union[float, int] = 0.0):
def __init__(self, sigma: Union[float, int] = 0.0, ot_method="exact"):
r"""Initialize the ConditionalFlowMatcher class.

It requires the hyper-parameter $\sigma$.
Parameters
----------
sigma : Union[float, int]
ot_sampler: exact OT method to draw couplings (x0, x1) (see Eq.(17) [1]).
ot_method: OT method to draw couplings (x0, x1) (see Eq.(17) [1]).
"exact" computes the exact OT plan; other solvers (e.g.
"sinkhorn") can be passed through to ``OTPlanSampler``.
"""
super().__init__(sigma)
self.ot_sampler = OTPlanSampler(method="exact")
self.ot_sampler = OTPlanSampler(method=ot_method)

def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False):
r"""
Expand Down
2 changes: 1 addition & 1 deletion torchcfm/models/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
from .models import MLP
from .models import MLP, GradModel
12 changes: 7 additions & 5 deletions torchcfm/optimal_transport.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,10 +86,11 @@ def get_map(self, x0, x1):
M = M / M.max() # should not be normalized when using minibatches
p = self.ot_fn(a, b, M.detach().cpu().numpy())
if not np.all(np.isfinite(p)):
print("ERROR: p is not finite")
print(p)
print("Cost mean, max", M.mean(), M.max())
print(x0, x1)
warnings.warn(
"Non-finite values in OT plan p. "
f"p={p}. Cost mean={M.mean()}, max={M.max()}. "
f"x0={x0}, x1={x1}."
)
if np.abs(p.sum()) < 1e-8:
Comment on lines 87 to 94

@sourcery-ai sourcery-ai Bot Jul 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion (performance): Warning message may be extremely verbose and slow when x0/x1 or p are large arrays.

Embedding full p, x0, and x1 in the warning can create very large strings and hurt performance, particularly in tight loops. Consider logging only shapes and basic stats (e.g., p.shape, x0.shape, x1.shape, dtypes, mins/maxes), or gating full dumps behind a debug flag.

Suggested change
p = self.ot_fn(a, b, M.detach().cpu().numpy())
if not np.all(np.isfinite(p)):
print("ERROR: p is not finite")
print(p)
print("Cost mean, max", M.mean(), M.max())
print(x0, x1)
warnings.warn(
"Non-finite values in OT plan p. "
f"p={p}. Cost mean={M.mean()}, max={M.max()}. "
f"x0={x0}, x1={x1}."
)
if np.abs(p.sum()) < 1e-8:
p = self.ot_fn(a, b, M.detach().cpu().numpy())
if not np.all(np.isfinite(p)):
# Log only summary statistics to avoid building very large warning
# strings when p, x0, or x1 are large arrays/tensors.
p_min = np.nanmin(p)
p_max = np.nanmax(p)
p_mean = np.nanmean(p)
x0_shape = getattr(x0, "shape", None)
x1_shape = getattr(x1, "shape", None)
x0_dtype = getattr(x0, "dtype", None)
x1_dtype = getattr(x1, "dtype", None)
warnings.warn(
"Non-finite values in OT plan p. "
f"p.shape={getattr(p, 'shape', None)}, "
f"p.dtype={getattr(p, 'dtype', None)}, "
f"p.min={p_min}, p.max={p_max}, p.mean={p_mean}. "
f"Cost mean={M.mean()}, max={M.max()}. "
f"x0.shape={x0_shape}, x0.dtype={x0_dtype}; "
f"x1.shape={x1_shape}, x1.dtype={x1_dtype}."
)
if np.abs(p.sum()) < 1e-8:

if self.warn:
warnings.warn("Numerical errors in OT plan, reverting to uniform plan.")
Expand Down Expand Up @@ -279,7 +280,8 @@ def wasserstein(
ret : float
Wasserstein distance
"""
assert power == 1 or power == 2
if power not in (1, 2):
raise ValueError(f"power must be 1 or 2, got {power}")
# ot_fn should take (a, b, M) as arguments where a, b are marginals and
# M is a cost matrix
if method == "exact" or method is None:
Expand Down
10 changes: 7 additions & 3 deletions torchcfm/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import math

import matplotlib.pyplot as plt
import numpy as np
import torch
from torchdyn.datasets import generate_moons
Expand All @@ -9,10 +8,12 @@


def eight_normal_sample(n, dim, scale=1, var=1):
if dim < 2:
raise ValueError(f"dim must be >= 2 (the eight normals are 2D), got {dim}")
m = torch.distributions.multivariate_normal.MultivariateNormal(
torch.zeros(dim), math.sqrt(var) * torch.eye(dim)
)
centers = [
centers_2d = [
(1, 0),
(-1, 0),
(0, 1),
Expand All @@ -22,7 +23,8 @@ def eight_normal_sample(n, dim, scale=1, var=1):
(-1.0 / np.sqrt(2), 1.0 / np.sqrt(2)),
(-1.0 / np.sqrt(2), -1.0 / np.sqrt(2)),
]
centers = torch.tensor(centers) * scale
centers = torch.zeros(8, dim)
centers[:, :2] = torch.tensor(centers_2d) * scale
noise = m.sample((n,))
multi = torch.multinomial(torch.ones(8), n, replacement=True)
data = []
Expand Down Expand Up @@ -54,6 +56,8 @@ def forward(self, t, x, *args, **kwargs):

def plot_trajectories(traj):
"""Plot trajectories of some selected samples."""
import matplotlib.pyplot as plt

n = 2000
plt.figure(figsize=(6, 6))
plt.scatter(traj[0, :n, 0], traj[0, :n, 1], s=10, alpha=0.8, c="black")
Expand Down
Loading