[MRG] Bures-Wasserstein Gradient Descent for Bures-Wasserstein Barycenters (#680)

* bw barycenter with batched sqrtm

* BWGD for barycenters

* sbwgd for barycenters

* Test fixed_point vs gradient_descent

* fix test bwgd

* nx exp_bures

* update doc

* fix merge

* doc exp bw

* First tests stochastic + exp

* exp_bures with einsum

* type Id test

* up test stochastic

* test weights

* Add BW distance with batchs

* step size SGD BW Barycenter

* batchable BW distance

* RELEASES.md

* precommit

* Add ot.gaussian.bures

* Add arg backend

* up stop criteria sgd Gaussian barycenter

* Fix release

* fix doc

* change API bw

* up test bures_wasserstein_distance

* up test bures_wasserstein_distance

* up test bures_wasserstein_distance

---------

Co-authored-by: Cédric Vincent-Cuaz <[email protected]>
Co-authored-by: Rémi Flamary <[email protected]>

Author: 3 people

README.md

@@ -390,3 +390,7 @@ Artificial Intelligence.
 [72] Thibault Séjourné, François-Xavier Vialard, and Gabriel Peyré (2021). [The Unbalanced Gromov Wasserstein Distance: Conic Formulation and Relaxation](https://proceedings.neurips.cc/paper/2021/file/4990974d150d0de5e6e15a1454fe6b0f-Paper.pdf). Neural Information Processing Systems (NeurIPS).
 
 [73] Séjourné, T., Vialard, F. X., & Peyré, G. (2022). [Faster Unbalanced Optimal Transport: Translation Invariant Sinkhorn and 1-D Frank-Wolfe](https://proceedings.mlr.press/v151/sejourne22a.html). In International Conference on Artificial Intelligence and Statistics (pp. 4995-5021). PMLR.
+
+[74] Chewi, S., Maunu, T., Rigollet, P., & Stromme, A. J. (2020). [Gradient descent algorithms for Bures-Wasserstein barycenters](https://proceedings.mlr.press/v125/chewi20a.html). In Conference on Learning Theory (pp. 1276-1304). PMLR.
+
+[75] Altschuler, J., Chewi, S., Gerber, P. R., & Stromme, A. (2021). [Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent](https://papers.neurips.cc/paper_files/paper/2021/hash/b9acb4ae6121c941324b2b1d3fac5c30-Abstract.html). Advances in Neural Information Processing Systems, 34, 22132-22145.

RELEASES.md

@@ -9,6 +9,10 @@
 - Reorganize sub-module `ot/lp/__init__.py` into separate files (PR #714)
 - Implement projected gradient descent solvers for entropic partial FGW (PR #702)
 - Fix documentation in the module `ot.gaussian` (PR #718)
+- Refactored `ot.bregman._convolutional` to improve readability (PR #709)
+- Added `ot.gaussian.bures_barycenter_gradient_descent` (PR #680)
+- Added `ot.gaussian.bures_wasserstein_distance` (PR #680)
+- `ot.gaussian.bures_wasserstein_distance` can be batched (PR #680)
 
 #### Closed issues
 - Fixed `ot.mapping` solvers which depended on deprecated `cvxpy` `ECOS` solver (PR #692, Issue #668)
@@ -44,7 +48,6 @@ This release also contains few bug fixes, concerning the support of any metric i
 - Notes before depreciating partial Gromov-Wasserstein function in `ot.partial` moved to ot.gromov  (PR #663)
 - Create `ot.gromov._partial` add new features `loss_fun = "kl_loss"` and `symmetry=False` to all solvers while increasing speed + updating adequatly `ot.solvers` (PR #663)
 - Added `ot.unbalanced.sinkhorn_unbalanced_translation_invariant` (PR #676)
-- Refactored `ot.bregman._convolutional` to improve readability (PR #709)
 
 #### Closed issues
 - Fixed `ot.gaussian` ignoring weights when computing means (PR #649, Issue #648)

ot/backend.py

@@ -1363,7 +1363,7 @@ def solve(self, a, b):
         return np.linalg.solve(a, b)
 
     def trace(self, a):
-        return np.trace(a)
+        return np.einsum("...ii", a)
 
     def inv(self, a):
         return scipy.linalg.inv(a)
@@ -1776,7 +1776,7 @@ def solve(self, a, b):
         return jnp.linalg.solve(a, b)
 
     def trace(self, a):
-        return jnp.trace(a)
+        return jnp.diagonal(a, axis1=-2, axis2=-1).sum(-1)
 
     def inv(self, a):
         return jnp.linalg.inv(a)
@@ -2309,7 +2309,7 @@ def solve(self, a, b):
         return torch.linalg.solve(a, b)
 
     def trace(self, a):
-        return torch.trace(a)
+        return torch.diagonal(a, dim1=-2, dim2=-1).sum(-1)
 
     def inv(self, a):
         return torch.linalg.inv(a)
@@ -2723,7 +2723,7 @@ def solve(self, a, b):
         return cp.linalg.solve(a, b)
 
     def trace(self, a):
-        return cp.trace(a)
+        return cp.trace(a, axis1=-2, axis2=-1)
 
     def inv(self, a):
         return cp.linalg.inv(a)