Merged Main back into geomloss modifications that allow slicing along columns.

Merge remote-tracking branch 'origin/master' into geomloss_update

Author: nathanneike

.github/workflows/build_doc.yml

@@ -15,11 +15,16 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     # Standard drop-in approach that should work for most people.
-
+    - name: Free Disk Space (Ubuntu)
+      uses: insightsengineering/disk-space-reclaimer@v1
+      with:
+        android: true
+        dotnet: true
     - name: Set up Python 3.10
       uses: actions/setup-python@v5
       with:
         python-version: "3.10"
+        cache: 'pip'
 
     - name: Get Python running
       run: |

.github/workflows/build_tests.yml

@@ -64,11 +64,17 @@ jobs:
         python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
+    - name: Free Disk Space (Ubuntu)
+      uses: insightsengineering/disk-space-reclaimer@v1
+      with:
+        android: true
+        dotnet: true
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+        cache: 'pip'
     - name: Install POT
       run: |
         pip install -e .
@@ -93,6 +99,7 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: "3.13"
+        cache: 'pip'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip setuptools
@@ -121,6 +128,7 @@ jobs:
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
+         cache: 'pip'
      - name: Install POT
        run: |
          pip install -e .
@@ -148,6 +156,7 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
+        cache: 'pip'
     - name: RC.exe
       run: |
         function Invoke-VSDevEnvironment {

RELEASES.md

@@ -7,11 +7,13 @@ This new release adds support for sparse cost matrices in the exact EMD solver.
 #### New features
 - Add support for sparse cost matrices in exact EMD solver `ot.emd` and `ot.emd2` (PR #778)
 - Migrate backend from deprecated `scipy.sparse.coo_matrix` to modern `scipy.sparse.coo_array` API (PR #TBD)
+- Geomloss function now handles both scalar and slice indices for i and j. Using backend agnostic reshaping. Allows to do plan[i,:] and plan[:,j]
 
 #### Closed issues
 - Add support for sparse cost matrices in EMD solver (PR #778, Issue #397)
 - Fix deprecated JAX function in `ot.backend.JaxBackend` (PR #771, Issue #770)
 - Add test for build from source (PR #772, Issue #764)
+- Fix device for batch Ot solver in `ot.batch` (PR #784, Issue #783)
 
 ## 0.9.6.post1
 

ot/batch/_linear.py

@@ -310,9 +310,9 @@ def solve_batch(
     B, n, m = M.shape
 
     if a is None:
-        a = nx.ones((B, n)) / n
+        a = nx.ones((B, n), type_as=M) / n
     if b is None:
-        b = nx.ones((B, m)) / m
+        b = nx.ones((B, m), type_as=M) / m
 
     if solver == "log_sinkhorn":
         K = -M / reg

ot/bregman/_geomloss.py

@@ -54,13 +54,32 @@ def get_sinkhorn_geomloss_lazytensor(
     shape = (X_a.shape[0], X_b.shape[0])
 
     def func(i, j, X_a, X_b, f, g, a, b, metric, blur):
+        X_a_i = X_a[i]
+        X_b_j = X_b[j]
+
+        if X_a_i.ndim == 1:
+            X_a_i = X_a_i[None, :]
+        if X_b_j.ndim == 1:
+            X_b_j = X_b_j[None, :]
+
         if metric == "sqeuclidean":
-            C = dist(X_a[i], X_b[j], metric=metric) / 2
+            C = dist(X_a_i, X_b_j, metric=metric) / 2
         else:
-            C = dist(X_a[i], X_b[j], metric=metric)
-        return nx.exp((f[i, None] + g[None, j] - C) / (blur**2)) * (
-            a[i, None] * b[None, j]
-        )
+            C = dist(X_a_i, X_b_j, metric=metric)
+
+        # Robust broadcasting using nx backend (handles both numpy and torch)
+        # For scalars, slice to keep 1D; for arrays, index directly
+        f_i = f[i : i + 1] if isinstance(i, int) else f[i]
+        g_j = g[j : j + 1] if isinstance(j, int) else g[j]
+        a_i = a[i : i + 1] if isinstance(i, int) else a[i]
+        b_j = b[j : j + 1] if isinstance(j, int) else b[j]
+
+        f_i = nx.reshape(f_i, (-1, 1))
+        g_j = nx.reshape(g_j, (1, -1))
+        a_i = nx.reshape(a_i, (-1, 1))
+        b_j = nx.reshape(b_j, (1, -1))
+
+        return nx.squeeze(nx.exp((f_i + g_j - C) / (blur**2)) * a_i * b_j)
 
     T = LazyTensor(
         shape, func, X_a=X_a, X_b=X_b, f=f, g=g, a=a, b=b, metric=metric, blur=blur