`AttributeError: 'tuple' object has no attribute 'divergence'`

Report

with ott-jax==0.4.9 the return type of ott.tools.sinkhorn_divergence.ott.tools.sinkhorn_divergence will no longer be a ott.tools.sinkhorn_divergence.SinkhornDivergenceOutput but a tuple of (jaxlib.xla_extension.ArrayImpl, ott.tools.sinkhorn_divergence.SinkhornDivergenceOutput). Thus, when calling .divergence on the new returned type we get the above error. This is in particularly related to this line in the cell_flowè_perturbation/src/cfp/metrics/_metrics.py file

To check it out, here's a minimally reproducible example.

./minimal_example.py

import warnings
from ott.tools.sinkhorn_divergence import sinkhorn_divergence
from ott.geometry import costs, pointcloud
import numpy as np

warnings.filterwarnings("ignore")

x=np.random.normal(size=(5, 20))
y=np.random.normal(size=(5, 20))

cost_fn = costs.SqEuclidean()
geom = pointcloud.PointCloud

div = sinkhorn_divergence(geom,x=x,y=y,cost_fn=cost_fn,epsilon=1.0,scale_cost=1.0,)
try:
    print(div.divergence)
except AttributeError as error:
    print(error)
finally:
    print(type(div))
    if isinstance(div, tuple):
        print([type(e) for e in div])

Then on a shell, we see that everything is good with ott-jax==0.4.8

(base) $ conda create -y -q -n ott-trial-0-4-8 python=3.12
(base) $ conda activate ott-trial-0-4-8
(ott-trial-0-4-8) $ pip install -q ott-jax==0.4.8
(ott-trial-0-4-8) $ python ./minimal_example_sinkhorn_divergence.py
>>> 24.664762
>>> <class 'ott.tools.sinkhorn_divergence.SinkhornDivergenceOutput'>
(ott-trial-0-4-8) $ conda deactivate

However, things break when switching to the more recent version

(base) $ conda create -y -q -n ott-trial-0-4-9 python=3.12
(base) $ conda activate ott-trial-0-4-9
(ott-trial-0-4-8) $ pip install -q ott-jax==0.4.9
(ott-trial-0-4-8) $ python ./minimal_example_sinkhorn_divergence.py
>>> 'tuple' object has no attribute 'divergence'
>>> <class 'tuple'>
>>> [<class 'jaxlib.xla_extension.ArrayImpl'>, <class 'ott.tools.sinkhorn_divergence.SinkhornDivergenceOutput'>

To fix this is sufficient to add the following two lines right after having computed the divergence

# ott-jax 0.4.9
if isinstance(div, tuple):
    div = div[1]