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

RELEASES.md

@@ -4,6 +4,7 @@
 
 #### New features
 
+- Use `ot.utils.check_marginal` (and shape-tuple support in `ot.utils.unif`) to fill and validate default marginals consistently across solvers (Gromov, low-rank, stochastic, barycenter, factored) (PR #856)
 - Add stereographic spherical sliced Wasserstein distance in `ot.sliced.stereographic_sliced_wasserstein_sphere`, with its rotationally invariant extension (PR #836)
 - Add Quasi-Monte Carlo sliced Wasserstein sampling (QSW/RQSW) via generalized
   spiral points, selectable with `sampling_slices` in `sliced_wasserstein_distance`,

ot/bregman/_barycenter.py

@@ -13,7 +13,7 @@
 import warnings
 import numpy as np
 
-from ..utils import dist, list_to_array, unif
+from ..utils import dist, list_to_array, unif, check_marginal
 from ..backend import get_backend
 
 from ._utils import geometricBar, geometricMean, projR, projC
@@ -352,8 +352,7 @@ def free_support_sinkhorn_barycenter(
     N = len(measures_locations)
     k = X_init.shape[0]
     d = X_init.shape[1]
-    if b is None:
-        b = nx.ones((k,), type_as=X_init) / k
+    b = check_marginal(b, k, type_as=X_init)
     if weights is None:
         weights = nx.ones((N,), type_as=X_init) / N
 

ot/bregman/_geomloss.py

@@ -14,7 +14,7 @@
     from geomloss import SamplesLoss
     import torch
     from torch.autograd import grad
-    from ..utils import get_backend, LazyTensor, dist
+    from ..utils import get_backend, LazyTensor, dist, check_marginal
 
     if geomloss.__version__ < "0.3.1":
         old_geomloss = True
@@ -199,10 +199,8 @@ def empirical_sinkhorn2_geomloss(
     if nx.__name__ not in ["torch", "numpy"]:
         raise ValueError("geomloss only support torch or numpy backend")
 
-    if a is None:
-        a = nx.ones(X_s.shape[0], type_as=X_s) / X_s.shape[0]
-    if b is None:
-        b = nx.ones(X_t.shape[0], type_as=X_t) / X_t.shape[0]
+    a = check_marginal(a, X_s.shape[0], type_as=X_s)
+    b = check_marginal(b, X_t.shape[0], type_as=X_t)
 
     if nx.__name__ == "numpy":
         X_s_torch = torch.tensor(X_s)

ot/factored.py

@@ -7,7 +7,7 @@
 # License: MIT License
 
 from .backend import get_backend
-from .utils import dist, get_lowrank_lazytensor
+from .utils import dist, get_lowrank_lazytensor, check_marginal
 from .lp import emd
 from .bregman import sinkhorn
 
@@ -104,10 +104,8 @@ def factored_optimal_transport(
     n_b = Xb.shape[0]
     d = Xa.shape[1]
 
-    if a is None:
-        a = nx.ones((n_a), type_as=Xa) / n_a
-    if b is None:
-        b = nx.ones((n_b), type_as=Xb) / n_b
+    a = check_marginal(a, n_a, type_as=Xa)
+    b = check_marginal(b, n_b, type_as=Xb)
 
     if X0 is None:
         X = nx.randn(r, d, type_as=Xa)

ot/gromov/_gw.py

@@ -17,7 +17,7 @@
 
 from ..utils import dist, UndefinedParameter, list_to_array
 from ..optim import cg, line_search_armijo, solve_1d_linesearch_quad
-from ..utils import check_random_state, unif
+from ..utils import check_marginal, check_random_state, unif
 from ..backend import get_backend, NumpyBackend
 
 from ._utils import init_matrix, gwloss, gwggrad
@@ -373,10 +373,8 @@ def gromov_wasserstein2(
     nx = get_backend(C1, C2)
 
     # init marginals if set as None
-    if p is None:
-        p = unif(C1.shape[0], type_as=C1)
-    if q is None:
-        q = unif(C2.shape[0], type_as=C1)
+    p = check_marginal(p, C1.shape[0], type_as=C1)
+    q = check_marginal(q, C2.shape[0], type_as=C1)
 
     T, log_gw = gromov_wasserstein(
         C1,
@@ -780,10 +778,8 @@ def fused_gromov_wasserstein2(
     nx = get_backend(C1, C2, M)
 
     # init marginals if set as None
-    if p is None:
-        p = unif(C1.shape[0], type_as=M)
-    if q is None:
-        q = unif(C2.shape[0], type_as=M)
+    p = check_marginal(p, C1.shape[0], type_as=M)
+    q = check_marginal(q, C2.shape[0], type_as=M)
 
     T, log_fgw = fused_gromov_wasserstein(
         M,

ot/gromov/_lowrank.py

@@ -7,7 +7,7 @@
 # License: MIT License
 
 import warnings
-from ..utils import unif, get_lowrank_lazytensor
+from ..utils import get_lowrank_lazytensor, check_marginal
 from ..backend import get_backend
 from ..lowrank import compute_lr_sqeuclidean_matrix, _init_lr_sinkhorn, _LR_Dysktra
 
@@ -180,10 +180,8 @@ def lowrank_gromov_wasserstein_samples(
     ns, nt = X_s.shape[0], X_t.shape[0]
 
     # Initialize weights a, b
-    if a is None:
-        a = unif(ns, type_as=X_s)
-    if b is None:
-        b = unif(nt, type_as=X_t)
+    a = check_marginal(a, ns, type_as=X_s)
+    b = check_marginal(b, nt, type_as=X_t)
 
     # Compute rank (see Section 3.1, def 1)
     r = rank

ot/gromov/_partial.py

@@ -9,7 +9,7 @@
 #
 # License: MIT License
 
-from ..utils import list_to_array, unif
+from ..utils import check_marginal, list_to_array, unif
 from ..backend import get_backend, NumpyBackend
 from ..partial import entropic_partial_wasserstein
 from ._utils import _transform_matrix, gwloss, gwggrad
@@ -451,10 +451,8 @@ def partial_gromov_wasserstein2(
     nx = get_backend(C1, C2)
 
     # init marginals if set as None
-    if p is None:
-        p = unif(C1.shape[0], type_as=C1)
-    if q is None:
-        q = unif(C2.shape[0], type_as=C1)
+    p = check_marginal(p, C1.shape[0], type_as=C1)
+    q = check_marginal(q, C2.shape[0], type_as=C1)
 
     T, log_pgw = partial_gromov_wasserstein(
         C1,
@@ -925,10 +923,8 @@ def partial_fused_gromov_wasserstein2(
     nx = get_backend(M, C1, C2)
 
     # init marginals if set as None
-    if p is None:
-        p = unif(C1.shape[0], type_as=C1)
-    if q is None:
-        q = unif(C2.shape[0], type_as=C1)
+    p = check_marginal(p, C1.shape[0], type_as=C1)
+    q = check_marginal(q, C2.shape[0], type_as=C1)
 
     T, log_pfgw = partial_fused_gromov_wasserstein(
         M,
@@ -1207,10 +1203,8 @@ def entropic_partial_gromov_wasserstein(
 
     nx = get_backend(*arr)
 
-    if p is None:
-        p = nx.ones(C1.shape[0], type_as=C1) / C1.shape[0]
-    if q is None:
-        q = nx.ones(C2.shape[0], type_as=C2) / C2.shape[0]
+    p = check_marginal(p, C1.shape[0], type_as=C1)
+    q = check_marginal(q, C2.shape[0], type_as=C2)
 
     if m is None:
         m = min(nx.sum(p), nx.sum(q))
@@ -1565,10 +1559,8 @@ def entropic_partial_fused_gromov_wasserstein(
 
     nx = get_backend(*arr)
 
-    if p is None:
-        p = nx.ones(C1.shape[0], type_as=C1) / C1.shape[0]
-    if q is None:
-        q = nx.ones(C2.shape[0], type_as=C2) / C2.shape[0]
+    p = check_marginal(p, C1.shape[0], type_as=C1)
+    q = check_marginal(q, C2.shape[0], type_as=C2)
 
     if m is None:
         m = min(nx.sum(p), nx.sum(q))

ot/gromov/_semirelaxed.py

@@ -11,7 +11,7 @@
 import numpy as np
 
 
-from ..utils import list_to_array, unif, dist, UndefinedParameter
+from ..utils import list_to_array, unif, dist, UndefinedParameter, check_marginal
 from ..optim import semirelaxed_cg, solve_1d_linesearch_quad
 from ..backend import get_backend
 
@@ -336,8 +336,7 @@ def semirelaxed_gromov_wasserstein2(
     nx = get_backend(C1, C2)
 
     # init marginals if set as None
-    if p is None:
-        p = unif(C1.shape[0], type_as=C1)
+    p = check_marginal(p, C1.shape[0], type_as=C1)
 
     T, log_srgw = semirelaxed_gromov_wasserstein(
         C1,
@@ -698,8 +697,7 @@ def semirelaxed_fused_gromov_wasserstein2(
     nx = get_backend(C1, C2)
 
     # init marginals if set as None
-    if p is None:
-        p = unif(C1.shape[0], type_as=C1)
+    p = check_marginal(p, C1.shape[0], type_as=C1)
 
     T, log_fgw = semirelaxed_fused_gromov_wasserstein(
         M,

ot/lowrank.py

@@ -8,7 +8,7 @@
 # License: MIT License
 
 import warnings
-from .utils import unif, dist, get_lowrank_lazytensor
+from .utils import dist, get_lowrank_lazytensor, check_marginal
 from .backend import get_backend
 from .bregman import sinkhorn
 
@@ -420,10 +420,8 @@ def lowrank_sinkhorn(
     ns, nt = X_s.shape[0], X_t.shape[0]
 
     # Initialize weights a, b
-    if a is None:
-        a = unif(ns, type_as=X_s)
-    if b is None:
-        b = unif(nt, type_as=X_t)
+    a = check_marginal(a, ns, type_as=X_s)
+    b = check_marginal(b, nt, type_as=X_t)
 
     # Compute rank (see Section 3.1, def 1)
     r = rank

ot/lp/_barycenter_solvers.py

@@ -9,7 +9,7 @@
 # License: MIT License
 
 from ..backend import get_backend
-from ..utils import dist
+from ..utils import dist, check_marginal
 from ._network_simplex import emd, emd2
 
 import numpy as np
@@ -240,8 +240,7 @@ def free_support_barycenter(
     N = len(measures_locations)
     k = X_init.shape[0]
     d = X_init.shape[1]
-    if b is None:
-        b = nx.ones((k,), type_as=X_init) / k
+    b = check_marginal(b, k, type_as=X_init)
     if weights is None:
         weights = nx.ones((N,), type_as=X_init) / N
 
@@ -396,8 +395,7 @@ def generalized_free_support_barycenter(
     if Y_init is None:
         Y_init = nx.randn(n_samples_bary, d, type_as=X_list[0])
 
-    if b is None:
-        b = nx.ones(n_samples_bary, type_as=X_list[0]) / n_samples_bary  # not optimized
+    b = check_marginal(b, n_samples_bary, type_as=X_list[0])  # not optimized
 
     out = free_support_barycenter(
         Z_list,

ot/stochastic.py

@@ -10,7 +10,7 @@
 # License: MIT License
 
 import numpy as np
-from .utils import dist, check_random_state
+from .utils import dist, check_random_state, check_marginal
 from .backend import get_backend
 
 ##############################################################################
@@ -653,11 +653,8 @@ def loss_dual_entropic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidea
 
     nx = get_backend(u, v, xs, xt)
 
-    if ws is None:
-        ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0]
-
-    if wt is None:
-        wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0]
+    ws = check_marginal(ws, xs.shape[0], type_as=xs)
+    wt = check_marginal(wt, xt.shape[0], type_as=xt)
 
     if callable(metric):
         M = metric(xs, xt)
@@ -711,11 +708,8 @@ def plan_dual_entropic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclidea
 
     nx = get_backend(u, v, xs, xt)
 
-    if ws is None:
-        ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0]
-
-    if wt is None:
-        wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0]
+    ws = check_marginal(ws, xs.shape[0], type_as=xs)
+    wt = check_marginal(wt, xt.shape[0], type_as=xt)
 
     if callable(metric):
         M = metric(xs, xt)
@@ -769,11 +763,8 @@ def loss_dual_quadratic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclide
 
     nx = get_backend(u, v, xs, xt)
 
-    if ws is None:
-        ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0]
-
-    if wt is None:
-        wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0]
+    ws = check_marginal(ws, xs.shape[0], type_as=xs)
+    wt = check_marginal(wt, xt.shape[0], type_as=xt)
 
     if callable(metric):
         M = metric(xs, xt)
@@ -827,11 +818,8 @@ def plan_dual_quadratic(u, v, xs, xt, reg=1, ws=None, wt=None, metric="sqeuclide
 
     nx = get_backend(u, v, xs, xt)
 
-    if ws is None:
-        ws = nx.ones(xs.shape[0], type_as=xs) / xs.shape[0]
-
-    if wt is None:
-        wt = nx.ones(xt.shape[0], type_as=xt) / xt.shape[0]
+    ws = check_marginal(ws, xs.shape[0], type_as=xs)
+    wt = check_marginal(wt, xt.shape[0], type_as=xt)
 
     if callable(metric):
         M = metric(xs, xt)

ot/utils.py

@@ -298,7 +298,7 @@ def check_marginal(a, shape, type_as=None, nx=None):
         raise ValueError(f"marginal has shape {tuple(a.shape)}, expected {size}")
     if type_as is not None:
         if nx is None:
-            nx = get_backend(type_as)
+            nx = get_backend(type_as, a)
         nx.assert_same_dtype_device(type_as, a)
     return a
 
@@ -2139,8 +2139,7 @@ def split_sample_ratio(
 
     n_a = X_a.shape[0]
 
-    if a is None:
-        a = nx.ones(n_a, type_as=X_a) / n_a
+    a = check_marginal(a, n_a, type_as=X_a)
 
     if random_split:
         if random_state is not None: