[MRG] speedup and update tests and wheels and prepare for release 0.9.6

.github/workflows/build_tests.yml

@@ -47,7 +47,7 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
     - uses: actions/checkout@v4
@@ -78,7 +78,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: "3.12"
+        python-version: "3.13"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip setuptools
@@ -98,7 +98,7 @@ jobs:
      strategy:
        max-parallel: 4
        matrix:
-         os: [macos-latest, macos-13]
+         os: [macos-latest]
          python-version: ["3.12"]
 
      steps:

.github/workflows/build_wheels.yml

@@ -30,7 +30,7 @@ jobs:
 
     - name: Install cibuildwheel
       run: |
-        python -m pip install cibuildwheel==2.23.3
+        python -m pip install cibuildwheel==3.1.4
 
     - name: Build wheels
       env:
@@ -65,7 +65,7 @@ jobs:
 
     - name: Install cibuildwheel
       run: |
-        python -m pip install cibuildwheel==2.16.4
+        python -m pip install cibuildwheel==3.1.4
 
     - name: Set up QEMU
       if: runner.os == 'Linux'

.github/workflows/build_wheels_weekly.yml

@@ -29,7 +29,7 @@ jobs:
 
     - name: Install cibuildwheel
       run: |
-        python -m pip install cibuildwheel==2.23.3
+        python -m pip install cibuildwheel==3.1.4
 
     - name: Set up QEMU
       if: runner.os == 'Linux'

RELEASES.md

@@ -1,6 +1,6 @@
 # Releases
 
