debug barycenter again

Author: rflamary

RELEASES.md

@@ -21,7 +21,7 @@
 + Add `stop_criterion` feature to (un)regularized (f)gw barycenter solvers (PR #578)
 + Add `fixed_structure` and `fixed_features` to entropic fgw barycenter solver (PR #578)
 + Add new BAPG solvers with KL projections for GW and FGW (PR #581)
-+ Add Bures-Wasserstein barycenter in `ot.gaussian` (PR #582)
++ Add Bures-Wasserstein barycenter in `ot.gaussian` and example (PR #582, PR #584)
 
 
 #### Closed issues

ot/gaussian.py

@@ -399,15 +399,15 @@ def bures_wasserstein_barycenter(m, C, weights=None, num_iter=1000, eps=1e-7, lo
     """
     nx = get_backend(*C, *m,)
 
+    if weights is None:
+        weights = nx.ones(C.shape[0], type_as=C[0]) / C.shape[0]
+
     # Compute the mean barycenter
-    mb = nx.dot(weights, m)
+    mb = nx.sum(m * weights[:, None], axis=0)
 
     # Init the covariance barycenter
     Cb = nx.mean(C * weights[:, None, None], axis=0)
 
-    if weights is None:
-        weights = nx.ones(len(C), type_as=C[0]) / len(C)
-
     for it in range(num_iter):
         # fixed point update
         Cb12 = nx.sqrtm(Cb)