Update tests to cover unknown loss_fun

Author: kachayev

test/test_gromov.py

@@ -311,7 +311,12 @@ def test_gw_helper_validation(loss_fun):
 
 @pytest.skip_backend("jax", reason="test very slow with jax backend")
 @pytest.skip_backend("tf", reason="test very slow with tf backend")
-def test_entropic_gromov(nx):
+@pytest.mark.parametrize('loss_fun', [
+    'square_loss',
+    'kl_loss',
+    pytest.param('unknown_loss', marks=pytest.mark.xfail(raises=ValueError)),
+])
+def test_entropic_gromov(nx, loss_fun):
     n_samples = 10  # nb samples
 
     mu_s = np.array([0, 0])
@@ -333,10 +338,10 @@ def test_entropic_gromov(nx):
     C1b, C2b, pb, qb, G0b = nx.from_numpy(C1, C2, p, q, G0)
 
     G, log = ot.gromov.entropic_gromov_wasserstein(
-        C1, C2, None, q, 'square_loss', symmetric=None, G0=G0,
+        C1, C2, None, q, loss_fun, symmetric=None, G0=G0,
         epsilon=1e-2, max_iter=10, verbose=True, log=True)
     Gb = nx.to_numpy(ot.gromov.entropic_gromov_wasserstein(
-        C1b, C2b, pb, None, 'square_loss', symmetric=True, G0=None,
+        C1b, C2b, pb, None, loss_fun, symmetric=True, G0=None,
         epsilon=1e-2, max_iter=10, verbose=True, log=False
     ))
 
@@ -347,11 +352,40 @@ def test_entropic_gromov(nx):
     np.testing.assert_allclose(
         q, Gb.sum(0), atol=1e-04)  # cf convergence gromov
 
+
+@pytest.skip_backend("jax", reason="test very slow with jax backend")
+@pytest.skip_backend("tf", reason="test very slow with tf backend")
+@pytest.mark.parametrize('loss_fun', [
+    'square_loss',
+    'kl_loss',
+    pytest.param('unknown_loss', marks=pytest.mark.xfail(raises=ValueError)),
+])
+def test_entropic_gromov2(nx, loss_fun):
+    n_samples = 10  # nb samples
+
+    mu_s = np.array([0, 0])
+    cov_s = np.array([[1, 0], [0, 1]])
+
+    xs = ot.datasets.make_2D_samples_gauss(n_samples, mu_s, cov_s, random_state=42)
+
+    xt = xs[::-1].copy()
+
+    p = ot.unif(n_samples)
+    q = ot.unif(n_samples)
+    G0 = p[:, None] * q[None, :]
+    C1 = ot.dist(xs, xs)
+    C2 = ot.dist(xt, xt)
+
+    C1 /= C1.max()
+    C2 /= C2.max()
+
+    C1b, C2b, pb, qb, G0b = nx.from_numpy(C1, C2, p, q, G0)
+
     gw, log = ot.gromov.entropic_gromov_wasserstein2(
-        C1, C2, p, None, 'kl_loss', symmetric=True, G0=None,
+        C1, C2, p, None, loss_fun, symmetric=True, G0=None,
         max_iter=10, epsilon=1e-2, log=True)
     gwb, logb = ot.gromov.entropic_gromov_wasserstein2(
-        C1b, C2b, None, qb, 'kl_loss', symmetric=None, G0=G0b,
+        C1b, C2b, None, qb, loss_fun, symmetric=None, G0=G0b,
         max_iter=10, epsilon=1e-2, log=True)
     gwb = nx.to_numpy(gwb)