-## 0.9.6dev
+## 0.9.6
 
 #### New features
 - Implement CG solvers for partial FGW (PR #687)
@@ -28,6 +28,7 @@
 - Removed release information from quickstart guide (PR #744)
 - Implement batch parallel solvers in ot.batch (PR #745)
 - Update REAMDE with new API and reorganize examples (PR #754)
+- Speedup and update tests and wheels (PR #759)
 
 #### Closed issues
 - Fixed `ot.mapping` solvers which depended on deprecated `cvxpy` `ECOS` solver (PR #692, Issue #668)

examples/backends/plot_ot_batch.py

@@ -54,7 +54,7 @@
         ot.dist(samples_source[i], samples_target[i])
     )  # List of cost matrices n_samples x n_samples
 # Batched approach
-M_batch = ot.batch.dist_batch(
+M_batch = ot.dist_batch(
     samples_source, samples_target
 )  # Array of cost matrices n_problems x n_samples x n_samples
 
@@ -88,7 +88,7 @@
     results_values_list.append(res.value_linear)
 
 # Batched approach
-results_batch = ot.batch.solve_batch(
+results_batch = ot.solve_batch(
     M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy"
 )
 results_values_batch = results_batch.value_linear
@@ -131,8 +131,8 @@ def benchmark_naive(samples_source, samples_target):
 
 def benchmark_batch(samples_source, samples_target):
     start = perf_counter()
-    M_batch = ot.batch.dist_batch(samples_source, samples_target)
-    res_batch = ot.batch.solve_batch(
+    M_batch = ot.dist_batch(samples_source, samples_target)
+    res_batch = ot.solve_batch(
         M=M_batch, reg=reg, max_iter=max_iter, tol=tol, reg_type="entropy"
     )
     end = perf_counter()
@@ -176,8 +176,7 @@ def benchmark_batch(samples_source, samples_target):
 #    If your data is on a GPU, :func:`ot.batch.solve_gromov_batch`
 #    is significantly faster AND provides better objective values.
 
-from ot import solve_gromov
-from ot.batch import solve_gromov_batch
+from ot import solve_gromov, solve_gromov_batch
 
 
 def benchmark_naive_gw(samples_source, samples_target):
@@ -195,8 +194,8 @@ def benchmark_naive_gw(samples_source, samples_target):
 
 def benchmark_batch_gw(samples_source, samples_target):
     start = perf_counter()
-    C1_batch = ot.batch.dist_batch(samples_source, samples_source)
-    C2_batch = ot.batch.dist_batch(samples_target, samples_target)
+    C1_batch = ot.dist_batch(samples_source, samples_source)
+    C2_batch = ot.dist_batch(samples_target, samples_target)
     res_batch = solve_gromov_batch(
         C1_batch, C2_batch, reg=1, max_iter=100, max_iter_inner=50, tol=tol
     )

ot/__init__.py

@@ -72,12 +72,12 @@
 from .solvers import solve, solve_gromov, solve_sample
 from .lowrank import lowrank_sinkhorn
 
-from .batch import solve_batch, solve_gromov_batch
+from .batch import solve_batch, solve_sample_batch, solve_gromov_batch, dist_batch
 
 # utils functions
 from .utils import dist, unif, tic, toc, toq
 
-__version__ = "0.9.6dev0"
+__version__ = "0.9.6"
 
 __all__ = [
     "emd",
@@ -139,4 +139,6 @@
     "lowrank_gromov_wasserstein_samples",
     "solve_batch",
     "solve_gromov_batch",
+    "solve_sample_batch",
+    "dist_batch",
 ]

test/gromov/test_partial.py

@@ -465,8 +465,8 @@ def test_partial_fgw2_gradients():
 @pytest.skip_backend("tf", reason="test very slow with tf backend")
 def test_entropic_partial_gromov_wasserstein(nx):
     rng = np.random.RandomState(42)
-    n_samples = 20  # nb samples
-    n_noise = 10  # nb of samples (noise)
+    n_samples = 10  # nb samples
+    n_noise = 5  # nb of samples (noise)
 
     p = ot.unif(n_samples + n_noise)
     psub = ot.unif(n_samples - 5 + n_noise)
@@ -516,6 +516,7 @@ def test_entropic_partial_gromov_wasserstein(nx):
             log=True,
             symmetric=list_sym[i],
             verbose=True,
+            numItermax=10,
         )
 
         resb, logb = ot.gromov.entropic_partial_gromov_wasserstein(
@@ -530,6 +531,7 @@ def test_entropic_partial_gromov_wasserstein(nx):
             log=True,
             symmetric=False,
             verbose=True,
+            numItermax=10,
         )
 
         resb_ = nx.to_numpy(resb)
@@ -552,6 +554,7 @@ def test_entropic_partial_gromov_wasserstein(nx):
         log=False,
         symmetric=list_sym[i],
         verbose=True,
+        numItermax=10,
     )
 
     resb = ot.gromov.entropic_partial_gromov_wasserstein(
@@ -564,6 +567,7 @@ def test_entropic_partial_gromov_wasserstein(nx):
         log=False,
         symmetric=False,
         verbose=True,
+        numItermax=10,
     )
 
     resb_ = nx.to_numpy(resb)
@@ -573,11 +577,25 @@ def test_entropic_partial_gromov_wasserstein(nx):
     # tests with different number of samples across spaces
     m = 0.5
     res, log = ot.gromov.entropic_partial_gromov_wasserstein(
-        C1, C1sub, p=p, q=psub, reg=1e4, m=m, log=True
+        C1,
+        C1sub,
+        p=p,
+        q=psub,
+        reg=1e4,
+        m=m,
+        log=True,
+        numItermax=10,
     )
 
     resb, logb = ot.gromov.entropic_partial_gromov_wasserstein(
-        C1b, C1subb, p=pb, q=psubb, reg=1e4, m=m, log=True
+        C1b,
+        C1subb,
+        p=pb,
+        q=psubb,
+        reg=1e4,
+        m=m,
+        log=True,
+        numItermax=10,
     )
 
     resb_ = nx.to_numpy(resb)
@@ -589,10 +607,26 @@ def test_entropic_partial_gromov_wasserstein(nx):
     # tests for pGW2
     for loss_fun in ["square_loss", "kl_loss"]:
         w0, log0 = ot.gromov.entropic_partial_gromov_wasserstein2(
-            C1, C2, p=None, q=q, reg=1e4, m=m, loss_fun=loss_fun, log=True
+            C1,
+            C2,
+            p=None,
+            q=q,
+            reg=1e4,
+            m=m,
+            loss_fun=loss_fun,
+            log=True,
+            numItermax=10,
         )
         w0_val = ot.gromov.entropic_partial_gromov_wasserstein2(
-            C1b, C2b, p=pb, q=None, reg=1e4, m=m, loss_fun=loss_fun, log=False
+            C1b,
+            C2b,
+            p=pb,
+            q=None,
+            reg=1e4,
+            m=m,
+            loss_fun=loss_fun,
+            log=False,
+            numItermax=10,
         )
         np.testing.assert_allclose(w0, w0_val, rtol=1e-8)
 
@@ -666,6 +700,7 @@ def test_entropic_partial_fused_gromov_wasserstein(nx):
             log=True,
             symmetric=list_sym[i],
             verbose=True,
+            numItermax=10,
         )
 
         resb, logb = ot.gromov.entropic_partial_fused_gromov_wasserstein(
@@ -681,6 +716,7 @@ def test_entropic_partial_fused_gromov_wasserstein(nx):
             log=True,
             symmetric=False,
             verbose=True,
+            numItermax=10,
         )
 
         resb_ = nx.to_numpy(resb)
@@ -704,6 +740,7 @@ def test_entropic_partial_fused_gromov_wasserstein(nx):
         log=False,
         symmetric=list_sym[i],
         verbose=True,
+        numItermax=10,
     )
 
     resb = ot.gromov.entropic_partial_fused_gromov_wasserstein(
@@ -717,6 +754,7 @@ def test_entropic_partial_fused_gromov_wasserstein(nx):
         log=False,
         symmetric=False,
         verbose=True,
+        numItermax=10,
     )
 
     resb_ = nx.to_numpy(resb)
@@ -726,11 +764,27 @@ def test_entropic_partial_fused_gromov_wasserstein(nx):
     # tests with different number of samples across spaces
     m = 0.5
     res, log = ot.gromov.entropic_partial_fused_gromov_wasserstein(
-        M11sub, C1, C1sub, p=p, q=psub, reg=1e4, m=m, log=True
+        M11sub,
+        C1,
+        C1sub,
+        p=p,
+        q=psub,
+        reg=1e4,
+        m=m,
+        log=True,
+        numItermax=10,
     )
 
     resb, logb = ot.gromov.entropic_partial_fused_gromov_wasserstein(
-        M11subb, C1b, C1subb, p=pb, q=psubb, reg=1e4, m=m, log=True
+        M11subb,
+        C1b,
+        C1subb,
+        p=pb,
+        q=psubb,
+        reg=1e4,
+        m=m,
+        log=True,
+        numItermax=10,
     )
 
     resb_ = nx.to_numpy(resb)
@@ -742,9 +796,27 @@ def test_entropic_partial_fused_gromov_wasserstein(nx):
     # tests for pGW2
     for loss_fun in ["square_loss", "kl_loss"]:
         w0, log0 = ot.gromov.entropic_partial_fused_gromov_wasserstein2(
-            M12, C1, C2, p=None, q=q, reg=1e4, m=m, loss_fun=loss_fun, log=True
+            M12,
+            C1,
+            C2,
+            p=None,
+            q=q,
+            reg=1e4,
+            m=m,
+            loss_fun=loss_fun,
+            log=True,
+            numItermax=10,
         )
         w0_val = ot.gromov.entropic_partial_fused_gromov_wasserstein2(
-            M12b, C1b, C2b, p=pb, q=None, reg=1e4, m=m, loss_fun=loss_fun, log=False
+            M12b,
+            C1b,
+            C2b,
+            p=pb,
+            q=None,
+            reg=1e4,
+            m=m,
+            loss_fun=loss_fun,
+            log=False,
+            numItermax=10,
         )
         np.testing.assert_allclose(w0, w0_val, rtol=1e-8)

test/test_da.py

@@ -912,8 +912,8 @@ def test_emd_laplace_class(nx):
 def test_nearest_brenier_potential(nx):
     X = nx.ones((2, 2))
     for ssnb in [
-        ot.da.NearestBrenierPotential(log=True),
-        ot.da.NearestBrenierPotential(log=False),
+        ot.da.NearestBrenierPotential(log=True, its=5),
+        ot.da.NearestBrenierPotential(log=False, its=5),
     ]:
         ssnb.fit(Xs=X, Xt=X)
         G_lu = ssnb.transform(Xs=X)