Bures.barycenter returns covariance matrix that has 0 values on the diagonal

[conorhassan]

Hi there,

I love the OTT package - it's fantastic!

I am trying to use Bures.barycenter to find the weighted average (Gaussian distribution) between multiple Gaussians, but it returns a covariance matrix with 0's on the diagonal. Here is a minimal example:

from ott.geometry.costs import Bures, mean_and_cov_to_x, x_to_means_and_covs
import jax.numpy as jnp

# first Gaussian 
mu1 = jnp.array([-0.8909, -0.3568, 0.2758, 0.0352, -0.1457])
Sigma1 = jnp.array([0.3206, 0.8825, 0.1113, 0.9454, 0.0052]) * jnp.eye(5)
# second Gaussian 
mu2 = jnp.array([-0.8862, -0.3652, 0.2751, 0.0349, -0.1486])
Sigma2 = jnp.array([0.3075, 0.8545, 0.1110, 0.9206, 0.0054]) * jnp.eye(5)

# initializing Bures instance 
weights = jnp.array([300./537., 237./537.])
bures = Bures(5)

# stacking parameter values
xs = jnp.vstack(
    (mean_and_cov_to_x(mu1, Sigma1, 5), 
    mean_and_cov_to_x(mu2, Sigma2, 5))
)

# print output
output = bures.barycenter(weights, xs)
mu, Sigma = x_to_means_and_covs(output, 5)
print(Sigma)

I am trying to use this operation as part of an iterative procedure and thus need to do operations such as Cholesky decompositions on this matrix, etc.

I dug into the covariance_fixpoint_iter function and set the rtol to 1e-6 as opposed to the default value of 1e-2, but this didn't change much.

Any thoughts would be appreciated!

[marcocuturi]

upgraded into issue #199 , thanks!