Skip to content

[MRG] Use ot.utils.check_marginal for default marginals across solvers - #856

Merged
rflamary merged 2 commits into
PythonOT:masterfrom
tomMoral:use-check-marginal-in-solvers
Sep 15, 2026
Merged

rflamary merged 2 commits into
PythonOT:masterfrom
tomMoral:use-check-marginal-in-solvers

Conversation

@tomMoral

Copy link
Copy Markdown
Contributor

Types of changes

  • Refactoring / consistency (non-breaking, no new user-facing feature)

Motivation and context / Related issue

Follow-up to #851, which introduced ot.utils.check_marginal (fill a uniform
marginal 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(...) / n idiom is repeated ~25×
across the rest of the library. Routing those through check_marginal makes
default-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_marginal at 25 backend-preserving
1D sites:

  • ot/gromov/: _gw.py, _partial.py, _semirelaxed.py, _lowrank.py
  • ot/lowrank.py, ot/factored.py, ot/stochastic.py
  • ot/bregman/_geomloss.py, ot/bregman/_barycenter.py
  • ot/lp/_barycenter_solvers.py, ot/utils.py (split_sample_ratio)

check_marginal stays in ot.utils (not exported at top-level).

Deliberately not changed:

  • ot/gaussian.py — weights are (n, 1) column vectors; unif/check_marginal
    normalize over the last axis, which would be wrong for that shape.
  • ot/gromov/_dictionary.py, ot/weak.py, ot/bregman/_empirical.py — numpy-by-design
    paths (they convert everything to numpy first), so a CPU default is intentional.

One small hardening to check_marginal: when validating a provided marginal it
now infers the backend with get_backend(type_as, a), so a cross-backend
marginal 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 + validation
change, 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=None path and for consistent-backend inputs (all existing tests) the
behavior is unchanged.

PR checklist

🤖 Generated with Claude Code

@codecov

codecov Bot commented Sep 14, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.85%. Comparing base (caa11a4) to head (60c1704).

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:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@rflamary rflamary left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks very much @tomMoral

@rflamary
rflamary merged commit 3073aa0 into PythonOT:master Sep 15, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants