[MRG] Use ot.utils.check_marginal for default marginals across solvers - #856
Merged
rflamary merged 2 commits intoSep 15, 2026
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #856 +/- ##
=======================================
Coverage 96.84% 96.85%
=======================================
Files 128 128
Lines 26195 26161 -34
=======================================
- Hits 25369 25337 -32
+ Misses 826 824 -2 🚀 New features to boost your workflow:
|
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
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.
Types of changes
Motivation and context / Related issue
Follow-up to #851, which introduced
ot.utils.check_marginal(fill a uniformmarginal on the inputs' backend/device when
None, otherwise validate shape +dtype/device) and used it in the batch solvers.
The same
if a is None: a = unif(...) / nx.ones(...) / nidiom is repeated ~25×across the rest of the library. Routing those through
check_marginalmakesdefault-marginal handling consistent and adds early shape/dtype/device
validation for user-provided marginals (previously a mismatch surfaced as a deep
error further down the solver).
Description
Replaced the default-marginal idiom with
check_marginalat 25 backend-preserving1D sites:
ot/gromov/:_gw.py,_partial.py,_semirelaxed.py,_lowrank.pyot/lowrank.py,ot/factored.py,ot/stochastic.pyot/bregman/_geomloss.py,ot/bregman/_barycenter.pyot/lp/_barycenter_solvers.py,ot/utils.py(split_sample_ratio)check_marginalstays inot.utils(not exported at top-level).Deliberately not changed:
ot/gaussian.py— weights are(n, 1)column vectors;unif/check_marginalnormalize over the last axis, which would be wrong for that shape.
ot/gromov/_dictionary.py,ot/weak.py,ot/bregman/_empirical.py— numpy-by-designpaths (they convert everything to numpy first), so a CPU default is intentional.
One small hardening to
check_marginal: when validating a provided marginal itnow infers the backend with
get_backend(type_as, a), so a cross-backendmarginal raises the standard "Backends should be the same" error instead of an
AttributeError.Note on behavior
These sites already passed
type_as, so this is a consistency + validationchange, not a device bugfix. The one behavior change is that a provided marginal
whose shape/dtype/device disagrees with the reference now raises early. For the
common
a=Nonepath and for consistent-backend inputs (all existing tests) thebehavior is unchanged.
PR checklist
🤖 Generated with Claude Code