[WIP] Generalize convolutional Wasserstein barycenters to N-D grids

README.md

@@ -40,7 +40,7 @@ POT provides the following generic OT solvers:
   Algorithm](https://pythonot.github.io/auto_examples/plot_OT_1D.html) \[2] ,
   stabilized version \[9] \[10] \[34], lazy CPU/GPU solver from geomloss \[60] \[61], greedy Sinkhorn \[22] and Screening
   Sinkhorn \[26].
-* Bregman projections for [Wasserstein barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_barycenter_lp_vs_entropic.html) \[3], [convolutional barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_convolutional_barycenter.html) \[21]  and unmixing \[4].
+* Bregman projections for [Wasserstein barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_barycenter_lp_vs_entropic.html) \[3], [convolutional barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_convolutional_barycenter.html) \[21] (`ot.bregman.convolutional_grid_barycenter` generalizes it to grids of any dimension, e.g. 1D signals or 3D volumes) and unmixing \[4].
 * Sinkhorn divergence \[23] and entropic regularization OT from empirical data.
 * Debiased Sinkhorn barycenters [Sinkhorn divergence barycenter](https://pythonot.github.io/auto_examples/barycenters/plot_debiased_barycenter.html) \[37]
 * Smooth optimal transport solvers (dual and semi-dual) for KL and squared L2 regularizations \[17].

RELEASES.md

@@ -4,6 +4,7 @@
 
 #### New features
 
+- Generalize the separable-kernel convolutional Wasserstein barycenter to regular grids of any dimension (1D signals, 3D volumes, ...) via `ot.bregman.convolutional_grid_barycenter` and `ot.bregman.convolutional_grid_barycenter_debiased`; `ot.bregman.convolutional_barycenter2d` and `ot.bregman.convolutional_barycenter2d_debiased` are now thin `A.ndim == 3` wrappers around them, with unchanged signatures, defaults and docstrings. As before, `log["U"]`/`log["V"]` are only populated by the `method="sinkhorn"` solvers, not `"sinkhorn_log"`, and this asymmetry is now consistent across all four grid functions
 - 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
@@ -15,6 +16,8 @@
 
 #### Closed issues
 
+- Vectorize `ot.bregman.convolutional_barycenter2d`'s log-domain solvers (`method="sinkhorn_log"`) across histograms instead of a per-histogram Python loop with in-place writes, which unblocks Jax and Tensorflow for `sinkhorn_log` (previously a `NotImplementedError`) and speeds up the exp-domain solvers too via a separable, batched kernel application (~7-8x for `sinkhorn` and ~55-65x for `sinkhorn_log` measured on the 4-image, `reg=0.004` example)
+- Fix `Backend.logsumexp` API conformance: `CupyBackend.logsumexp` and `TensorflowBackend.logsumexp` did not accept `keepdims`, unlike the NumPy/Jax/Torch backends and the base class signature; `CupyBackend.logsumexp` now delegates to `cupyx.scipy.special.logsumexp` instead of a hand-rolled port (Issue #867)
 - Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860)
 - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859)
 - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)

docs/source/user_guide.rst

@@ -575,7 +575,11 @@ accelerate the estimation of Wasserstein barycenter when the support has a
 separable structure [21]_. In the case of 2D images for instance one can replace
 the matrix vector production in the Bregman projections by convolution
 operators. We provide an implementation of this algorithm in function
-:any:`ot.bregman.convolutional_barycenter2d`.
+:any:`ot.bregman.convolutional_barycenter2d`, and its debiased variant [37]_ in
+:any:`ot.bregman.convolutional_barycenter2d_debiased`. The same separable-kernel
+algorithm generalizes to a regular grid of any dimension (1D signals, 3D
+volumes, ...) in :any:`ot.bregman.convolutional_grid_barycenter` and
+:any:`ot.bregman.convolutional_grid_barycenter_debiased`.
 
 
 

ot/backend.py

@@ -174,6 +174,7 @@ def norm_1d_jax_jvp(primals, tangents):
     try:
         import cupy as cp
         import cupyx
+        import cupyx.scipy.special
 
         cp_type = cp.ndarray
     except ImportError:
@@ -2946,22 +2947,8 @@ def diag(self, a, k=0):
     def unique(self, a, return_inverse=False):
         return cp.unique(a, return_inverse=return_inverse)
 
-    def logsumexp(self, a, axis=None):
-        # Taken from
-        # https://github.com/scipy/scipy/blob/v1.7.1/scipy/special/_logsumexp.py#L7-L127
-        a_max = cp.amax(a, axis=axis, keepdims=True)
-
-        if a_max.ndim > 0:
-            a_max[~cp.isfinite(a_max)] = 0
-        elif not cp.isfinite(a_max):
-            a_max = 0
-
-        tmp = cp.exp(a - a_max)
-        s = cp.sum(tmp, axis=axis)
-        out = cp.log(s)
-        a_max = cp.squeeze(a_max, axis=axis)
-        out += a_max
-        return out
+    def logsumexp(self, a, axis=None, keepdims=False):
+        return cupyx.scipy.special.logsumexp(a, axis=axis, keepdims=keepdims)
 
     def stack(self, arrays, axis=0):
         return cp.stack(arrays, axis)
@@ -3410,8 +3397,8 @@ def unique(self, a, return_inverse=False):
         else:
             return y_prime
 
-    def logsumexp(self, a, axis=None):
-        return tf.math.reduce_logsumexp(a, axis=axis)
+    def logsumexp(self, a, axis=None, keepdims=False):
+        return tf.math.reduce_logsumexp(a, axis=axis, keepdims=keepdims)
 
     def stack(self, arrays, axis=0):
         return tnp.stack(arrays, axis)

ot/bregman/__init__.py

@@ -32,6 +32,8 @@
 from ._convolutional import (
     convolutional_barycenter2d,
     convolutional_barycenter2d_debiased,
+    convolutional_grid_barycenter,
+    convolutional_grid_barycenter_debiased,
 )
 
 from ._empirical import (
@@ -69,6 +71,8 @@
     "jcpot_barycenter",
     "convolutional_barycenter2d",
     "convolutional_barycenter2d_debiased",
+    "convolutional_grid_barycenter",
+    "convolutional_grid_barycenter_debiased",
     "empirical_sinkhorn",
     "empirical_sinkhorn2",
     "empirical_sinkhorn2_geomloss",