Conversation
- conditional_flow_matching.py: compute_lambda now explicitly handles sigma=0 (returns ones_like) instead of relying on a 1e-8 denominator guard that masked the degenerate case - conditional_flow_matching.py: ExactOptimalTransportConditionalFlowMatcher gains a backward-compatible ot_method='exact' parameter (mirrors SchrodingerBridgeCFM), fixing the docstring/signature mismatch - optimal_transport.py: replace stdout 'print' debug calls in get_map with a proper warnings.warn on non-finite OT plans - optimal_transport.py: replace bare 'assert power == 1 or power == 2' with a ValueError that names the invalid value - utils.py: move 'import matplotlib.pyplot as plt' inside plot_trajectories (lazy import) so the module imports without matplotlib installed - models/__init__.py: export GradModel alongside MLP
Contributor
Reviewer's GuideThis PR applies a set of small, focused fixes across conditional flow matching, optimal transport utilities, plotting helpers, and model exports, improving numerical robustness, API/doc alignment, logging behavior, optional matplotlib usage, and exposing an existing GradModel symbol. File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 1 issue, and left some high level feedback:
- The
get_mapwarning currently interpolates full arrays (p,x0,x1) into the message, which can be extremely verbose; consider logging only shapes and summary statistics (e.g., min/mean/max) to keep diagnostics readable. - In
compute_lambda, you might want to guard against very small non-zerosigmavalues (e.g., via a tolerance) rather than onlysigma == 0, to avoid unstable scaling whensigmais close to zero but not exactly zero.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `get_map` warning currently interpolates full arrays (`p`, `x0`, `x1`) into the message, which can be extremely verbose; consider logging only shapes and summary statistics (e.g., min/mean/max) to keep diagnostics readable.
- In `compute_lambda`, you might want to guard against very small non-zero `sigma` values (e.g., via a tolerance) rather than only `sigma == 0`, to avoid unstable scaling when `sigma` is close to zero but not exactly zero.
## Individual Comments
### Comment 1
<location path="torchcfm/optimal_transport.py" line_range="87-94" />
<code_context>
- 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:
</code_context>
<issue_to_address>
**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.
```suggestion
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:
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
87
to
94
| 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: |
Contributor
There was a problem hiding this comment.
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: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes six small bugs and rough edges in
torchcfm/.Changes
1.
conditional_flow_matching.py—compute_lambdanumerical guardcompute_lambdaused2 * sigma_t / (self.sigma**2 + 1e-8)which masked the degeneratesigma=0case with a tiny denominator. Replaced with an explicitif self.sigma == 0: return torch.ones_like(t)branch. Usest(always a tensor per docstring) rather thansigma_t(which the base class returns as a Python int whensigma=0).2.
conditional_flow_matching.py—ExactOptimalTransportConditionalFlowMatcherdocstring/signature mismatchThe docstring referenced an
ot_samplerparameter that didn't appear in__init__. Added a backward-compatibleot_method='exact'parameter (mirroringSchrodingerBridgeConditionalFlowMatcher) and updated the docstring to describe it.3.
optimal_transport.py— stdout debug prints inget_mapReplaced
print('ERROR: p is not finite')+ three moreprint()calls with a singlewarnings.warn(...)that preserves the same information (p, cost mean/max, x0, x1).4.
optimal_transport.py— bareassertinwassersteinReplaced
assert power == 1 or power == 2withif power not in (1, 2): raise ValueError(f"power must be 1 or 2, got {power}").5.
utils.py— unconditionalmatplotlibimportMoved
import matplotlib.pyplot as pltfrom the top of the module to insideplot_trajectories(the only function that uses it). The module now imports without matplotlib installed — useful for headless test environments.6.
models/__init__.py— missingGradModelexportChanged
from .models import MLPtofrom .models import MLP, GradModel.GradModelwas defined inmodels.pybut not exported.Motivation
These are small correctness and quality bugs found while exploring the codebase. Each fix is localized and backward compatible.
Validation
All changes tested as part of a 4-PR batch (community-docs, add-tests, bugfixes, ci-updates) merged locally and run with pytest:
Bugs discovered and fixed during local testing
During local pytest validation, two bugs surfaced and were fixed in this same branch:
compute_lambdainitially returnedtorch.ones_like(sigma_t)which raisedTypeErrorwhen the base class returnedself.sigmaas a Python int. Fixed by usingtorch.ones_like(t)instead.eight_normal_sample(n, dim, ...)hardcoded 2D centers but accepted adimparameter, raising a shape error whendim != 2. Addedif dim < 2: raise ValueErrorand zero-padded centers fordim > 2.Summary by Sourcery
Fix numerical robustness issues, clean up optimal transport diagnostics, make plotting utilities optional on import, and expose the missing GradModel symbol from the models package.
Bug Fixes:
Enhancements: