From eaefbf1e359a642458b7ef9e324caa4251636a6f Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Thu, 2 Apr 2026 10:33:44 +0800
Subject: [PATCH 01/20] optimal transport

---
 torchmorph/csrc/ot_kernel.cu | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 torchmorph/csrc/ot_kernel.cu

diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
new file mode 100644
index 0000000..ea028fa
--- /dev/null
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -0,0 +1 @@
+// Kernel implementation for Optimal Transport
\ No newline at end of file

From 55603aa7d65b18d78eb0dc055ace7de314dd196e Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Wed, 22 Apr 2026 13:06:08 +0800
Subject: [PATCH 02/20] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E4=B8=80=E7=BB=B4?=
 =?UTF-8?q?=E6=88=90=E6=9C=AC=E7=9F=A9=E9=98=B5=E6=9E=84=E5=BB=BA=E5=92=8C?=
 =?UTF-8?q?Sinkhorn=E8=BF=AD=E4=BB=A3=E7=9A=84=E5=AE=9E=E7=8E=B0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 test/test_optimal_transport.py  | 18 +++++++
 torchmorph/optimal_transport.py | 84 +++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 test/test_optimal_transport.py
 create mode 100644 torchmorph/optimal_transport.py

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
new file mode 100644
index 0000000..99e5ec6
--- /dev/null
+++ b/test/test_optimal_transport.py
@@ -0,0 +1,18 @@
+# test_optimal_transport.py
+import sys
+import importlib.util
+from pathlib import Path
+
+# 项目根目录（test 的父目录）
+project_root = Path(__file__).parent.parent
+
+# 目标模块路径
+module_path = project_root / "torchmorph" / "optimal_transport.py"
+
+# 加载模块
+spec = importlib.util.spec_from_file_location("optimal_transport", module_path)
+ot = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(ot)
+
+a=ot.build_cost_matrix_1d(2,3,'cuda')
+print(a)
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
new file mode 100644
index 0000000..1623de6
--- /dev/null
+++ b/torchmorph/optimal_transport.py
@@ -0,0 +1,84 @@
+import torch
+from torch import Tensor
+
+def build_cost_matrix_1d(
+    len_a:      int,
+    len_b:      int,
+    device:     str = 'cuda',
+    dtype:      torch.dtype = torch.float32,
+    p:          int = 2                         #wasserstein distance order, default is 2 
+) -> torch.Tensor:
+
+    pos_a = torch.arange(len_a, device=device, dtype=dtype)   
+    pos_b = torch.arange(len_b, device=device, dtype=dtype)   
+
+    diff = pos_a[:, None] - pos_b[None, :]
+
+    if p == 1:
+        M = torch.abs(diff)
+    elif p == 2:
+        M = diff.pow(2)
+    else:
+        M = torch.abs(diff).pow(p)
+
+    return M
+
+
+#sinkhorn iteration
+def sinkhorn_binary(
+        source:         Tensor,          #(B,C,*Spatial)
+        target:         Tensor,          #(B,C,*Spatial)
+        cost_matrix:    Tensor,
+        reg:            float=1.0,
+        itrstep:        int=0,
+        threshold:      float=0,
+        device=         'cuda'
+        ):
+    
+    ''' Calculate the optimal transport plan using Sinkhorn iterations in the binary case. 
+    
+    Args:
+        source is the input tensor, 
+        target is the output tensor,
+        cost_matrix is the cost of transporting mass from source to target,
+        reg is the regularization parameter, and itrstep is the number of iterations to perform.
+    '''
+    #Device check
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA device not available.")
+
+    #check dimensions,and reshape to (B,C,N)
+    if not source.ndim == target.ndim:
+        raise ValueError("Source and target must have the same number of dimensions.")
+    else:
+        source=source.view(source.size(0),source.size(1),-1)
+        target=target.view(target.size(0),target.size(1),-1)
+        
+    #prepare for iteration
+    cost_matrix=build_cost_matrix_1d(source.size(-1),target.size(-1),device=device)
+
+    source/=source.sum(dim = -1,keepdim = True).to(device)
+    target/=target.sum(dim = -1,keepdim = True).to(device)
+    
+    u = torch.ones_like(source, device = device)
+    v = torch.ones_like(target, device = device)
+    k = torch.exp(- cost_matrix / reg)
+
+    #generate sinkhorn iteration judged by threshold or iteration times
+    if itrstep == 0:
+        step = 1000
+    for i in range(step):
+        if ((source / (k @ v)) - u).abs().sum().mean() < threshold:
+            convergence_flag = True
+            break
+        else:
+            u = source / (k @ v)
+            v = target / (k.t() @ u)
+            P = torch.diag(u) @ k @ torch.diag(v)
+            convergence_flag = False
+
+    if threshold == 0:
+        print(f"Sinkhorn iteration completed in {i+1} steps.")
+    else:
+        print(f"Sinkhorn iteration completed in {i+1} steps. Convergence: {convergence_flag}")
+    return P

From e668c242eee0709f5af8007d6af1a0c702af59bd Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Thu, 23 Apr 2026 14:26:45 +0800
Subject: [PATCH 03/20] debug ot function and add test

---
 test/test_optimal_transport.py  | 157 ++++++++++++++++++++++++++++++--
 torchmorph/optimal_transport.py | 106 +++++++++++++--------
 2 files changed, 216 insertions(+), 47 deletions(-)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 99e5ec6..73c45d1 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -2,17 +2,160 @@
 import sys
 import importlib.util
 from pathlib import Path
+import torch
+import ot
+import numpy as np
 
-# 项目根目录（test 的父目录）
+# Project root directory (parent of test)
 project_root = Path(__file__).parent.parent
 
-# 目标模块路径
+# Target module path
 module_path = project_root / "torchmorph" / "optimal_transport.py"
 
-# 加载模块
+# Load module
 spec = importlib.util.spec_from_file_location("optimal_transport", module_path)
-ot = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(ot)
+tr = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(tr)
 
-a=ot.build_cost_matrix_1d(2,3,'cuda')
-print(a)
+def test_build_cost_matrix():
+    """Simple test for the correctness of the build_cost_matrix function."""
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # Test 1D shape
+    shape_1d = (5,)
+    cost_1d = tr.build_cost_matrix(shape_1d, device=device, p=2)
+    N_1d = shape_1d[0]
+    assert cost_1d.shape == (N_1d, N_1d), f"Expected shape {(N_1d, N_1d)}, got {cost_1d.shape}"
+    # Diagonal should be 0
+    assert torch.allclose(cost_1d.diag(), torch.zeros(N_1d, device=device)), "Diagonal not zero"
+    # Symmetry
+    assert torch.allclose(cost_1d, cost_1d.T), "Matrix not symmetric"
+    # Distance between adjacent elements should be 1
+    assert torch.allclose(cost_1d[0, 1], torch.tensor(1.0, device=device)), "Adjacent distance should be 1"
+    
+    # ---------- 2. Test 2D shape ----------
+    shape_2d = (3, 4)  # H=3, W=4
+    cost_2d = tr.build_cost_matrix(shape_2d, device=device, p=2)
+    N_2d = shape_2d[0] * shape_2d[1]
+    assert cost_2d.shape == (N_2d, N_2d), f"Expected shape {(N_2d, N_2d)}, got {cost_2d.shape}"
+    # Diagonal should be 0
+    assert torch.allclose(cost_2d.diag(), torch.zeros(N_2d, device=device)), "Diagonal not zero"
+    # Symmetry
+    assert torch.allclose(cost_2d, cost_2d.T), "Matrix not symmetric"
+    
+    idx_00 = 0 * shape_2d[1] + 0  # Row 0, Column 0
+    idx_01 = 0 * shape_2d[1] + 1  # Row 0, Column 1
+    assert torch.allclose(cost_2d[idx_00, idx_01], torch.tensor(1.0, device=device)), \
+           f"Expected distance 1.0 between adjacent columns, got {cost_2d[idx_00, idx_01]}"
+    
+    # ---------- 3. Test 3D shape ----------
+    shape_3d = (2, 2, 2)  # A small 2x2x2 volume
+    cost_3d = tr.build_cost_matrix(shape_3d, device=device, p=2)
+    N_3d = 8
+    assert cost_3d.shape == (N_3d, N_3d), f"Expected shape {(N_3d, N_3d)}, got {cost_3d.shape}"
+    assert torch.allclose(cost_3d.diag(), torch.zeros(N_3d, device=device))
+    assert torch.allclose(cost_3d, cost_3d.T)
+    
+    # ---------- 4. Test different p values (Manhattan distance) ----------
+    shape = (3,)
+    cost_l1 = tr.build_cost_matrix(shape, device=device, p=1)
+    cost_l2 = tr.build_cost_matrix(shape, device=device, p=2)
+
+    expected = torch.tensor([[0.,1.,2.],[1.,0.,1.],[2.,1.,0.]], device=device)
+    assert torch.allclose(cost_l1, expected), "L1 distance matrix incorrect"
+    assert torch.allclose(cost_l2, expected), "L2 distance matrix incorrect"
+
+
+def test_sinkhorn_balanced_print():
+
+    print("=" * 50)
+    print("Testing sinkhorn_balanced function")
+    print(f"Input shape: ")
+           
+    # Check CUDA availability
+    if not torch.cuda.is_available():
+        print("Warning: CUDA not available, tests may fail.")
+        device = 'cpu'
+    else:
+        device = 'cuda'
+        print("CUDA available, running on GPU.")
+    
+    # Create two random distributions
+    source = torch.rand(1, 1, 2, 2, device=device)
+    target = torch.rand(1, 1, 2, 2, device=device)
+    
+    print("\n--- Input source ---")
+    print(source)
+    print("\n--- Input target ---")
+    print(target)
+    
+    print("\n--- Calling sinkhorn_balanced ---")
+    try:
+        P = tr.sinkhorn_balanced(
+            source, target,
+            cost_matrix=None,   
+            reg=0.1,
+            itrstep=100,
+            threshold=1e-5,
+            returngrad=False,
+            device=device
+        )
+        print("\n--- Output P (transport plan) shape ---")
+        print(P.shape)
+        print("\n--- Output P content (partial) ---")
+        print(P)
+    except Exception as e:
+        print("\n--- Error occurred ---")
+        print(f"Error type: {type(e).__name__}")
+        print(f"Error message: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+
+def test_sinkhorn_vs_pot():
+    """Compare custom sinkhorn_balanced with POT's ot.sinkhorn."""
+    print("\n" + "=" * 50)
+    print("Comparing custom Sinkhorn with POT library")
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"Device: {device}")
+
+    source = torch.rand(1, 1, 2, 2, device=device)
+    target = torch.rand(1, 1, 2, 2, device=device)
+
+    # Prepare data for POT (CPU numpy float64)
+    source_flat = source.view(-1)
+    target_flat = target.view(-1)
+    a = (source_flat / source_flat.sum()).cpu().numpy().astype(np.float64)
+    b = (target_flat / target_flat.sum()).cpu().numpy().astype(np.float64)
+
+    cost_matrix = tr.build_cost_matrix((2, 2), device='cpu', p=2)
+    M = cost_matrix.cpu().numpy().astype(np.float64)
+
+    # Custom
+    P_custom = tr.sinkhorn_balanced(
+        source, target,
+        cost_matrix=None,
+        reg=0.1,
+        itrstep=200,
+        threshold=1e-6,
+        device=device
+    )
+    P_custom_np = P_custom[0, 0].cpu().numpy()
+
+    # POT
+    P_pot = ot.sinkhorn(a, b, M, reg=0.1, numItermax=200, stopThr=1e-6)
+
+    max_diff = np.abs(P_custom_np - P_pot).max()
+    print(f"Max absolute difference: {max_diff:.2e}")
+
+    if max_diff < 1e-2:
+        print("✅ Custom implementation matches POT.")
+    else:
+        print("⚠️ Difference exceeds tolerance.")
+
+if __name__ == "__main__":
+    test_sinkhorn_balanced_print()
+    test_sinkhorn_vs_pot()
+    test_build_cost_matrix()
\ No newline at end of file
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 1623de6..cfd9aa2 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -1,84 +1,110 @@
 import torch
+from typing import Optional
 from torch import Tensor
 
-def build_cost_matrix_1d(
-    len_a:      int,
-    len_b:      int,
+#generate cost matrix for multi-dimensional tensor like 2D or 3D tensor
+def build_cost_matrix(
+    shape,
     device:     str = 'cuda',
-    dtype:      torch.dtype = torch.float32,
-    p:          int = 2                         #wasserstein distance order, default is 2 
-) -> torch.Tensor:
-
-    pos_a = torch.arange(len_a, device=device, dtype=dtype)   
-    pos_b = torch.arange(len_b, device=device, dtype=dtype)   
-
-    diff = pos_a[:, None] - pos_b[None, :]
-
-    if p == 1:
-        M = torch.abs(diff)
-    elif p == 2:
-        M = diff.pow(2)
-    else:
-        M = torch.abs(diff).pow(p)
-
-    return M
+    p           =2 #Euclidean distance by default
+    )-> torch.Tensor:
+ 
+    coords = torch.stack(torch.meshgrid([torch.arange(s) for s in shape], indexing="ij"), dim=-1)
+    coords_flat = coords.view(-1, coords.size(-1)).float()
+    
 
+    cost_matrix = torch.cdist(coords_flat, coords_flat, p=p).to(device)
+    return cost_matrix
 
 #sinkhorn iteration
-def sinkhorn_binary(
+def sinkhorn_balanced(
         source:         Tensor,          #(B,C,*Spatial)
         target:         Tensor,          #(B,C,*Spatial)
-        cost_matrix:    Tensor,
+        cost_matrix:    Optional[Tensor] = None,
         reg:            float=1.0,
         itrstep:        int=0,
         threshold:      float=0,
+        returngrad:     bool=False,      #whether to calculate gradient
         device=         'cuda'
         ):
     
-    ''' Calculate the optimal transport plan using Sinkhorn iterations in the binary case. 
+    ''' Calculate the optimal transport plan using Sinkhorn iterations in the balanced case. 
     
     Args:
         source is the input tensor, 
         target is the output tensor,
         cost_matrix is the cost of transporting mass from source to target,
         reg is the regularization parameter, and itrstep is the number of iterations to perform.
+    Returns:
+        The optimal transport plan P, and optionally the gradient if returngrad is True.
     '''
     #Device check
     if not torch.cuda.is_available():
         raise ValueError("CUDA device not available.")
 
-    #check dimensions,and reshape to (B,C,N)
+    #get batch and channel dimensions
+    B,C = source.shape[:2]
+    #check dimensions,reshape and choose strategy
     if not source.ndim == target.ndim:
         raise ValueError("Source and target must have the same number of dimensions.")
-    else:
-        source=source.view(source.size(0),source.size(1),-1)
-        target=target.view(target.size(0),target.size(1),-1)
+    elif source.ndim == 2:
+        if cost_matrix is None:
+            cost_matrix=build_cost_matrix(source, device=device, p=2)
+        source=source.view(-1)
+        target=target.view(-1)
+    elif source.ndim > 2:
+        #cut and reshape source and target to coordinate tensors, then calculate cost matrix
+        if cost_matrix is None:
+            spatial_shape = source.shape[2:]
+            cost_matrix=build_cost_matrix(spatial_shape, device=device, p=2)
+        source=source.view(B,C,-1)
+        target=target.view(B,C,-1)
         
-    #prepare for iteration
-    cost_matrix=build_cost_matrix_1d(source.size(-1),target.size(-1),device=device)
 
     source/=source.sum(dim = -1,keepdim = True).to(device)
     target/=target.sum(dim = -1,keepdim = True).to(device)
     
+    #initialize u and v, and expand dimensions for batch and channel
     u = torch.ones_like(source, device = device)
     v = torch.ones_like(target, device = device)
     k = torch.exp(- cost_matrix / reg)
 
+    k=k.unsqueeze(0).unsqueeze(0).expand(B,C,-1,-1)
+
     #generate sinkhorn iteration judged by threshold or iteration times
-    if itrstep == 0:
+    if itrstep == 0 and threshold > 0:
         step = 1000
+    elif itrstep == 0 and threshold == 0:
+        raise ValueError("Invalid iteration or step settings. ")
+    else:
+        step = itrstep
+    
+    #iteration
+    convergence_flag = False
     for i in range(step):
-        if ((source / (k @ v)) - u).abs().sum().mean() < threshold:
+        u_old = u.clone()
+        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-9)
+        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-9)
+        if (u_old - u).abs().sum().mean() < threshold:
             convergence_flag = True
             break
-        else:
-            u = source / (k @ v)
-            v = target / (k.t() @ u)
-            P = torch.diag(u) @ k @ torch.diag(v)
-            convergence_flag = False
 
-    if threshold == 0:
-        print(f"Sinkhorn iteration completed in {i+1} steps.")
+
+    #diagonalize u and v, then calculate optimal transport plan P
+    u = torch.diag_embed(u)
+    v = torch.diag_embed(v)      
+    P = u @ k @ v
+
+    if convergence_flag == True:
+        print(f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}")
+    else:
+        print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete." )
+
+    #calculate gradient after iteration covergence ,return results    
+    if returngrad == True:
+        log_u = torch.log(u)
+        mean_log_u = sum(log_u) / k
+        grad = (log_u - mean_log_u) / reg #gradient
+        return P, grad
     else:
-        print(f"Sinkhorn iteration completed in {i+1} steps. Convergence: {convergence_flag}")
-    return P
+        return P

From 660ff961415e4385eddc409a9974f213c4e03a59 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Fri, 24 Apr 2026 16:58:23 +0800
Subject: [PATCH 04/20] add basic benchmark

---
 benchmark/optimal_transport.py  | 142 ++++++++++++++++++++++++++++++
 test/test_optimal_transport.py  | 100 +++++++--------------
 torchmorph/optimal_transport.py | 148 ++++++++++++++++++--------------
 3 files changed, 259 insertions(+), 131 deletions(-)
 create mode 100644 benchmark/optimal_transport.py

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
new file mode 100644
index 0000000..c832cf0
--- /dev/null
+++ b/benchmark/optimal_transport.py
@@ -0,0 +1,142 @@
+import importlib.util
+from pathlib import Path
+
+import ot
+import torch
+from torch.utils import benchmark
+
+# Project root directory (parent of test)
+project_root = Path(__file__).parent.parent
+
+# Target module path and Load module
+module_path = project_root / "torchmorph" / "optimal_transport.py"
+spec = importlib.util.spec_from_file_location("optimal_transport", module_path)
+tr = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(tr)
+
+torch.manual_seed(42)
+
+
+def run_sinkhorn_balanced_benchmark():
+    # Set parameters
+    B, C, H, W = 1, 1, 32, 32
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+
+    # Prepare input data
+    source = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    target = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    source = source / source.sum()
+    target = target / target.sum()
+
+    globals_dict = {
+        "tr": tr,
+        "source": source,
+        "target": target,
+    }
+
+    stmt = "tr.sinkhorn_balanced(source, target, reg=10.0, itrstep=100)"
+
+    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
+    # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
+    number_of_runs = 50
+    result_sinkhorn_balanced = timer.timeit(number_of_runs)
+    return result_sinkhorn_balanced
+
+
+def run_ot_sinkhorn_benchmark():
+    # Set parameters
+    H, W = 32, 32
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+
+    # Prepare input data
+    source = torch.rand(H * W, device=device, dtype=dtype)
+    target = torch.rand(H * W, device=device, dtype=dtype)
+    source = source / source.sum()
+    target = target / target.sum()
+    M = tr.build_cost_matrix((H, W), device=device, p=2)
+
+    globals_dict = {
+        "ot": ot,
+        "source": source,
+        "target": target,
+        "M": M,
+    }
+
+    stmt = "ot.sinkhorn(source, target, M, reg=10.0, numItermax=100)"
+
+    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
+    # timer.blocked_autorange(min_run_time=1)
+    number_of_runs = 50
+    result_ot = timer.timeit(number_of_runs)
+    return result_ot
+
+
+def large_scale_sinkhorn_balanced_benchmark():
+    # Set parameters
+    B, C, H, W = 1, 1, 128, 128
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+
+    # Prepare input data
+    source = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    target = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    source = source / source.sum()
+    target = target / target.sum()
+
+    globals_dict = {
+        "tr": tr,
+        "source": source,
+        "target": target,
+    }
+
+    stmt = "tr.sinkhorn_balanced(source, target, reg=10.0, itrstep=100)"
+
+    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
+    # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
+    number_of_runs = 10
+    result_sinkhorn_balanced = timer.timeit(number_of_runs)
+    return result_sinkhorn_balanced
+
+
+def batch_channel_sinkhorn_balanced_benchmark():
+    # Set parameters
+    B, C, H, W = 32, 32, 32, 32
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+
+    # Prepare input data
+    source = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    target = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    source = source / source.sum()
+    target = target / target.sum()
+
+    globals_dict = {
+        "tr": tr,
+        "source": source,
+        "target": target,
+    }
+
+    stmt = "tr.sinkhorn_balanced(source, target, reg=10.0, itrstep=100)"
+
+    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
+    # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
+    number_of_runs = 10
+    result_sinkhorn_balanced = timer.timeit(number_of_runs)
+    return result_sinkhorn_balanced
+
+
+if __name__ == "__main__":
+    result_sinkhorn_balanced = run_sinkhorn_balanced_benchmark()
+    result_ot = run_ot_sinkhorn_benchmark()
+    result_large_scale = large_scale_sinkhorn_balanced_benchmark()
+    result_batch_channel = batch_channel_sinkhorn_balanced_benchmark()
+
+print(result_sinkhorn_balanced)
+print(result_ot)
+print(result_large_scale)
+print(result_batch_channel)
+
+speedup = result_ot.mean / result_sinkhorn_balanced.mean
+print(f"My sinkhorn to OT sinkhorn Speedup: {speedup:.1f}x")
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 73c45d1..bb44fbf 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -1,10 +1,7 @@
-# test_optimal_transport.py
-import sys
 import importlib.util
 from pathlib import Path
+
 import torch
-import ot
-import numpy as np
 
 # Project root directory (parent of test)
 project_root = Path(__file__).parent.parent
@@ -17,22 +14,28 @@
 tr = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(tr)
 
+
 def test_build_cost_matrix():
     """Simple test for the correctness of the build_cost_matrix function."""
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    
+
     # Test 1D shape
     shape_1d = (5,)
     cost_1d = tr.build_cost_matrix(shape_1d, device=device, p=2)
     N_1d = shape_1d[0]
-    assert cost_1d.shape == (N_1d, N_1d), f"Expected shape {(N_1d, N_1d)}, got {cost_1d.shape}"
+    assert cost_1d.shape == (
+        N_1d,
+        N_1d,
+    ), f"Expected shape {(N_1d, N_1d)}, got {cost_1d.shape}"
     # Diagonal should be 0
     assert torch.allclose(cost_1d.diag(), torch.zeros(N_1d, device=device)), "Diagonal not zero"
     # Symmetry
     assert torch.allclose(cost_1d, cost_1d.T), "Matrix not symmetric"
     # Distance between adjacent elements should be 1
-    assert torch.allclose(cost_1d[0, 1], torch.tensor(1.0, device=device)), "Adjacent distance should be 1"
-    
+    assert torch.allclose(
+        cost_1d[0, 1], torch.tensor(1.0, device=device)
+    ), "Adjacent distance should be 1"
+
     # ---------- 2. Test 2D shape ----------
     shape_2d = (3, 4)  # H=3, W=4
     cost_2d = tr.build_cost_matrix(shape_2d, device=device, p=2)
@@ -42,12 +45,13 @@ def test_build_cost_matrix():
     assert torch.allclose(cost_2d.diag(), torch.zeros(N_2d, device=device)), "Diagonal not zero"
     # Symmetry
     assert torch.allclose(cost_2d, cost_2d.T), "Matrix not symmetric"
-    
+
     idx_00 = 0 * shape_2d[1] + 0  # Row 0, Column 0
     idx_01 = 0 * shape_2d[1] + 1  # Row 0, Column 1
-    assert torch.allclose(cost_2d[idx_00, idx_01], torch.tensor(1.0, device=device)), \
-           f"Expected distance 1.0 between adjacent columns, got {cost_2d[idx_00, idx_01]}"
-    
+    assert torch.allclose(
+        cost_2d[idx_00, idx_01], torch.tensor(1.0, device=device)
+    ), f"Expected distance 1.0 between adjacent columns, got {cost_2d[idx_00, idx_01]}"
+
     # ---------- 3. Test 3D shape ----------
     shape_3d = (2, 2, 2)  # A small 2x2x2 volume
     cost_3d = tr.build_cost_matrix(shape_3d, device=device, p=2)
@@ -55,13 +59,14 @@ def test_build_cost_matrix():
     assert cost_3d.shape == (N_3d, N_3d), f"Expected shape {(N_3d, N_3d)}, got {cost_3d.shape}"
     assert torch.allclose(cost_3d.diag(), torch.zeros(N_3d, device=device))
     assert torch.allclose(cost_3d, cost_3d.T)
-    
+
     # ---------- 4. Test different p values (Manhattan distance) ----------
-    shape = (3,)
+    shape = (1, 1, 1, 3)  # Correctly define shape before usage
+    print(f"Input shape: {shape}")
     cost_l1 = tr.build_cost_matrix(shape, device=device, p=1)
     cost_l2 = tr.build_cost_matrix(shape, device=device, p=2)
 
-    expected = torch.tensor([[0.,1.,2.],[1.,0.,1.],[2.,1.,0.]], device=device)
+    expected = torch.tensor([[0.0, 1.0, 2.0], [1.0, 0.0, 1.0], [2.0, 1.0, 0.0]], device=device)
     assert torch.allclose(cost_l1, expected), "L1 distance matrix incorrect"
     assert torch.allclose(cost_l2, expected), "L2 distance matrix incorrect"
 
@@ -70,8 +75,8 @@ def test_sinkhorn_balanced_print():
 
     print("=" * 50)
     print("Testing sinkhorn_balanced function")
-    print(f"Input shape: ")
-           
+    print("Input shape: ")
+
     # Check CUDA availability
     if not torch.cuda.is_available():
         print("Warning: CUDA not available, tests may fail.")
@@ -79,26 +84,28 @@ def test_sinkhorn_balanced_print():
     else:
         device = 'cuda'
         print("CUDA available, running on GPU.")
-    
+
     # Create two random distributions
     source = torch.rand(1, 1, 2, 2, device=device)
     target = torch.rand(1, 1, 2, 2, device=device)
-    
+
     print("\n--- Input source ---")
     print(source)
     print("\n--- Input target ---")
     print(target)
-    
+
     print("\n--- Calling sinkhorn_balanced ---")
     try:
         P = tr.sinkhorn_balanced(
-            source, target,
-            cost_matrix=None,   
+            source,
+            target,
+            cost_matrix=None,
             reg=0.1,
             itrstep=100,
             threshold=1e-5,
             returngrad=False,
-            device=device
+            device=device,
+            p=2,
         )
         print("\n--- Output P (transport plan) shape ---")
         print(P.shape)
@@ -109,53 +116,10 @@ def test_sinkhorn_balanced_print():
         print(f"Error type: {type(e).__name__}")
         print(f"Error message: {e}")
         import traceback
-        traceback.print_exc()
-
-
-
-def test_sinkhorn_vs_pot():
-    """Compare custom sinkhorn_balanced with POT's ot.sinkhorn."""
-    print("\n" + "=" * 50)
-    print("Comparing custom Sinkhorn with POT library")
-
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    print(f"Device: {device}")
 
-    source = torch.rand(1, 1, 2, 2, device=device)
-    target = torch.rand(1, 1, 2, 2, device=device)
+        traceback.print_exc()
 
-    # Prepare data for POT (CPU numpy float64)
-    source_flat = source.view(-1)
-    target_flat = target.view(-1)
-    a = (source_flat / source_flat.sum()).cpu().numpy().astype(np.float64)
-    b = (target_flat / target_flat.sum()).cpu().numpy().astype(np.float64)
-
-    cost_matrix = tr.build_cost_matrix((2, 2), device='cpu', p=2)
-    M = cost_matrix.cpu().numpy().astype(np.float64)
-
-    # Custom
-    P_custom = tr.sinkhorn_balanced(
-        source, target,
-        cost_matrix=None,
-        reg=0.1,
-        itrstep=200,
-        threshold=1e-6,
-        device=device
-    )
-    P_custom_np = P_custom[0, 0].cpu().numpy()
-
-    # POT
-    P_pot = ot.sinkhorn(a, b, M, reg=0.1, numItermax=200, stopThr=1e-6)
-
-    max_diff = np.abs(P_custom_np - P_pot).max()
-    print(f"Max absolute difference: {max_diff:.2e}")
-
-    if max_diff < 1e-2:
-        print("✅ Custom implementation matches POT.")
-    else:
-        print("⚠️ Difference exceeds tolerance.")
 
 if __name__ == "__main__":
     test_sinkhorn_balanced_print()
-    test_sinkhorn_vs_pot()
-    test_build_cost_matrix()
\ No newline at end of file
+    test_build_cost_matrix()
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index cfd9aa2..070a51f 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -1,110 +1,132 @@
-import torch
 from typing import Optional
+
+import torch
 from torch import Tensor
 
-#generate cost matrix for multi-dimensional tensor like 2D or 3D tensor
+
+# generate cost matrix for multi-dimensional tensor like 2D or 3D tensor
 def build_cost_matrix(
-    shape,
-    device:     str = 'cuda',
-    p           =2 #Euclidean distance by default
-    )-> torch.Tensor:
- 
+    shape, device: str = 'cuda', p=2  # Euclidean distance by default
+) -> torch.Tensor:
+
     coords = torch.stack(torch.meshgrid([torch.arange(s) for s in shape], indexing="ij"), dim=-1)
     coords_flat = coords.view(-1, coords.size(-1)).float()
-    
 
     cost_matrix = torch.cdist(coords_flat, coords_flat, p=p).to(device)
     return cost_matrix
 
-#sinkhorn iteration
+
+# sinkhorn iteration
 def sinkhorn_balanced(
-        source:         Tensor,          #(B,C,*Spatial)
-        target:         Tensor,          #(B,C,*Spatial)
-        cost_matrix:    Optional[Tensor] = None,
-        reg:            float=1.0,
-        itrstep:        int=0,
-        threshold:      float=0,
-        returngrad:     bool=False,      #whether to calculate gradient
-        device=         'cuda'
-        ):
-    
-    ''' Calculate the optimal transport plan using Sinkhorn iterations in the balanced case. 
-    
+    source: Tensor,  # (B,C,*Spatial)
+    target: Tensor,  # (B,C,*Spatial)
+    cost_matrix: Optional[Tensor] = None,
+    p: int = 2,  # only effective when cost_matrix is None
+    reg: float = 1.0,
+    itrstep: int = 0,
+    threshold: float = 0,
+    returngrad: bool = False,  # whether to calculate gradient
+    device: str = 'cuda',
+):
+    """Compute the balanced optimal transport plan via Sinkhorn iterations.
+
+    Solves the entropic regularized optimal transport problem between
+    source and target distributions. Supports arbitrary spatial dimensions
+    and can optionally return the gradient of the dual potential.
+
     Args:
-        source is the input tensor, 
-        target is the output tensor,
-        cost_matrix is the cost of transporting mass from source to target,
-        reg is the regularization parameter, and itrstep is the number of iterations to perform.
+        source (Tensor): Input distribution, shape (B, C, *Spatial) or (*Spatial).
+        target (Tensor): Target distribution, same shape as source.
+        cost_matrix (Tensor, optional): Pairwise cost matrix of shape (N, N).
+            If None, it is built from the spatial coordinates using Lp distance.
+        p (int): Lp norm order for cost matrix construction. Defaults to 2.
+        reg (float): Entropic regularization strength. Defaults to 1.0.
+        itrstep (int): Number of Sinkhorn iterations. If 0 and threshold > 0,
+            a maximum of 1000 steps is performed. Defaults to 0.
+        threshold (float): Convergence threshold on the mean absolute change
+            of u. Ignored if itrstep > 0. Defaults to 0.
+        returngrad (bool): If True, also returns the gradient of the dual
+            potential (useful for barycenter tasks). Defaults to False.
+        device (str): Device for computation; typically 'cuda'. Defaults to 'cuda'.
+
     Returns:
-        The optimal transport plan P, and optionally the gradient if returngrad is True.
-    '''
-    #Device check
+        Tensor or tuple:
+            - **P**  Optimal transport plan of shape (B, C, N, N).
+            - **grad**  Gradient of the dual potential, returned only if
+              `returngrad` is True.
+
+    Raises:
+        ValueError: If CUDA is not available, or source/target have mismatched
+            dimensions, or non positive total mass, or invalid iteration settings.
+    """
+    # Device check
     if not torch.cuda.is_available():
         raise ValueError("CUDA device not available.")
 
-    #get batch and channel dimensions
-    B,C = source.shape[:2]
-    #check dimensions,reshape and choose strategy
+    # get batch and channel dimensions
+    B, C = source.shape[:2]
+    # check dimensions,reshape and choose strategy
     if not source.ndim == target.ndim:
         raise ValueError("Source and target must have the same number of dimensions.")
     elif source.ndim == 2:
         if cost_matrix is None:
-            cost_matrix=build_cost_matrix(source, device=device, p=2)
-        source=source.view(-1)
-        target=target.view(-1)
+            cost_matrix = build_cost_matrix(source, device=device, p=p)
+        source = source.view(-1)
+        target = target.view(-1)
     elif source.ndim > 2:
-        #cut and reshape source and target to coordinate tensors, then calculate cost matrix
+        # cut and reshape source and target to coordinate tensors, then calculate cost matrix
         if cost_matrix is None:
             spatial_shape = source.shape[2:]
-            cost_matrix=build_cost_matrix(spatial_shape, device=device, p=2)
-        source=source.view(B,C,-1)
-        target=target.view(B,C,-1)
-        
-
-    source/=source.sum(dim = -1,keepdim = True).to(device)
-    target/=target.sum(dim = -1,keepdim = True).to(device)
-    
-    #initialize u and v, and expand dimensions for batch and channel
-    u = torch.ones_like(source, device = device)
-    v = torch.ones_like(target, device = device)
-    k = torch.exp(- cost_matrix / reg)
-
-    k=k.unsqueeze(0).unsqueeze(0).expand(B,C,-1,-1)
-
-    #generate sinkhorn iteration judged by threshold or iteration times
+            cost_matrix = build_cost_matrix(spatial_shape, device=device, p=p)
+        source = source.view(B, C, -1)
+        target = target.view(B, C, -1)
+
+    if not source.sum().item() > 0 or not target.sum().item() > 0:
+        raise ValueError("Source and target must have positive mass.")
+
+    source /= source.sum(dim=-1, keepdim=True).to(device)
+    target /= target.sum(dim=-1, keepdim=True).to(device)
+
+    # initialize u and v, and expand dimensions for batch and channel
+    u = torch.ones_like(source, device=device)
+    v = torch.ones_like(target, device=device)
+    k = torch.exp(-cost_matrix / reg)
+
+    k = k.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
+
+    # generate sinkhorn iteration judged by threshold or iteration times
     if itrstep == 0 and threshold > 0:
         step = 1000
     elif itrstep == 0 and threshold == 0:
         raise ValueError("Invalid iteration or step settings. ")
     else:
         step = itrstep
-    
-    #iteration
+
+    # iteration
     convergence_flag = False
     for i in range(step):
         u_old = u.clone()
-        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-9)
-        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-9)
+        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-15)
+        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-15)
         if (u_old - u).abs().sum().mean() < threshold:
             convergence_flag = True
             break
 
-
-    #diagonalize u and v, then calculate optimal transport plan P
+    # diagonalize u and v, then calculate optimal transport plan P
     u = torch.diag_embed(u)
-    v = torch.diag_embed(v)      
+    v = torch.diag_embed(v)
     P = u @ k @ v
 
-    if convergence_flag == True:
+    if convergence_flag:
         print(f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}")
     else:
-        print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete." )
+        print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete.")
 
-    #calculate gradient after iteration covergence ,return results    
-    if returngrad == True:
+    # calculate gradient after iteration covergence ,return results
+    if returngrad:
         log_u = torch.log(u)
         mean_log_u = sum(log_u) / k
-        grad = (log_u - mean_log_u) / reg #gradient
+        grad = (log_u - mean_log_u) / reg  # gradient
         return P, grad
     else:
         return P

From 47ce04e52904388eb170674309bf13ae9d83fca6 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Tue, 28 Apr 2026 16:07:43 +0800
Subject: [PATCH 05/20] decouple the iteration from the derivation

---
 benchmark/optimal_transport.py  | 42 ++++++++------
 test/test_optimal_transport.py  |  3 +-
 torchmorph/optimal_transport.py | 99 +++++++++++++++++++--------------
 3 files changed, 81 insertions(+), 63 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index c832cf0..9089c7f 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -14,32 +14,36 @@
 tr = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(tr)
 
-torch.manual_seed(42)
-
 
 def run_sinkhorn_balanced_benchmark():
     # Set parameters
     B, C, H, W = 1, 1, 32, 32
+    myreg = 0.5
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32
+    torch.manual_seed(42)
 
     # Prepare input data
     source = torch.rand(B, C, H, W, device=device, dtype=dtype)
     target = torch.rand(B, C, H, W, device=device, dtype=dtype)
     source = source / source.sum()
     target = target / target.sum()
+    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
 
     globals_dict = {
         "tr": tr,
         "source": source,
         "target": target,
+        "myreg": myreg,
+        "cost_matrix": cost_matrix,
     }
 
-    stmt = "tr.sinkhorn_balanced(source, target, reg=10.0, itrstep=100)"
+    stmt = "tr.sinkhorn_balanced(source, target, cost_matrix=cost_matrix, reg=myreg,"
+    "itrstep=100,threshold=1e-5)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
-    number_of_runs = 50
+    number_of_runs = 100
     result_sinkhorn_balanced = timer.timeit(number_of_runs)
     return result_sinkhorn_balanced
 
@@ -47,8 +51,10 @@ def run_sinkhorn_balanced_benchmark():
 def run_ot_sinkhorn_benchmark():
     # Set parameters
     H, W = 32, 32
+    myreg = 0.5
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32
+    torch.manual_seed(42)
 
     # Prepare input data
     source = torch.rand(H * W, device=device, dtype=dtype)
@@ -62,13 +68,14 @@ def run_ot_sinkhorn_benchmark():
         "source": source,
         "target": target,
         "M": M,
+        "myreg": myreg,
     }
 
-    stmt = "ot.sinkhorn(source, target, M, reg=10.0, numItermax=100)"
+    stmt = "ot.sinkhorn(source, target, M, reg=1/myreg, numItermax=100, stopThr=1e-5)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # timer.blocked_autorange(min_run_time=1)
-    number_of_runs = 50
+    number_of_runs = 100
     result_ot = timer.timeit(number_of_runs)
     return result_ot
 
@@ -91,7 +98,7 @@ def large_scale_sinkhorn_balanced_benchmark():
         "target": target,
     }
 
-    stmt = "tr.sinkhorn_balanced(source, target, reg=10.0, itrstep=100)"
+    stmt = "tr.sinkhorn_balanced(source, target, reg=25.0, itrstep=100)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
@@ -118,7 +125,7 @@ def batch_channel_sinkhorn_balanced_benchmark():
         "target": target,
     }
 
-    stmt = "tr.sinkhorn_balanced(source, target, reg=10.0, itrstep=100)"
+    stmt = "tr.sinkhorn_balanced(source, target, reg=25.0, itrstep=100)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
@@ -130,13 +137,12 @@ def batch_channel_sinkhorn_balanced_benchmark():
 if __name__ == "__main__":
     result_sinkhorn_balanced = run_sinkhorn_balanced_benchmark()
     result_ot = run_ot_sinkhorn_benchmark()
-    result_large_scale = large_scale_sinkhorn_balanced_benchmark()
-    result_batch_channel = batch_channel_sinkhorn_balanced_benchmark()
-
-print(result_sinkhorn_balanced)
-print(result_ot)
-print(result_large_scale)
-print(result_batch_channel)
-
-speedup = result_ot.mean / result_sinkhorn_balanced.mean
-print(f"My sinkhorn to OT sinkhorn Speedup: {speedup:.1f}x")
+    # result_large_scale = large_scale_sinkhorn_balanced_benchmark()
+    # result_batch_channel = batch_channel_sinkhorn_balanced_benchmark()
+
+    print(result_sinkhorn_balanced)
+    print(result_ot)
+    speedup = result_ot.mean / result_sinkhorn_balanced.mean
+    print(f"My sinkhorn to OT sinkhorn Speedup: {speedup:.1f}x")
+    # print(result_large_scale)
+    # print(result_batch_channel)
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index bb44fbf..04ff5e3 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -79,8 +79,7 @@ def test_sinkhorn_balanced_print():
 
     # Check CUDA availability
     if not torch.cuda.is_available():
-        print("Warning: CUDA not available, tests may fail.")
-        device = 'cpu'
+        raise ValueError("Warning: CUDA not available, tests may fail.")
     else:
         device = 'cuda'
         print("CUDA available, running on GPU.")
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 070a51f..5a5306e 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -16,6 +16,10 @@ def build_cost_matrix(
     return cost_matrix
 
 
+def distribution_preprocess():
+    pass
+
+
 # sinkhorn iteration
 def sinkhorn_balanced(
     source: Tensor,  # (B,C,*Spatial)
@@ -25,75 +29,74 @@ def sinkhorn_balanced(
     reg: float = 1.0,
     itrstep: int = 0,
     threshold: float = 0,
-    returngrad: bool = False,  # whether to calculate gradient
+    returngrad: bool = False,  # whether to return the gradient of the dual potential
     device: str = 'cuda',
+    debugflag: bool = False,
 ):
-    """Compute the balanced optimal transport plan via Sinkhorn iterations.
+    """Compute a balanced entropic optimal transport plan with Sinkhorn iterations.
 
-    Solves the entropic regularized optimal transport problem between
-    source and target distributions. Supports arbitrary spatial dimensions
-    and can optionally return the gradient of the dual potential.
+    Uses the Gibbs kernel K = exp(-reg * C), where `reg` is the inverse
+    regularization scale. Source and target are clamped to be non-negative,
+    normalized over the flattened spatial dimension, and matched by iterative
+    Sinkhorn scaling.
 
     Args:
-        source (Tensor): Input distribution, shape (B, C, *Spatial) or (*Spatial).
-        target (Tensor): Target distribution, same shape as source.
+        source (Tensor): Source distribution, shape (B, C, *Spatial) or (*Spatial).
+        target (Tensor): Target distribution with the same shape as source.
         cost_matrix (Tensor, optional): Pairwise cost matrix of shape (N, N).
-            If None, it is built from the spatial coordinates using Lp distance.
-        p (int): Lp norm order for cost matrix construction. Defaults to 2.
-        reg (float): Entropic regularization strength. Defaults to 1.0.
-        itrstep (int): Number of Sinkhorn iterations. If 0 and threshold > 0,
-            a maximum of 1000 steps is performed. Defaults to 0.
-        threshold (float): Convergence threshold on the mean absolute change
-            of u. Ignored if itrstep > 0. Defaults to 0.
-        returngrad (bool): If True, also returns the gradient of the dual
-            potential (useful for barycenter tasks). Defaults to False.
-        device (str): Device for computation; typically 'cuda'. Defaults to 'cuda'.
+        p (int): Lp norm used to build the cost matrix when not provided.
+        reg (float): Positive inverse regularization parameter.
+        itrstep (int): Number of Sinkhorn iterations.
+        threshold (float): Stopping threshold on the change of the source scaling.
+        returngrad (bool): If True, also return scaling vectors u and v.
+        device (str): Computation device, typically 'cuda'.
 
     Returns:
-        Tensor or tuple:
-            - **P**  Optimal transport plan of shape (B, C, N, N).
-            - **grad**  Gradient of the dual potential, returned only if
-              `returngrad` is True.
+        dict: Contains "plan" and, when requested, "u" and "v".
 
     Raises:
-        ValueError: If CUDA is not available, or source/target have mismatched
-            dimensions, or non positive total mass, or invalid iteration settings.
+        ValueError: If dimensions, regularization, device, or stopping settings
+            are invalid.
     """
     # Device check
     if not torch.cuda.is_available():
         raise ValueError("CUDA device not available.")
 
-    # get batch and channel dimensions
-    B, C = source.shape[:2]
     # check dimensions,reshape and choose strategy
     if not source.ndim == target.ndim:
         raise ValueError("Source and target must have the same number of dimensions.")
     elif source.ndim == 2:
         if cost_matrix is None:
-            cost_matrix = build_cost_matrix(source, device=device, p=p)
-        source = source.view(-1)
-        target = target.view(-1)
+            cost_matrix = build_cost_matrix(source.shape, device=device, p=p)
+        source = source.view(1, 1, -1)
+        target = target.view(1, 1, -1)
     elif source.ndim > 2:
         # cut and reshape source and target to coordinate tensors, then calculate cost matrix
+        B, C = source.shape[:2]  # get batch and channel dimensions
         if cost_matrix is None:
             spatial_shape = source.shape[2:]
             cost_matrix = build_cost_matrix(spatial_shape, device=device, p=p)
         source = source.view(B, C, -1)
         target = target.view(B, C, -1)
 
-    if not source.sum().item() > 0 or not target.sum().item() > 0:
-        raise ValueError("Source and target must have positive mass.")
+    # ensure non-negativity
+    source = torch.clamp(source, min=0)
+    target = torch.clamp(target, min=0)
 
-    source /= source.sum(dim=-1, keepdim=True).to(device)
-    target /= target.sum(dim=-1, keepdim=True).to(device)
+    # normalize source and target to make them valid probability distributions
+    source /= torch.clamp(source.sum(dim=-1, keepdim=True).to(device), min=1e-12)
+    target /= torch.clamp(target.sum(dim=-1, keepdim=True).to(device), min=1e-12)
 
     # initialize u and v, and expand dimensions for batch and channel
     u = torch.ones_like(source, device=device)
     v = torch.ones_like(target, device=device)
-    k = torch.exp(-cost_matrix / reg)
+    k = torch.exp(-cost_matrix * reg)
 
     k = k.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
 
+    if reg <= 0:
+        raise ValueError("Regularization parameter must be positive.")
+
     # generate sinkhorn iteration judged by threshold or iteration times
     if itrstep == 0 and threshold > 0:
         step = 1000
@@ -106,16 +109,16 @@ def sinkhorn_balanced(
     convergence_flag = False
     for i in range(step):
         u_old = u.clone()
-        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-15)
-        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-15)
+        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
+        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
         if (u_old - u).abs().sum().mean() < threshold:
             convergence_flag = True
             break
 
     # diagonalize u and v, then calculate optimal transport plan P
-    u = torch.diag_embed(u)
-    v = torch.diag_embed(v)
-    P = u @ k @ v
+    u_diag = torch.diag_embed(u)
+    v_diag = torch.diag_embed(v)
+    P = u_diag @ k @ v_diag
 
     if convergence_flag:
         print(f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}")
@@ -124,9 +127,19 @@ def sinkhorn_balanced(
 
     # calculate gradient after iteration covergence ,return results
     if returngrad:
-        log_u = torch.log(u)
-        mean_log_u = sum(log_u) / k
-        grad = (log_u - mean_log_u) / reg  # gradient
-        return P, grad
+        return {"plan": P, "u": u, "v": v}
     else:
-        return P
+        return {"plan": P}
+
+
+def sinkhorn_gradient(u: Tensor, v: Tensor, reg: float):
+    u = torch.log(torch.clamp(u, min=1e-12))
+    v = torch.log(torch.clamp(v, min=1e-12))
+
+    u = u - u.mean(dim=-1, keepdim=True)  # remove constants
+    v = v - v.mean(dim=-1, keepdim=True)
+
+    grad_source = u / reg
+    grad_target = v / reg
+
+    return grad_source, grad_target

From bc47d27551eac4938d64be4b2330dfb153bcb2a0 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Mon, 4 May 2026 16:56:58 +0800
Subject: [PATCH 06/20] Implement Sinkhorn iterations in CUDA

---
 benchmark/optimal_transport.py  |  14 ++-
 test/test_optimal_transport.py  |   2 +-
 torchmorph/csrc/ot_kernel.cu    | 122 ++++++++++++++++++++++++-
 torchmorph/csrc/torchmorph.cpp  |  25 +++++
 torchmorph/optimal_transport.py | 156 ++++++++++++++++++++++++--------
 5 files changed, 274 insertions(+), 45 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index 9089c7f..eb13787 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -30,6 +30,10 @@ def run_sinkhorn_balanced_benchmark():
     target = target / target.sum()
     cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
 
+    source, target, cost_matrix = tr.data_preprocess(
+        source, target, cost_matrix=cost_matrix, p=2, device=device
+    )
+
     globals_dict = {
         "tr": tr,
         "source": source,
@@ -38,8 +42,10 @@ def run_sinkhorn_balanced_benchmark():
         "cost_matrix": cost_matrix,
     }
 
-    stmt = "tr.sinkhorn_balanced(source, target, cost_matrix=cost_matrix, reg=myreg,"
-    "itrstep=100,threshold=1e-5)"
+    stmt = (
+        "tr.sinkhorn_balanced(source, target, cost_matrix=cost_matrix, "
+        "reg=myreg, itrstep=100, threshold=1e-5)"
+    )
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
@@ -98,7 +104,7 @@ def large_scale_sinkhorn_balanced_benchmark():
         "target": target,
     }
 
-    stmt = "tr.sinkhorn_balanced(source, target, reg=25.0, itrstep=100)"
+    stmt = "tr.sinkhorn_balanced_full(source, target, reg=25.0, itrstep=100, threshold=1e-5)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
@@ -125,7 +131,7 @@ def batch_channel_sinkhorn_balanced_benchmark():
         "target": target,
     }
 
-    stmt = "tr.sinkhorn_balanced(source, target, reg=25.0, itrstep=100)"
+    stmt = "tr.sinkhorn_balanced_full(source, target, reg=25.0, itrstep=100, threshold=1e-5)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 04ff5e3..c309d2c 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -95,7 +95,7 @@ def test_sinkhorn_balanced_print():
 
     print("\n--- Calling sinkhorn_balanced ---")
     try:
-        P = tr.sinkhorn_balanced(
+        P = tr.sinkhorn_balanced_full(
             source,
             target,
             cost_matrix=None,
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
index ea028fa..839b5ce 100644
--- a/torchmorph/csrc/ot_kernel.cu
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -1 +1,121 @@
-// Kernel implementation for Optimal Transport
\ No newline at end of file
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <tuple>
+#include <cmath>
+
+__global__ void u_update(
+    const float* a,
+    const float* b,
+    const float* K,
+    float* u,
+    float* v,
+    int N
+){
+    int bidx = blockIdx.x;
+    int blk = blockDim.x;
+    int tid = threadIdx.x;
+    int total = N;
+    float local_sum = 0.0f;
+
+    if(bidx >= total){
+        return;
+    }
+
+    extern __shared__ float sdata[];
+
+    //get ready for sinkhorn:u= a / (k v)
+    for (int j = tid; j < N; j += blk){
+        local_sum += K[bidx*N+j] * v[j];
+    }//calculate one column
+
+    sdata[tid]=local_sum;//get one element of one column
+    __syncthreads();
+
+    
+    for (int stride=blk/2; stride > 0; stride /= 2){
+        if (tid < stride){
+        sdata[tid]+=sdata[tid+stride];
+        }
+        __syncthreads();
+    }
+
+    if (tid == 0){
+        u[bidx] = a[bidx] / (sdata[0]+1e-12f);
+    }
+}
+
+__global__ void v_update(
+    const float* a,
+    const float* b,
+    const float* K,
+    float* u,
+    float* v,
+    int N
+){
+    int bidx = blockIdx.x;
+    int blk = blockDim.x;
+    int tid = threadIdx.x;
+    int total = N;
+    float local_sum = 0.0f;
+
+    if(bidx >= total){
+        return;
+    }
+
+    extern __shared__ float sdata[];
+
+    //get ready for sinkhorn:u= b / (k^t u)
+    for (int j = tid; j < N; j += blk){
+        local_sum += K[bidx+j*N] * u[j];
+    }//calculate one column
+
+    sdata[tid]=local_sum;//get one element of one column
+    __syncthreads();
+
+    
+    for (int stride=blk/2; stride > 0; stride /= 2){
+        if (tid < stride){
+        sdata[tid]+=sdata[tid+stride];
+        }
+        __syncthreads();
+    }
+
+    if (tid == 0){
+        v[bidx] = b[bidx] / (sdata[0]+1e-12f);
+    }
+}
+
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_uv_cuda(
+    const torch::Tensor source,
+    const torch::Tensor target,
+    const torch::Tensor k,
+    torch::Tensor u,
+    torch::Tensor v,
+    int itrstep,
+    int N
+){
+    int threads = 256;
+    int blocks = N;
+    int shared_memory = threads * sizeof(float);
+
+    for (int i = 0; i < itrstep; i++){
+        u_update<<<blocks, threads, shared_memory>>>(
+            source.data_ptr<float>(),
+            target.data_ptr<float>(),
+            k.data_ptr<float>(),
+            u.data_ptr<float>(),
+            v.data_ptr<float>(),
+            N
+        );
+        v_update<<<blocks, threads, shared_memory>>>(
+            source.data_ptr<float>(),
+            target.data_ptr<float>(),
+            k.data_ptr<float>(),
+            u.data_ptr<float>(),
+            v.data_ptr<float>(),
+            N
+        );
+    }
+return std::make_tuple(u, v);
+}
\ No newline at end of file
diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp
index 0adac87..5241510 100644
--- a/torchmorph/csrc/torchmorph.cpp
+++ b/torchmorph/csrc/torchmorph.cpp
@@ -27,6 +27,18 @@ std::tuple<torch::Tensor, torch::Tensor> bfdt_cuda(
     bool return_indices
 );
 
+// Optimal Transport functions
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_uv_cuda(
+    const torch::Tensor source,
+    const torch::Tensor target,
+    const torch::Tensor k,
+    torch::Tensor u,
+    torch::Tensor v,
+    int itrstep,
+    int N
+);
+
+
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("add_cuda", &add_cuda, "Add tensor with scalar");
@@ -53,4 +65,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           py::arg("sampling") = std::vector<float>(),
           py::arg("return_distances") = true,
           py::arg("return_indices") = false);
+
+    // Optimal Transport
+    m.def("sinkhorn_uv_cuda", &sinkhorn_uv_cuda,
+          "Sinkhorn UV CUDA",
+          py::arg("source"),
+          py::arg("target"),
+          py::arg("k"),
+          py::arg("u"),
+          py::arg("v"),
+          py::arg("itrstep"),
+          py::arg("N"));
 }
+
+
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 5a5306e..60c7810 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -3,6 +3,8 @@
 import torch
 from torch import Tensor
 
+# from torchmorph import _C
+
 
 # generate cost matrix for multi-dimensional tensor like 2D or 3D tensor
 def build_cost_matrix(
@@ -16,22 +18,95 @@ def build_cost_matrix(
     return cost_matrix
 
 
-def distribution_preprocess():
-    pass
+def data_preprocess(
+    source: Tensor,
+    target: Tensor,
+    cost_matrix: Optional[Tensor] = None,
+    p: int = 2,
+    device: str = 'cuda',
+):
+    # check dimensions,reshape and choose strategy
+    if not source.ndim == target.ndim or not source.shape == target.shape:
+        raise ValueError(
+            "Source and target must have the same number of dimensions and the same shape."
+        )
+    elif source.ndim == 2:
+        if cost_matrix is None:
+            cost_matrix = build_cost_matrix(source.shape, device=device, p=p)
+        source = source.view(1, 1, -1)
+        target = target.view(1, 1, -1)
+    elif source.ndim > 2:
+        # cut and reshape source and target to coordinate tensors, then calculate cost matrix
+        B, C = source.shape[:2]  # get batch and channel dimensions
+        if cost_matrix is None:
+            spatial_shape = source.shape[2:]
+            cost_matrix = build_cost_matrix(spatial_shape, device=device, p=p)
+        source = source.view(B, C, -1)
+        target = target.view(B, C, -1)
+
+    # move to specified device(cuda) and ensure dtype float32
+    source = source.to(device=device, dtype=torch.float32)
+    target = target.to(device=device, dtype=torch.float32)
+    cost_matrix = cost_matrix.to(device=device, dtype=torch.float32)
+
+    # ensure non-negativity
+    source = torch.clamp(source, min=0)
+    target = torch.clamp(target, min=0)
+
+    # normalize source and target to make them valid probability distributions
+    source /= torch.clamp(source.sum(dim=-1, keepdim=True).to(device), min=1e-12)
+    target /= torch.clamp(target.sum(dim=-1, keepdim=True).to(device), min=1e-12)
+
+    return source, target, cost_matrix
+
+
+def sinkhorn_balanced_full(
+    source: Tensor,
+    target: Tensor,
+    cost_matrix: Optional[Tensor] = None,
+    p: int = 2,
+    reg: float = 1.0,
+    itrstep: int = 0,
+    threshold: float = 0,
+    returngrad: bool = False,
+    device: str = 'cuda',
+    verbose: bool = False,
+):
+    source, target, cost_matrix = data_preprocess(
+        source,
+        target,
+        cost_matrix=cost_matrix,
+        p=p,
+        device=device,
+    )
+    result = sinkhorn_balanced(
+        source,
+        target,
+        cost_matrix=cost_matrix,
+        reg=reg,
+        itrstep=itrstep,
+        threshold=threshold,
+        returnuv=returngrad,
+        device=device,
+        verbose=verbose,
+    )
+
+    if returngrad:
+        return result
+    return result["plan"]
 
 
 # sinkhorn iteration
 def sinkhorn_balanced(
     source: Tensor,  # (B,C,*Spatial)
     target: Tensor,  # (B,C,*Spatial)
-    cost_matrix: Optional[Tensor] = None,
-    p: int = 2,  # only effective when cost_matrix is None
+    cost_matrix: Tensor,
     reg: float = 1.0,
     itrstep: int = 0,
     threshold: float = 0,
-    returngrad: bool = False,  # whether to return the gradient of the dual potential
+    returnuv: bool = False,  # whether to return the gradient of the dual potential
     device: str = 'cuda',
-    debugflag: bool = False,
+    verbose: bool = False,
 ):
     """Compute a balanced entropic optimal transport plan with Sinkhorn iterations.
 
@@ -62,36 +137,15 @@ def sinkhorn_balanced(
     if not torch.cuda.is_available():
         raise ValueError("CUDA device not available.")
 
-    # check dimensions,reshape and choose strategy
-    if not source.ndim == target.ndim:
-        raise ValueError("Source and target must have the same number of dimensions.")
-    elif source.ndim == 2:
-        if cost_matrix is None:
-            cost_matrix = build_cost_matrix(source.shape, device=device, p=p)
-        source = source.view(1, 1, -1)
-        target = target.view(1, 1, -1)
-    elif source.ndim > 2:
-        # cut and reshape source and target to coordinate tensors, then calculate cost matrix
-        B, C = source.shape[:2]  # get batch and channel dimensions
-        if cost_matrix is None:
-            spatial_shape = source.shape[2:]
-            cost_matrix = build_cost_matrix(spatial_shape, device=device, p=p)
-        source = source.view(B, C, -1)
-        target = target.view(B, C, -1)
-
-    # ensure non-negativity
-    source = torch.clamp(source, min=0)
-    target = torch.clamp(target, min=0)
-
-    # normalize source and target to make them valid probability distributions
-    source /= torch.clamp(source.sum(dim=-1, keepdim=True).to(device), min=1e-12)
-    target /= torch.clamp(target.sum(dim=-1, keepdim=True).to(device), min=1e-12)
+    if source.ndim != 3 or target.ndim != 3:
+        raise ValueError("source and target must be preprocessed to shape (B, C, N).")
 
     # initialize u and v, and expand dimensions for batch and channel
     u = torch.ones_like(source, device=device)
     v = torch.ones_like(target, device=device)
     k = torch.exp(-cost_matrix * reg)
 
+    B, C = source.shape[:2]
     k = k.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
 
     if reg <= 0:
@@ -115,18 +169,17 @@ def sinkhorn_balanced(
             convergence_flag = True
             break
 
-    # diagonalize u and v, then calculate optimal transport plan P
-    u_diag = torch.diag_embed(u)
-    v_diag = torch.diag_embed(v)
-    P = u_diag @ k @ v_diag
+    # use broadcasting to calculate the optimal transport plan P from u, v and K
+    P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
 
-    if convergence_flag:
-        print(f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}")
-    else:
-        print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete.")
+    if verbose:
+        if convergence_flag:
+            print(f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}")
+        else:
+            print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete.")
 
     # calculate gradient after iteration covergence ,return results
-    if returngrad:
+    if returnuv:
         return {"plan": P, "u": u, "v": v}
     else:
         return {"plan": P}
@@ -143,3 +196,28 @@ def sinkhorn_gradient(u: Tensor, v: Tensor, reg: float):
     grad_target = v / reg
 
     return grad_source, grad_target
+
+
+def sinkhorn_balanced_cuda(
+    source: Tensor,
+    target: Tensor,
+    cost_matrix: Tensor,
+    reg: float = 1.0,
+    itrstep: int = 100,
+    returnuv: bool = False,
+):
+    from torchmorph import _C
+
+    N = source.shape[-1]
+    K = torch.exp(-cost_matrix * reg)
+    u = torch.ones_like(source)
+    v = torch.ones_like(target)
+
+    u, v = _C.sinkhorn_uv_cuda(source, target, K, u, v, int(itrstep), int(N))
+
+    P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
+
+    if returnuv:
+        return {"plan": P, "u": u, "v": v}
+    else:
+        return {"plan": P}

From 10d34330c1cb68b0e8fbee1f1f40eee1f08e4d54 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Thu, 7 May 2026 22:52:27 +0800
Subject: [PATCH 07/20] Add Sinkhorn CUDA benchmark and correctness tests

---
 benchmark/optimal_transport.py  | 161 ++++++++++++++++++++++++++++----
 test/test_optimal_transport.py  |  42 +++++++++
 torchmorph/csrc/ot_kernel.cu    |  14 ++-
 torchmorph/csrc/torchmorph.cpp  |   6 +-
 torchmorph/optimal_transport.py | 121 +++++++++++++++++-------
 5 files changed, 286 insertions(+), 58 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index eb13787..c2a97b2 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -15,10 +15,16 @@
 spec.loader.exec_module(tr)
 
 
-def run_sinkhorn_balanced_benchmark():
+def run_sinkhorn_balanced_benchmark(
+    B: int = 1,
+    C: int = 1,
+    H: int = 32,
+    W: int = 32,
+    myreg: float = 0.02,
+    itrstep: int = 0,
+    threshold: float = 1e-5,
+):
     # Set parameters
-    B, C, H, W = 1, 1, 32, 32
-    myreg = 0.5
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32
     torch.manual_seed(42)
@@ -40,24 +46,25 @@ def run_sinkhorn_balanced_benchmark():
         "target": target,
         "myreg": myreg,
         "cost_matrix": cost_matrix,
+        "itrstep": itrstep,
+        "threshold": threshold,
     }
 
     stmt = (
-        "tr.sinkhorn_balanced(source, target, cost_matrix=cost_matrix, "
-        "reg=myreg, itrstep=100, threshold=1e-5)"
+        "tr.sinkhorn_balanced_batch(source, target, cost_matrix=cost_matrix, "
+        "reg=myreg, itrstep=itrstep, threshold=threshold)"
     )
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
-    number_of_runs = 100
-    result_sinkhorn_balanced = timer.timeit(number_of_runs)
+    # number_of_runs = 100
+    result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=2)
     return result_sinkhorn_balanced
 
 
-def run_ot_sinkhorn_benchmark():
+def run_ot_sinkhorn_benchmark(H: int = 32, W: int = 32, myreg: float = 0.02):
     # Set parameters
-    H, W = 32, 32
-    myreg = 0.5
+
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32
     torch.manual_seed(42)
@@ -81,14 +88,58 @@ def run_ot_sinkhorn_benchmark():
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # timer.blocked_autorange(min_run_time=1)
-    number_of_runs = 100
-    result_ot = timer.timeit(number_of_runs)
+    # number_of_runs = 100
+    result_ot = timer.blocked_autorange(min_run_time=2)
     return result_ot
 
 
-def large_scale_sinkhorn_balanced_benchmark():
+@torch.no_grad()
+def run_sinkhorn_relative_error(H: int = 32, W: int = 32, myreg: float = 0.02):
+    # Set parameters
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float32
+    torch.manual_seed(42)
+
+    # Prepare the same input data for both implementations.
+    source = torch.rand(1, 1, H, W, device=device, dtype=dtype)
+    target = torch.rand(1, 1, H, W, device=device, dtype=dtype)
+    source = source / source.sum()
+    target = target / target.sum()
+    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
+
+    source_preprocessed, target_preprocessed, cost_matrix = tr.data_preprocess(
+        source, target, cost_matrix=cost_matrix, p=2, device=device
+    )
+
+    sinkhorn_plan = (
+        tr.sinkhorn_balanced_batch(
+            source_preprocessed,
+            target_preprocessed,
+            cost_matrix=cost_matrix,
+            reg=myreg,
+            itrstep=100,
+            threshold=1e-5,
+        )["plan"]
+        .squeeze(0)
+        .squeeze(0)
+    )
+
+    ot_plan = ot.sinkhorn(
+        source_preprocessed.squeeze(0).squeeze(0),
+        target_preprocessed.squeeze(0).squeeze(0),
+        cost_matrix,
+        reg=1 / myreg,
+        numItermax=100,
+        stopThr=1e-5,
+    )
+
+    return torch.linalg.norm(sinkhorn_plan - ot_plan) / torch.clamp(
+        torch.linalg.norm(ot_plan), min=1e-12
+    )
+
+
+def large_scale_sinkhorn_balanced_benchmark(B: int = 1, C: int = 1, H: int = 128, W: int = 128):
     # Set parameters
-    B, C, H, W = 1, 1, 128, 128
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32
 
@@ -141,14 +192,84 @@ def batch_channel_sinkhorn_balanced_benchmark():
 
 
 if __name__ == "__main__":
-    result_sinkhorn_balanced = run_sinkhorn_balanced_benchmark()
-    result_ot = run_ot_sinkhorn_benchmark()
+    # sizes=[10,20,50,100]
+    # for size in sizes:
+    #     i = 0
+    #     result_sinkhorn_balanced[] = run_sinkhorn_balanced_benchmark(B=1, C=1, H=size, W=size)
+    #     result_ot[] = run_ot_sinkhorn_benchmark(H=size, W=size)
+    #     i += 1
     # result_large_scale = large_scale_sinkhorn_balanced_benchmark()
     # result_batch_channel = batch_channel_sinkhorn_balanced_benchmark()
 
-    print(result_sinkhorn_balanced)
-    print(result_ot)
-    speedup = result_ot.mean / result_sinkhorn_balanced.mean
-    print(f"My sinkhorn to OT sinkhorn Speedup: {speedup:.1f}x")
+    # print(result_sinkhorn_balanced)
+    # print(result_ot)
+    # speedup = result_ot.mean / result_sinkhorn_balanced.mean
+    # print(f"My sinkhorn to OT sinkhorn Speedup: {speedup:.1f}x")
     # print(result_large_scale)
     # print(result_batch_channel)
+
+    print_size_benchmark = False
+    print_lambda_benchmark = True
+
+    if print_size_benchmark:
+        # sizes = [16, 32, 64, 100]
+        sizes = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
+        result_sinkhorn_balanced = []
+        result_ot = []
+        result_lambda = []
+        result_relative_error = []
+
+        for size in sizes:
+            sinkhorn_result = run_sinkhorn_balanced_benchmark(B=1, C=1, H=size, W=size)
+            ot_result = run_ot_sinkhorn_benchmark(H=size, W=size)
+            lambda_result = run_sinkhorn_balanced_benchmark(B=1, C=1, H=size, W=size, myreg=0.01)
+            relative_error = run_sinkhorn_relative_error(H=size, W=size)
+            result_sinkhorn_balanced.append(sinkhorn_result)
+            result_ot.append(ot_result)
+            result_lambda.append(lambda_result)
+            result_relative_error.append(relative_error)
+
+        print(
+            "size,my_sinkhorn_ms,ot_sinkhorn_ms,lambda_100_sinkhorn_ms,"
+            "ot_speedup,lambda_100_vs_lambda_50,plan_relative_error"
+        )
+        for size, sinkhorn_result, ot_result, lambda_result, relative_error in zip(
+            sizes, result_sinkhorn_balanced, result_ot, result_lambda, result_relative_error
+        ):
+            speedup = ot_result.mean / sinkhorn_result.mean
+            lambda_vs_default = lambda_result.mean / sinkhorn_result.mean
+            print(
+                f"{size}x{size},"
+                f"{sinkhorn_result.mean * 1000:.4f},"
+                f"{ot_result.mean * 1000:.4f},"
+                f"{lambda_result.mean * 1000:.4f},"
+                f"{speedup:.2f}x,"
+                f"{lambda_vs_default:.2f}x,"
+                f"{relative_error.item():.4e}"
+            )
+
+    if print_lambda_benchmark:
+        H, W = 64, 64
+        lambdas = [1, 1.5, 2, 5, 10, 25, 50, 100]
+        result_sinkhorn_balanced = []
+
+        for lambda_value in lambdas:
+            myreg = 1 / lambda_value
+            sinkhorn_result = run_sinkhorn_balanced_benchmark(
+                B=1, C=1, H=H, W=W, myreg=myreg, threshold=1e-5
+            )
+            result_sinkhorn_balanced.append(sinkhorn_result)
+
+        lambda_50_index = lambdas.index(50)
+        lambda_50_mean = result_sinkhorn_balanced[lambda_50_index].mean
+
+        print("lambda,myreg,size,my_sinkhorn_ms,vs_lambda_50")
+        for lambda_value, sinkhorn_result in zip(lambdas, result_sinkhorn_balanced):
+            vs_lambda_50 = sinkhorn_result.mean / lambda_50_mean
+            print(
+                f"{lambda_value},"
+                f"{1 / lambda_value:.4f},"
+                f"{H}x{W},"
+                f"{sinkhorn_result.mean * 1000:.4f},"
+                f"{vs_lambda_50:.2f}x"
+            )
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index c309d2c..43b2f4b 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -119,6 +119,48 @@ def test_sinkhorn_balanced_print():
         traceback.print_exc()
 
 
+def test_sinkhorn_cuda():
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA not available, skipping CUDA test.")
+
+    device = 'cuda'
+    # Create two random distributions
+    torch.manual_seed(42)
+    source = torch.rand(2, 2, device=device)
+    target = torch.rand(2, 2, device=device)
+
+    print("\n--- Input source ---")
+    print(source)
+    print("\n--- Input target ---")
+    print(target)
+
+    source, target, cost_matrix = tr.data_preprocess(source, target, device=device)
+
+    torch_result = tr.sinkhorn_balanced(
+        source=source,
+        target=target,
+        cost_matrix=cost_matrix,
+        reg=0.02,
+        itrstep=100,
+    )
+
+    cuda_result = tr.sinkhorn_balanced_cuda(
+        source=source,
+        target=target,
+        cost_matrix=cost_matrix,
+        reg=0.02,
+        itrstep=100,
+    )
+
+    relative_error = torch.linalg.norm(cuda_result["plan"] - torch_result["plan"]) / torch.clamp(
+        torch.linalg.norm(torch_result["plan"]), min=1e-12
+    )
+
+    print(f"CUDA relative error: {relative_error.item():.4e}")
+    assert relative_error < 1e-4
+
+
 if __name__ == "__main__":
     test_sinkhorn_balanced_print()
     test_build_cost_matrix()
+    test_sinkhorn_cuda()
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
index 839b5ce..1ad28c6 100644
--- a/torchmorph/csrc/ot_kernel.cu
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -86,7 +86,19 @@ __global__ void v_update(
     }
 }
 
-std::tuple<torch::Tensor, torch::Tensor> sinkhorn_uv_cuda(
+__global__ void sinkhorn_large(
+    const float* a,
+    const float* b,
+    const float* K,
+    float* u,
+    float* v,
+    int N
+){
+    //sinkhorn for large scale matrix
+}
+
+
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
     const torch::Tensor source,
     const torch::Tensor target,
     const torch::Tensor k,
diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp
index 5241510..9e9e2fb 100644
--- a/torchmorph/csrc/torchmorph.cpp
+++ b/torchmorph/csrc/torchmorph.cpp
@@ -28,7 +28,7 @@ std::tuple<torch::Tensor, torch::Tensor> bfdt_cuda(
 );
 
 // Optimal Transport functions
-std::tuple<torch::Tensor, torch::Tensor> sinkhorn_uv_cuda(
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
     const torch::Tensor source,
     const torch::Tensor target,
     const torch::Tensor k,
@@ -67,8 +67,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           py::arg("return_indices") = false);
 
     // Optimal Transport
-    m.def("sinkhorn_uv_cuda", &sinkhorn_uv_cuda,
-          "Sinkhorn UV CUDA",
+    m.def("sinkhorn_fastiter", &sinkhorn_fastiter,
+          "Sinkhorn Fast Iteration CUDA",
           py::arg("source"),
           py::arg("target"),
           py::arg("k"),
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 60c7810..070e39c 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -23,18 +23,29 @@ def data_preprocess(
     target: Tensor,
     cost_matrix: Optional[Tensor] = None,
     p: int = 2,
+    force_batched: bool = False,
     device: str = 'cuda',
+    dtype=torch.float32,
 ):
+    """Validate, move, reshape, and normalize Sinkhorn input distributions.
+
+    If batch_flag is True, 1D/2D inputs are reshaped to (1, 1, N).
+    Higher-dimensional inputs are reshaped from (B, C, *Spatial) to (B, C, N).
+    """
     # check dimensions,reshape and choose strategy
     if not source.ndim == target.ndim or not source.shape == target.shape:
         raise ValueError(
             "Source and target must have the same number of dimensions and the same shape."
         )
-    elif source.ndim == 2:
+    elif source.ndim <= 2:
         if cost_matrix is None:
             cost_matrix = build_cost_matrix(source.shape, device=device, p=p)
-        source = source.view(1, 1, -1)
-        target = target.view(1, 1, -1)
+        if force_batched:
+            source = source.view(1, 1, -1)
+            target = target.view(1, 1, -1)
+        else:
+            source = source.view(-1)
+            target = target.view(-1)
     elif source.ndim > 2:
         # cut and reshape source and target to coordinate tensors, then calculate cost matrix
         B, C = source.shape[:2]  # get batch and channel dimensions
@@ -45,9 +56,9 @@ def data_preprocess(
         target = target.view(B, C, -1)
 
     # move to specified device(cuda) and ensure dtype float32
-    source = source.to(device=device, dtype=torch.float32)
-    target = target.to(device=device, dtype=torch.float32)
-    cost_matrix = cost_matrix.to(device=device, dtype=torch.float32)
+    source = source.to(device=device, dtype=dtype)
+    target = target.to(device=device, dtype=dtype)
+    cost_matrix = cost_matrix.to(device=device, dtype=dtype)
 
     # ensure non-negativity
     source = torch.clamp(source, min=0)
@@ -78,8 +89,9 @@ def sinkhorn_balanced_full(
         cost_matrix=cost_matrix,
         p=p,
         device=device,
+        force_batched=True,
     )
-    result = sinkhorn_balanced(
+    result = sinkhorn_balanced_batch(
         source,
         target,
         cost_matrix=cost_matrix,
@@ -96,8 +108,7 @@ def sinkhorn_balanced_full(
     return result["plan"]
 
 
-# sinkhorn iteration
-def sinkhorn_balanced(
+def sinkhorn_balanced_batch(
     source: Tensor,  # (B,C,*Spatial)
     target: Tensor,  # (B,C,*Spatial)
     cost_matrix: Tensor,
@@ -110,28 +121,20 @@ def sinkhorn_balanced(
 ):
     """Compute a balanced entropic optimal transport plan with Sinkhorn iterations.
 
-    Uses the Gibbs kernel K = exp(-reg * C), where `reg` is the inverse
-    regularization scale. Source and target are clamped to be non-negative,
-    normalized over the flattened spatial dimension, and matched by iterative
-    Sinkhorn scaling.
+    Builds K = exp(-reg * C), iteratively updates scaling vectors u and v,
+    and returns the transport plan P.
 
     Args:
-        source (Tensor): Source distribution, shape (B, C, *Spatial) or (*Spatial).
-        target (Tensor): Target distribution with the same shape as source.
-        cost_matrix (Tensor, optional): Pairwise cost matrix of shape (N, N).
-        p (int): Lp norm used to build the cost matrix when not provided.
-        reg (float): Positive inverse regularization parameter.
+        source (Tensor): Preprocessed source distribution with shape (B, C, N).
+        target (Tensor): Preprocessed target distribution with shape (B, C, N).
+        cost_matrix (Tensor): Pairwise cost matrix with shape (N, N).
+        reg (float): Inverse regularization parameter.
         itrstep (int): Number of Sinkhorn iterations.
-        threshold (float): Stopping threshold on the change of the source scaling.
-        returngrad (bool): If True, also return scaling vectors u and v.
-        device (str): Computation device, typically 'cuda'.
+        threshold (float): Optional convergence threshold.
+        returnuv (bool): If True, also return u and v.
 
     Returns:
-        dict: Contains "plan" and, when requested, "u" and "v".
-
-    Raises:
-        ValueError: If dimensions, regularization, device, or stopping settings
-            are invalid.
+        dict: Transport plan, and optionally scaling vectors u and v.
     """
     # Device check
     if not torch.cuda.is_available():
@@ -152,20 +155,22 @@ def sinkhorn_balanced(
         raise ValueError("Regularization parameter must be positive.")
 
     # generate sinkhorn iteration judged by threshold or iteration times
-    if itrstep == 0 and threshold > 0:
-        step = 1000
+    if itrstep < 0 or threshold < 0:
+        raise ValueError("iterstep and threshold must be non-negative.")
     elif itrstep == 0 and threshold == 0:
         raise ValueError("Invalid iteration or step settings. ")
-    else:
-        step = itrstep
+
+    if threshold > 0 and itrstep == 0:
+        itrstep = 500
 
     # iteration
     convergence_flag = False
-    for i in range(step):
-        u_old = u.clone()
+    convergence_check = 10
+    for i in range(itrstep):
+        u_old = u
         u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
         v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
-        if (u_old - u).abs().sum().mean() < threshold:
+        if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < threshold:
             convergence_flag = True
             break
 
@@ -185,6 +190,54 @@ def sinkhorn_balanced(
         return {"plan": P}
 
 
+def sinkhorn_balanced(
+    source: Tensor,  # (N,)
+    target: Tensor,  # (N,)
+    cost_matrix: Tensor,
+    reg: float = 1.0,
+    itrstep: int = 0,
+    threshold: float = 0,
+    returnuv: bool = False,  # whether to return the gradient of the dual potential
+    device: str = 'cuda',
+):
+    device = device if torch.cuda.is_available() else 'cpu'
+
+    # initialize u and v, and expand dimensions for batch and channel
+    u = torch.ones_like(source, device=device)
+    v = torch.ones_like(target, device=device)
+    k = torch.exp(-cost_matrix * reg)
+
+    if reg <= 0:
+        raise ValueError("Regularization parameter must be positive.")
+
+    # generate sinkhorn iteration judged by threshold or iteration times
+    if itrstep < 0 or threshold < 0:
+        raise ValueError("iterstep and threshold must be non-negative.")
+    elif itrstep == 0 and threshold == 0:
+        raise ValueError("Invalid iteration or step settings. ")
+
+    if threshold > 0 and itrstep == 0:
+        itrstep = 500
+
+    # iteration
+    convergence_check = 10
+    for i in range(itrstep):
+        u_old = u
+        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
+        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
+        if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < threshold:
+            break
+
+    # use broadcasting to calculate the optimal transport plan P from u, v and K
+    P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
+
+    # calculate gradient after iteration covergence ,return results
+    if returnuv:
+        return {"plan": P, "u": u, "v": v}
+    else:
+        return {"plan": P}
+
+
 def sinkhorn_gradient(u: Tensor, v: Tensor, reg: float):
     u = torch.log(torch.clamp(u, min=1e-12))
     v = torch.log(torch.clamp(v, min=1e-12))
@@ -213,7 +266,7 @@ def sinkhorn_balanced_cuda(
     u = torch.ones_like(source)
     v = torch.ones_like(target)
 
-    u, v = _C.sinkhorn_uv_cuda(source, target, K, u, v, int(itrstep), int(N))
+    u, v = _C.sinkhorn_fastiter(source, target, K, u, v, int(itrstep), int(N))
 
     P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
 

From c631ecf6338e902ef00dc77229c1fdfa5a79efe3 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Thu, 14 May 2026 15:11:09 +0800
Subject: [PATCH 08/20] Add CUDA support for Sinkhorn log-domain iterations and
 related tests.

---
 benchmark/optimal_transport.py  |  27 +----
 test/test_optimal_transport.py  |  47 +++++++-
 torchmorph/csrc/ot_kernel.cu    | 185 ++++++++++++++++++++++++++++++--
 torchmorph/csrc/torchmorph.cpp  |  25 ++++-
 torchmorph/optimal_transport.py |  28 ++++-
 5 files changed, 274 insertions(+), 38 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index c2a97b2..fd93035 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -26,12 +26,11 @@ def run_sinkhorn_balanced_benchmark(
 ):
     # Set parameters
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype = torch.float32
     torch.manual_seed(42)
 
     # Prepare input data
-    source = torch.rand(B, C, H, W, device=device, dtype=dtype)
-    target = torch.rand(B, C, H, W, device=device, dtype=dtype)
+    source = torch.rand(B, C, H, W, device=device, dtype=torch.float32)
+    target = torch.rand(B, C, H, W, device=device, dtype=torch.float32)
     source = source / source.sum()
     target = target / target.sum()
     cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
@@ -192,28 +191,12 @@ def batch_channel_sinkhorn_balanced_benchmark():
 
 
 if __name__ == "__main__":
-    # sizes=[10,20,50,100]
-    # for size in sizes:
-    #     i = 0
-    #     result_sinkhorn_balanced[] = run_sinkhorn_balanced_benchmark(B=1, C=1, H=size, W=size)
-    #     result_ot[] = run_ot_sinkhorn_benchmark(H=size, W=size)
-    #     i += 1
-    # result_large_scale = large_scale_sinkhorn_balanced_benchmark()
-    # result_batch_channel = batch_channel_sinkhorn_balanced_benchmark()
-
-    # print(result_sinkhorn_balanced)
-    # print(result_ot)
-    # speedup = result_ot.mean / result_sinkhorn_balanced.mean
-    # print(f"My sinkhorn to OT sinkhorn Speedup: {speedup:.1f}x")
-    # print(result_large_scale)
-    # print(result_batch_channel)
-
-    print_size_benchmark = False
-    print_lambda_benchmark = True
+    print_size_benchmark = True
+    print_lambda_benchmark = False
 
     if print_size_benchmark:
         # sizes = [16, 32, 64, 100]
-        sizes = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
+        sizes = [100]
         result_sinkhorn_balanced = []
         result_ot = []
         result_lambda = []
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 43b2f4b..9e90ffe 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -160,7 +160,48 @@ def test_sinkhorn_cuda():
     assert relative_error < 1e-4
 
 
+def test_sinkhorn_log():
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA not available, skipping CUDA test.")
+
+    device = 'cuda'
+    # Create two random distributions
+    torch.manual_seed(42)
+    source = torch.rand(2, 2, device=device)
+    target = torch.rand(2, 2, device=device)
+
+    print("\n--- Input source ---")
+    print(source)
+    print("\n--- Input target ---")
+    print(target)
+
+    source, target, cost_matrix = tr.data_preprocess(source, target, device=device)
+
+    u, v = tr.sinkhorn_balanced_log(
+        source=source,
+        target=target,
+        cost_matrix=cost_matrix,
+        reg=0.02,
+        itrstep=100,
+    )
+
+    print("\n--- Output u (log-domain) ---")
+    print(u)
+    print("\n--- Output v (log-domain) ---")
+    print(v)
+
+    return u, v
+
+
+def test_allclose(u1, v1, u2, v2):
+    u_close = torch.allclose(u1, u2, atol=1e-4)
+    v_close = torch.allclose(v1, v2, atol=1e-4)
+    print(f"u close: {u_close}, v close: {v_close}")
+    assert u_close and v_close, "u or v not close enough"
+
+
 if __name__ == "__main__":
-    test_sinkhorn_balanced_print()
-    test_build_cost_matrix()
-    test_sinkhorn_cuda()
+    # test_sinkhorn_balanced_print()
+    # test_build_cost_matrix()
+    test_sinkhorn_log()
+    # test_sinkhorn_cuda()
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
index 1ad28c6..66a57ab 100644
--- a/torchmorph/csrc/ot_kernel.cu
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -3,6 +3,8 @@
 #include <cuda_runtime.h>
 #include <tuple>
 #include <cmath>
+#include <cfloat>
+
 
 __global__ void u_update(
     const float* a,
@@ -45,6 +47,7 @@ __global__ void u_update(
     }
 }
 
+
 __global__ void v_update(
     const float* a,
     const float* b,
@@ -76,7 +79,7 @@ __global__ void v_update(
     
     for (int stride=blk/2; stride > 0; stride /= 2){
         if (tid < stride){
-        sdata[tid]+=sdata[tid+stride];
+            sdata[tid]+=sdata[tid+stride];
         }
         __syncthreads();
     }
@@ -86,17 +89,144 @@ __global__ void v_update(
     }
 }
 
-__global__ void sinkhorn_large(
-    const float* a,
-    const float* b,
-    const float* K,
-    float* u,
-    float* v,
-    int N
+
+// Update one log-domain Sinkhorn scaling vector.
+//
+// Each block owns one row i and computes:
+//   log_u[i] = log_a[i] - logsumexp_j(-reg * M[i, j] + log_v[j])
+// The logsumexp is evaluated with a two-pass block reduction:
+// first max(z_j), then sum(exp(z_j - max)).
+__global__ void u_update_log(
+    const float* log_a,
+    const float* M,
+    float* log_u,
+    const float* log_v,
+    int N,
+    float reg
 ){
-    //sinkhorn for large scale matrix
+    //sinkhorn for log domain iteration
+    int blk = blockDim.x;
+    int bidx = blockIdx.x;
+    int tid = threadIdx.x;
+    int total = N;
+    float local_max = -INFINITY;
+
+    if(bidx >= total){
+        return;
+    }
+
+    extern __shared__ float sdata[];
+
+    //generate each log_kv
+    for (int j = tid; j < N; j += blk){
+        float log_k = - reg * M[bidx*N+j];
+        if(local_max <= (log_k + log_v[j])){
+            local_max = log_k + log_v[j];
+        }     
+    }
+
+    sdata[tid] = local_max;
+    __syncthreads();
+
+    //first reduction to find max(log_k+log_v)
+    for (int stride=blk/2; stride > 0; stride /= 2){
+        if (tid < stride){
+            sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);//get larger num
+        }
+        __syncthreads();
+    }
+
+    float row_max = sdata[0];
+    __syncthreads();
+    float local_sum = 0.0f;
+
+    //log-domain iteration log_u = log_a - LSE(log_k+log_v)
+    for (int j = tid; j < N; j += blk){
+        float log_k = - reg * M[bidx*N+j];
+        local_sum += expf(log_k + log_v[j] - row_max);
+    }
+
+    sdata[tid] = local_sum;
+    __syncthreads();
+
+    for (int stride=blk/2; stride > 0; stride /= 2){
+        if (tid < stride){
+            sdata[tid] += sdata[tid + stride];
+        }
+        __syncthreads();
+    }
+
+    if (tid == 0){
+        log_u[bidx] = log_a[bidx] - row_max - logf(sdata[0]);
+    }
+
 }
 
+__global__ void v_update_log(
+    const float* log_b,
+    const float* M,
+    const float* log_u,
+    float* log_v,
+    int N,
+    float reg
+){
+    //sinkhorn for log domain iteration
+    int blk = blockDim.x;
+    int bidx = blockIdx.x;
+    int tid = threadIdx.x;
+    int total = N;
+    float local_max = -INFINITY;
+
+    if(bidx >= total){
+        return;
+    }
+
+    extern __shared__ float sdata[];
+
+    //generate each log_kTu
+    for (int j = tid; j < N; j += blk){
+        float log_k = - reg * M[bidx+j*N];
+        if(local_max <= (log_k + log_u[j])){
+            local_max = log_k + log_u[j];
+        }     
+    }
+
+    sdata[tid] = local_max;
+    __syncthreads();
+
+    //first reduction to find max(log_k+log_v)
+    for (int stride=blk/2; stride > 0; stride /= 2){
+        if (tid < stride){
+            sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);//get larger num
+        }
+        __syncthreads();
+    }
+
+    float row_max = sdata[0];
+    __syncthreads();
+    float local_sum = 0.0f;
+
+    //log-domain iteration log_v = log_b - LSE(log_k+log_u)
+    for (int j = tid; j < N; j += blk){
+        float log_k = - reg * M[bidx+j*N];
+        local_sum += expf(log_k + log_u[j] - row_max);
+    }
+
+    sdata[tid] = local_sum;
+    __syncthreads();
+
+    for (int stride=blk/2; stride > 0; stride /= 2){
+        if (tid < stride){
+            sdata[tid] += sdata[tid + stride];
+        }
+        __syncthreads();
+    }
+
+    if (tid == 0){
+        log_v[bidx] = log_b[bidx] - row_max - logf(sdata[0]);
+    }
+
+}
 
 std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
     const torch::Tensor source,
@@ -130,4 +260,39 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
         );
     }
 return std::make_tuple(u, v);
-}
\ No newline at end of file
+}
+
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
+    const torch::Tensor log_a,
+    const torch::Tensor log_b,
+    const torch::Tensor M,
+    torch::Tensor log_u,
+    torch::Tensor log_v,
+    int itrstep,
+    int N,
+    float reg
+){
+    int threads = 256;
+    int blocks = N;
+    int shared_memory = threads * sizeof(float);
+
+    for (int i = 0; i < itrstep; i++){
+        u_update_log<<<blocks, threads, shared_memory>>>(
+            log_a.data_ptr<float>(),
+            M.data_ptr<float>(),
+            log_u.data_ptr<float>(),
+            log_v.data_ptr<float>(),
+            N,
+            reg
+        );
+        v_update_log<<<blocks, threads, shared_memory>>>(
+            log_b.data_ptr<float>(),
+            M.data_ptr<float>(),
+            log_u.data_ptr<float>(),
+            log_v.data_ptr<float>(),
+            N,
+            reg
+        );
+    }
+return std::make_tuple(log_u, log_v);
+}
diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp
index 9e9e2fb..0e507a8 100644
--- a/torchmorph/csrc/torchmorph.cpp
+++ b/torchmorph/csrc/torchmorph.cpp
@@ -38,6 +38,18 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
     int N
 );
 
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
+    const torch::Tensor log_a,
+    const torch::Tensor log_b,
+    const torch::Tensor M,
+    torch::Tensor log_u,
+    torch::Tensor log_v,
+    int itrstep,
+    int N,
+    float reg
+);
+
+
 
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
@@ -76,6 +88,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           py::arg("v"),
           py::arg("itrstep"),
           py::arg("N"));
-}
 
+    m.def("sinkhorn_logiter", &sinkhorn_logiter,
+          "Sinkhorn Log-Domain Iteration CUDA",
+          py::arg("log_a"),
+          py::arg("log_b"),
+          py::arg("M"),
+          py::arg("log_u"),
+          py::arg("log_v"),
+          py::arg("itrstep"),
+          py::arg("N"),
+          py::arg("reg")
+    );
+}
 
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 070e39c..4dcff11 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -71,7 +71,7 @@ def data_preprocess(
     return source, target, cost_matrix
 
 
-def sinkhorn_balanced_full(
+def sinkhorn_balanced(
     source: Tensor,
     target: Tensor,
     cost_matrix: Optional[Tensor] = None,
@@ -190,7 +190,7 @@ def sinkhorn_balanced_batch(
         return {"plan": P}
 
 
-def sinkhorn_balanced(
+def sinkhorn_balanced_nobatch(
     source: Tensor,  # (N,)
     target: Tensor,  # (N,)
     cost_matrix: Tensor,
@@ -274,3 +274,27 @@ def sinkhorn_balanced_cuda(
         return {"plan": P, "u": u, "v": v}
     else:
         return {"plan": P}
+
+
+def sinkhorn_balanced_log(
+    source: Tensor,
+    target: Tensor,
+    cost_matrix: Tensor,
+    reg: float = 1.0,
+    itrstep: int = 100,
+    dtype=torch.float32,
+):
+    from torchmorph import _C
+
+    source = source.to(dtype=dtype).log()
+    target = target.to(dtype=dtype).log()
+    N = source.shape[-1]
+    u = torch.zeros_like(source)
+    v = torch.zeros_like(target)
+
+    u, v = _C.sinkhorn_logiter(source, target, cost_matrix, u, v, int(itrstep), int(N), float(reg))
+
+    u = u.exp()
+    v = v.exp()
+
+    return u, v

From 25f6b5d8d02322ad3478f7cf161b8949f9b1e174 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Fri, 22 May 2026 16:09:48 +0800
Subject: [PATCH 09/20] Refactor Sinkhorn functions and test gradient accuracy

---
 benchmark/optimal_transport.py  | 340 ++++++++++++++-----
 test/test_optimal_transport.py  | 101 ++++--
 torchmorph/csrc/ot_kernel.cu    |   9 +-
 torchmorph/csrc/torchmorph.cpp  |   2 +
 torchmorph/optimal_transport.py | 577 ++++++++++++++++----------------
 5 files changed, 624 insertions(+), 405 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index fd93035..44fd210 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -21,9 +21,10 @@ def run_sinkhorn_balanced_benchmark(
     H: int = 32,
     W: int = 32,
     myreg: float = 0.02,
-    itrstep: int = 0,
-    threshold: float = 1e-5,
+    itrstep: int = 100,
+    threshold: float = 0,
 ):
+    torch.cuda.reset_peak_memory_stats()
     # Set parameters
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.manual_seed(42)
@@ -35,34 +36,41 @@ def run_sinkhorn_balanced_benchmark(
     target = target / target.sum()
     cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
 
-    source, target, cost_matrix = tr.data_preprocess(
-        source, target, cost_matrix=cost_matrix, p=2, device=device
+    solver = tr.SinkhornSolver(
+        reg=myreg,
+        itrstep=itrstep,
+        threshold=threshold,
+        p=2,
+        device=device,
     )
+    source, target, cost_matrix, _ = solver.data_preprocess(source, target, cost_matrix=cost_matrix)
 
     globals_dict = {
-        "tr": tr,
+        "solver": solver,
         "source": source,
         "target": target,
-        "myreg": myreg,
         "cost_matrix": cost_matrix,
-        "itrstep": itrstep,
-        "threshold": threshold,
     }
 
-    stmt = (
-        "tr.sinkhorn_balanced_batch(source, target, cost_matrix=cost_matrix, "
-        "reg=myreg, itrstep=itrstep, threshold=threshold)"
-    )
+    stmt = "solver.sinkhorn_batch(source, target, cost_matrix=cost_matrix)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
     # number_of_runs = 100
-    result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=2)
+    result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=3)
+
+    peak_allocated = torch.cuda.max_memory_allocated()
+    peak_reserved = torch.cuda.max_memory_reserved()
+
+    print(f"peak allocated: {peak_allocated / 1024**2:.2f} MB")
+    print(f"peak reserved:   {peak_reserved / 1024**2:.2f} MB")
+
     return result_sinkhorn_balanced
 
 
 def run_ot_sinkhorn_benchmark(H: int = 32, W: int = 32, myreg: float = 0.02):
     # Set parameters
+    import ot
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     dtype = torch.float32
@@ -106,18 +114,22 @@ def run_sinkhorn_relative_error(H: int = 32, W: int = 32, myreg: float = 0.02):
     target = target / target.sum()
     cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
 
-    source_preprocessed, target_preprocessed, cost_matrix = tr.data_preprocess(
-        source, target, cost_matrix=cost_matrix, p=2, device=device
+    solver = tr.SinkhornSolver(
+        reg=myreg,
+        itrstep=100,
+        threshold=1e-5,
+        p=2,
+        device=device,
+    )
+    source_preprocessed, target_preprocessed, cost_matrix, _ = solver.data_preprocess(
+        source, target, cost_matrix=cost_matrix
     )
 
     sinkhorn_plan = (
-        tr.sinkhorn_balanced_batch(
+        solver.sinkhorn_batch(
             source_preprocessed,
             target_preprocessed,
             cost_matrix=cost_matrix,
-            reg=myreg,
-            itrstep=100,
-            threshold=1e-5,
         )["plan"]
         .squeeze(0)
         .squeeze(0)
@@ -147,14 +159,20 @@ def large_scale_sinkhorn_balanced_benchmark(B: int = 1, C: int = 1, H: int = 128
     target = torch.rand(B, C, H, W, device=device, dtype=dtype)
     source = source / source.sum()
     target = target / target.sum()
+    solver = tr.SinkhornSolver(
+        reg=25.0,
+        itrstep=100,
+        threshold=1e-5,
+        device=device,
+    )
 
     globals_dict = {
-        "tr": tr,
+        "solver": solver,
         "source": source,
         "target": target,
     }
 
-    stmt = "tr.sinkhorn_balanced_full(source, target, reg=25.0, itrstep=100, threshold=1e-5)"
+    stmt = "solver.sinkhorn(source, target)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
@@ -174,14 +192,20 @@ def batch_channel_sinkhorn_balanced_benchmark():
     target = torch.rand(B, C, H, W, device=device, dtype=dtype)
     source = source / source.sum()
     target = target / target.sum()
+    solver = tr.SinkhornSolver(
+        reg=25.0,
+        itrstep=100,
+        threshold=1e-5,
+        device=device,
+    )
 
     globals_dict = {
-        "tr": tr,
+        "solver": solver,
         "source": source,
         "target": target,
     }
 
-    stmt = "tr.sinkhorn_balanced_full(source, target, reg=25.0, itrstep=100, threshold=1e-5)"
+    stmt = "solver.sinkhorn(source, target)"
 
     timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
     # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
@@ -190,69 +214,213 @@ def batch_channel_sinkhorn_balanced_benchmark():
     return result_sinkhorn_balanced
 
 
-if __name__ == "__main__":
-    print_size_benchmark = True
-    print_lambda_benchmark = False
-
-    if print_size_benchmark:
-        # sizes = [16, 32, 64, 100]
-        sizes = [100]
-        result_sinkhorn_balanced = []
-        result_ot = []
-        result_lambda = []
-        result_relative_error = []
-
-        for size in sizes:
-            sinkhorn_result = run_sinkhorn_balanced_benchmark(B=1, C=1, H=size, W=size)
-            ot_result = run_ot_sinkhorn_benchmark(H=size, W=size)
-            lambda_result = run_sinkhorn_balanced_benchmark(B=1, C=1, H=size, W=size, myreg=0.01)
-            relative_error = run_sinkhorn_relative_error(H=size, W=size)
-            result_sinkhorn_balanced.append(sinkhorn_result)
-            result_ot.append(ot_result)
-            result_lambda.append(lambda_result)
-            result_relative_error.append(relative_error)
+def run_sinkhorn_log_benchmark(
+    H: int = 32,
+    W: int = 32,
+    myreg: float = 0.02,
+    itrstep: int = 100,
+):
+    torch.cuda.reset_peak_memory_stats()
+    # Set parameters
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA device not available.")
 
+    device = "cuda"
+    torch.manual_seed(42)
+
+    # Prepare input data
+    source = torch.rand(H, W, device=device, dtype=torch.float32)
+    target = torch.rand(H, W, device=device, dtype=torch.float32)
+    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
+
+    solver = tr.SinkhornSolver(reg=myreg, itrstep=itrstep, p=2, device=device)
+    source, target, cost_matrix, cost_matrix_T = solver.data_preprocess(
+        source, target, cost_matrix=cost_matrix
+    )
+
+    globals_dict = {
+        "solver": solver,
+        "source": source,
+        "target": target,
+        "cost_matrix": cost_matrix,
+        "cost_matrix_T": cost_matrix_T,
+    }
+
+    stmt = "solver.sinkhorn_log_cuda(source, target, cost_matrix, cost_matrix_T)"
+
+    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
+    result_sinkhorn_log = timer.blocked_autorange(min_run_time=3)
+
+    peak_allocated = torch.cuda.max_memory_allocated()
+    peak_reserved = torch.cuda.max_memory_reserved()
+
+    print(f"peak allocated: {peak_allocated / 1024**2:.2f} MB")
+    print(f"peak reserved:   {peak_reserved / 1024**2:.2f} MB")
+
+    return result_sinkhorn_log
+
+
+def run_sinkhorn_cuda_benchmark(
+    H: int = 32,
+    W: int = 32,
+    myreg: float = 0.02,
+    itrstep: int = 100,
+):
+    torch.cuda.reset_peak_memory_stats()
+    # Set parameters
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA device not available.")
+
+    device = "cuda"
+    torch.manual_seed(42)
+
+    # Prepare input data
+    source = torch.rand(H, W, device=device, dtype=torch.float32)
+    target = torch.rand(H, W, device=device, dtype=torch.float32)
+    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
+
+    solver = tr.SinkhornSolver(reg=myreg, itrstep=itrstep, p=2, device=device)
+    source, target, cost_matrix, _ = solver.data_preprocess(source, target, cost_matrix=cost_matrix)
+
+    globals_dict = {
+        "solver": solver,
+        "source": source,
+        "target": target,
+        "cost_matrix": cost_matrix,
+    }
+
+    stmt = "solver.sinkhorn_cuda(source, target, cost_matrix=cost_matrix)"
+
+    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
+    result_sinkhorn_cuda = timer.blocked_autorange(min_run_time=3)
+
+    peak_allocated = torch.cuda.max_memory_allocated()
+    peak_reserved = torch.cuda.max_memory_reserved()
+
+    print(f"peak allocated: {peak_allocated / 1024**2:.2f} MB")
+    print(f"peak reserved:   {peak_reserved / 1024**2:.2f} MB")
+
+    return result_sinkhorn_cuda
+
+
+def reference_error_check(
+    size: int = 32, myreg: float = 0.1, itrstep: int = 200, verbose: bool = False
+):
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA device not available.")
+
+    device = "cuda"
+    torch.manual_seed(2)
+
+    # Prepare input data
+    source = torch.rand(size, size, device=device, dtype=torch.float32)
+    target = torch.rand(size, size, device=device, dtype=torch.float32)
+    cost_matrix = tr.build_cost_matrix((size, size), device=device, p=2)
+
+    solver = tr.SinkhornSolver(reg=myreg, itrstep=itrstep, p=2, device=device)
+    source, target, cost_matrix, cost_matrix_T = solver.data_preprocess(
+        source, target, cost_matrix=cost_matrix
+    )
+
+    _, log = ot.bregman.sinkhorn_log(
+        source,
+        target,
+        cost_matrix,
+        reg=1 / myreg,
+        numItermax=itrstep,
+        stopThr=0,
+        log=True,
+    )
+
+    u, v = solver.sinkhorn_log_cuda(
+        source,
+        target,
+        cost_matrix,
+        cost_matrix_T,
+    )
+
+    grad_f, grad_g = solver.gradient(u, v)
+
+    log_u = log["log_u"]
+    log_v = log["log_v"]
+
+    f = log_u / myreg
+    g = log_v / myreg
+    f = f - f.mean()
+    g = g - g.mean()
+
+    if verbose:
+        print(f)
+        print(g)
+        print(grad_f)
+        print(grad_g)
+
+    f_abs_error = (f - grad_f).abs()
+    g_abs_error = (g - grad_g).abs()
+    f_rel_l2 = torch.linalg.norm(f - grad_f) / torch.clamp(torch.linalg.norm(f), min=1e-12)
+    g_rel_l2 = torch.linalg.norm(g - grad_g) / torch.clamp(torch.linalg.norm(g), min=1e-12)
+
+    return {
+        "size": size,
+        "points": size * size,
+        "reg": myreg,
+        "itrstep": itrstep,
+        "f_grad_max_abs": f_abs_error.max().item(),
+        "f_grad_mean_abs": f_abs_error.mean().item(),
+        "f_grad_rel_l2": f_rel_l2.item(),
+        "g_grad_max_abs": g_abs_error.max().item(),
+        "g_grad_mean_abs": g_abs_error.mean().item(),
+        "g_grad_rel_l2": g_rel_l2.item(),
+        "f_allclose": torch.allclose(f, grad_f, rtol=1e-4, atol=1e-4),
+        "g_allclose": torch.allclose(g, grad_g, rtol=1e-4, atol=1e-4),
+    }
+
+
+def print_reference_error_table(
+    sizes=(2, 4, 8, 16, 32, 64),
+    itrsteps=(100, 200, 500),
+    myreg: float = 0.1,
+):
+    """Print a size and iteration sweep for POT-vs-CUDA gradient error."""
+    rows = [
+        reference_error_check(size=size, myreg=myreg, itrstep=itrstep)
+        for size in sizes
+        for itrstep in itrsteps
+    ]
+
+    headers = (
+        "size",
+        "N",
+        "reg",
+        "itr",
+        "f_max",
+        "f_mean",
+        "f_rel_l2",
+        "g_max",
+        "g_mean",
+        "g_rel_l2",
+        "close",
+    )
+    print("| " + " | ".join(headers) + " |")
+    print("| " + " | ".join(["---"] * len(headers)) + " |")
+    for row in rows:
+        close = f"{row['f_allclose']}/{row['g_allclose']}"
         print(
-            "size,my_sinkhorn_ms,ot_sinkhorn_ms,lambda_100_sinkhorn_ms,"
-            "ot_speedup,lambda_100_vs_lambda_50,plan_relative_error"
+            f"| {row['size']}x{row['size']} "
+            f"| {row['points']} "
+            f"| {row['reg']:.3g} "
+            f"| {row['itrstep']} "
+            f"| {row['f_grad_max_abs']:.3e} "
+            f"| {row['f_grad_mean_abs']:.3e} "
+            f"| {row['f_grad_rel_l2']:.3e} "
+            f"| {row['g_grad_max_abs']:.3e} "
+            f"| {row['g_grad_mean_abs']:.3e} "
+            f"| {row['g_grad_rel_l2']:.3e} "
+            f"| {close} |"
         )
-        for size, sinkhorn_result, ot_result, lambda_result, relative_error in zip(
-            sizes, result_sinkhorn_balanced, result_ot, result_lambda, result_relative_error
-        ):
-            speedup = ot_result.mean / sinkhorn_result.mean
-            lambda_vs_default = lambda_result.mean / sinkhorn_result.mean
-            print(
-                f"{size}x{size},"
-                f"{sinkhorn_result.mean * 1000:.4f},"
-                f"{ot_result.mean * 1000:.4f},"
-                f"{lambda_result.mean * 1000:.4f},"
-                f"{speedup:.2f}x,"
-                f"{lambda_vs_default:.2f}x,"
-                f"{relative_error.item():.4e}"
-            )
-
-    if print_lambda_benchmark:
-        H, W = 64, 64
-        lambdas = [1, 1.5, 2, 5, 10, 25, 50, 100]
-        result_sinkhorn_balanced = []
-
-        for lambda_value in lambdas:
-            myreg = 1 / lambda_value
-            sinkhorn_result = run_sinkhorn_balanced_benchmark(
-                B=1, C=1, H=H, W=W, myreg=myreg, threshold=1e-5
-            )
-            result_sinkhorn_balanced.append(sinkhorn_result)
-
-        lambda_50_index = lambdas.index(50)
-        lambda_50_mean = result_sinkhorn_balanced[lambda_50_index].mean
-
-        print("lambda,myreg,size,my_sinkhorn_ms,vs_lambda_50")
-        for lambda_value, sinkhorn_result in zip(lambdas, result_sinkhorn_balanced):
-            vs_lambda_50 = sinkhorn_result.mean / lambda_50_mean
-            print(
-                f"{lambda_value},"
-                f"{1 / lambda_value:.4f},"
-                f"{H}x{W},"
-                f"{sinkhorn_result.mean * 1000:.4f},"
-                f"{vs_lambda_50:.2f}x"
-            )
+
+    return rows
+
+
+if __name__ == "__main__":
+    print_reference_error_table()
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 9e90ffe..ec5d82c 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -95,17 +95,19 @@ def test_sinkhorn_balanced_print():
 
     print("\n--- Calling sinkhorn_balanced ---")
     try:
-        P = tr.sinkhorn_balanced_full(
-            source,
-            target,
-            cost_matrix=None,
+        solver = tr.SinkhornSolver(
             reg=0.1,
             itrstep=100,
             threshold=1e-5,
-            returngrad=False,
             device=device,
             p=2,
         )
+        P = solver.sinkhorn(
+            source,
+            target,
+            cost_matrix=None,
+            returngrad=False,
+        )
         print("\n--- Output P (transport plan) shape ---")
         print(P.shape)
         print("\n--- Output P content (partial) ---")
@@ -134,26 +136,25 @@ def test_sinkhorn_cuda():
     print("\n--- Input target ---")
     print(target)
 
-    source, target, cost_matrix = tr.data_preprocess(source, target, device=device)
+    solver = tr.SinkhornSolver(reg=0.02, itrstep=100, device=device)
+    source, target, cost_matrix, _ = solver.data_preprocess(source, target)
 
-    torch_result = tr.sinkhorn_balanced(
+    torch_result = solver.sinkhorn(
         source=source,
         target=target,
         cost_matrix=cost_matrix,
-        reg=0.02,
-        itrstep=100,
+        returngrad=True,
     )
 
-    cuda_result = tr.sinkhorn_balanced_cuda(
+    cuda_result = solver.sinkhorn_cuda(
         source=source,
         target=target,
         cost_matrix=cost_matrix,
-        reg=0.02,
-        itrstep=100,
     )
 
-    relative_error = torch.linalg.norm(cuda_result["plan"] - torch_result["plan"]) / torch.clamp(
-        torch.linalg.norm(torch_result["plan"]), min=1e-12
+    torch_plan = torch_result["plan"].squeeze(0).squeeze(0)
+    relative_error = torch.linalg.norm(cuda_result["plan"] - torch_plan) / torch.clamp(
+        torch.linalg.norm(torch_plan), min=1e-12
     )
 
     print(f"CUDA relative error: {relative_error.item():.4e}")
@@ -167,41 +168,77 @@ def test_sinkhorn_log():
     device = 'cuda'
     # Create two random distributions
     torch.manual_seed(42)
-    source = torch.rand(2, 2, device=device)
-    target = torch.rand(2, 2, device=device)
+    source = torch.rand(2, 2, device=device).abs()
+    target = torch.rand(2, 2, device=device).abs()
+
+    solver = tr.SinkhornSolver(reg=0.02, itrstep=100, device=device)
+    source, target, cost_matrix, cost_matrix_T = solver.data_preprocess(source, target)
 
     print("\n--- Input source ---")
     print(source)
     print("\n--- Input target ---")
     print(target)
 
-    source, target, cost_matrix = tr.data_preprocess(source, target, device=device)
-
-    u, v = tr.sinkhorn_balanced_log(
+    u, v = solver.sinkhorn_log_cuda(
         source=source,
         target=target,
         cost_matrix=cost_matrix,
-        reg=0.02,
-        itrstep=100,
+        cost_matrix_T=cost_matrix_T,
     )
 
-    print("\n--- Output u (log-domain) ---")
+    print("\n--- Output u  ---")
     print(u)
-    print("\n--- Output v (log-domain) ---")
+    print("\n--- Output v  ---")
     print(v)
 
-    return u, v
 
+def loss_test():
+    if not torch.cuda.is_available():
+        raise ValueError("CUDA not available, skipping loss test.")
+
+    device = "cuda"
+    torch.manual_seed(42)
+
+    reg = 5
+    itrstep = 100
+    solver = tr.SinkhornSolver(reg=reg, itrstep=itrstep, device=device)
+
+    source = torch.nn.Parameter(torch.randn(4, device=device).abs())
+    optimizer = torch.optim.Adam([source], lr=1e-2)
+
+    target = torch.rand(4, device=device).abs()
+
+    for step in range(500):
+        optimizer.zero_grad()
+
+        source_prob, target, cost_matrix, cost_matrix_T = solver.data_preprocess(source, target)
+
+        with torch.no_grad():
+            u, v = solver.sinkhorn_log_cuda(
+                source=source_prob.detach(),
+                target=target.detach(),
+                cost_matrix=cost_matrix,
+                cost_matrix_T=cost_matrix_T,
+            )
+
+            K = torch.exp(-cost_matrix * reg)
+            P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
+            loss_value = (P * cost_matrix).sum()
+
+            grad_source, _ = solver.gradient(u, v)
+
+        source_prob.backward(grad_source)
+
+        optimizer.step()
 
-def test_allclose(u1, v1, u2, v2):
-    u_close = torch.allclose(u1, u2, atol=1e-4)
-    v_close = torch.allclose(v1, v2, atol=1e-4)
-    print(f"u close: {u_close}, v close: {v_close}")
-    assert u_close and v_close, "u or v not close enough"
+        print(f"step {step}, loss = {loss_value.item():.6f}")
+        print("source:", source.detach())
+        print("target:", target)
 
 
 if __name__ == "__main__":
-    # test_sinkhorn_balanced_print()
-    # test_build_cost_matrix()
+    test_sinkhorn_balanced_print()
+    test_build_cost_matrix()
     test_sinkhorn_log()
-    # test_sinkhorn_cuda()
+    test_sinkhorn_cuda()
+    loss_test()
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
index 66a57ab..13c4ec5 100644
--- a/torchmorph/csrc/ot_kernel.cu
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -164,7 +164,7 @@ __global__ void u_update_log(
 
 __global__ void v_update_log(
     const float* log_b,
-    const float* M,
+    const float* M_T,
     const float* log_u,
     float* log_v,
     int N,
@@ -185,7 +185,7 @@ __global__ void v_update_log(
 
     //generate each log_kTu
     for (int j = tid; j < N; j += blk){
-        float log_k = - reg * M[bidx+j*N];
+        float log_k = - reg * M_T[bidx*N+j];
         if(local_max <= (log_k + log_u[j])){
             local_max = log_k + log_u[j];
         }     
@@ -208,7 +208,7 @@ __global__ void v_update_log(
 
     //log-domain iteration log_v = log_b - LSE(log_k+log_u)
     for (int j = tid; j < N; j += blk){
-        float log_k = - reg * M[bidx+j*N];
+        float log_k = - reg * M_T[bidx*N+j];
         local_sum += expf(log_k + log_u[j] - row_max);
     }
 
@@ -266,6 +266,7 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
     const torch::Tensor log_a,
     const torch::Tensor log_b,
     const torch::Tensor M,
+    const torch::Tensor M_T,
     torch::Tensor log_u,
     torch::Tensor log_v,
     int itrstep,
@@ -287,7 +288,7 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
         );
         v_update_log<<<blocks, threads, shared_memory>>>(
             log_b.data_ptr<float>(),
-            M.data_ptr<float>(),
+            M_T.data_ptr<float>(),
             log_u.data_ptr<float>(),
             log_v.data_ptr<float>(),
             N,
diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp
index 0e507a8..073982c 100644
--- a/torchmorph/csrc/torchmorph.cpp
+++ b/torchmorph/csrc/torchmorph.cpp
@@ -42,6 +42,7 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
     const torch::Tensor log_a,
     const torch::Tensor log_b,
     const torch::Tensor M,
+    const torch::Tensor M_T,
     torch::Tensor log_u,
     torch::Tensor log_v,
     int itrstep,
@@ -94,6 +95,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           py::arg("log_a"),
           py::arg("log_b"),
           py::arg("M"),
+          py::arg("M_T"),
           py::arg("log_u"),
           py::arg("log_v"),
           py::arg("itrstep"),
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 4dcff11..5efb1c6 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -3,14 +3,13 @@
 import torch
 from torch import Tensor
 
-# from torchmorph import _C
 
+def build_cost_matrix(shape, p=2, device='cuda') -> torch.Tensor:  # Euclidean distance by default
+    """Build the pairwise ground-cost matrix for a spatial grid.
 
-# generate cost matrix for multi-dimensional tensor like 2D or 3D tensor
-def build_cost_matrix(
-    shape, device: str = 'cuda', p=2  # Euclidean distance by default
-) -> torch.Tensor:
-
+    The shape defines grid coordinates, and p selects the distance norm.
+    Returns an (N, N) matrix for N flattened grid points.
+    """
     coords = torch.stack(torch.meshgrid([torch.arange(s) for s in shape], indexing="ij"), dim=-1)
     coords_flat = coords.view(-1, coords.size(-1)).float()
 
@@ -18,283 +17,295 @@ def build_cost_matrix(
     return cost_matrix
 
 
-def data_preprocess(
-    source: Tensor,
-    target: Tensor,
-    cost_matrix: Optional[Tensor] = None,
-    p: int = 2,
-    force_batched: bool = False,
-    device: str = 'cuda',
-    dtype=torch.float32,
-):
-    """Validate, move, reshape, and normalize Sinkhorn input distributions.
-
-    If batch_flag is True, 1D/2D inputs are reshaped to (1, 1, N).
-    Higher-dimensional inputs are reshaped from (B, C, *Spatial) to (B, C, N).
-    """
-    # check dimensions,reshape and choose strategy
-    if not source.ndim == target.ndim or not source.shape == target.shape:
-        raise ValueError(
-            "Source and target must have the same number of dimensions and the same shape."
+class SinkhornSolver:
+    def __init__(self, reg=0.02, itrstep=100, threshold=0, p=2, device='cuda'):
+        self.reg = reg
+        self.itrstep = itrstep
+        self.threshold = threshold
+        self.device = device
+        self.p = p
+
+    def data_preprocess(
+        self,
+        source: Tensor,
+        target: Tensor,
+        cost_matrix: Optional[Tensor] = None,
+        force_batched: bool = False,
+        dtype=torch.float32,
+    ):
+        """Preprocess source and target distributions for Sinkhorn solvers.
+
+        This checks matching shapes, flattens spatial axes, clamps negatives,
+        normalizes mass, and prepares the cost matrix transpose.
+        """
+        # check dimensions,reshape and choose strategy
+        if not source.ndim == target.ndim or not source.shape == target.shape:
+            raise ValueError(
+                "Source and target must have the same number of dimensions and the same shape."
+            )
+        elif source.ndim <= 2:
+            if cost_matrix is None:
+                cost_matrix = build_cost_matrix(source.shape, self.p, self.device)
+            if force_batched:
+                source = source.reshape(1, 1, -1)
+                target = target.reshape(1, 1, -1)
+            else:
+                source = source.reshape(-1)
+                target = target.reshape(-1)
+        elif source.ndim > 2:
+            # cut and reshape source and target to coordinate tensors, then calculate cost matrix
+            B, C = source.shape[:2]  # get batch and channel dimensions
+            if cost_matrix is None:
+                spatial_shape = source.shape[2:]
+                cost_matrix = build_cost_matrix(spatial_shape, self.p, self.device)
+            source = source.reshape(B, C, -1)
+            target = target.reshape(B, C, -1)
+
+        # move to specified device(cuda) and ensure dtype float32
+        source = source.to(device=self.device, dtype=dtype)
+        target = target.to(device=self.device, dtype=dtype)
+        cost_matrix = cost_matrix.to(device=self.device, dtype=dtype)
+
+        # ensure non-negativity
+        source = torch.clamp(source, min=0)
+        target = torch.clamp(target, min=0)
+
+        # normalize source and target to make them valid probability distributions
+        source /= torch.clamp(source.sum(dim=-1, keepdim=True).to(self.device), min=1e-12)
+        target /= torch.clamp(target.sum(dim=-1, keepdim=True).to(self.device), min=1e-12)
+
+        if not torch.allclose(cost_matrix, cost_matrix.transpose(-2, -1)):
+            cost_matrix_T = cost_matrix.transpose(-2, -1).contiguous()
+        else:
+            cost_matrix_T = cost_matrix
+
+        return source, target, cost_matrix, cost_matrix_T
+
+    def sinkhorn(
+        self,
+        source: Tensor,
+        target: Tensor,
+        cost_matrix: Optional[Tensor] = None,
+        returngrad: bool = False,
+        verbose: bool = False,
+    ):
+        """Compute a balanced Sinkhorn transport plan from raw inputs.
+
+        The function preprocesses inputs, runs the batched solver, and returns
+        either the plan or auxiliary variables when gradients are requested.
+        """
+        source, target, cost_matrix, _ = self.data_preprocess(
+            source,
+            target,
+            cost_matrix=cost_matrix,
+            force_batched=True,
         )
-    elif source.ndim <= 2:
-        if cost_matrix is None:
-            cost_matrix = build_cost_matrix(source.shape, device=device, p=p)
-        if force_batched:
-            source = source.view(1, 1, -1)
-            target = target.view(1, 1, -1)
+        result = self.sinkhorn_batch(
+            source, target, cost_matrix=cost_matrix, returnuv=returngrad, verbose=verbose
+        )
+
+        if returngrad:
+            return result
+        return result["plan"]
+
+    def sinkhorn_batch(
+        self,
+        source: Tensor,  # (B,C,*Spatial)
+        target: Tensor,  # (B,C,*Spatial)
+        cost_matrix: Tensor,
+        returnuv: bool = False,  # whether to return the gradient of the dual potential
+        verbose: bool = False,
+    ):
+        """Compute batched balanced Sinkhorn iterations in PyTorch.
+
+        Inputs should be preprocessed to shape (B, C, N). The solver builds
+        K = exp(-reg * C), updates u and v, and returns the transport plan.
+        """
+        # Device check
+        if not torch.cuda.is_available():
+            raise ValueError("CUDA device not available.")
+
+        if source.ndim != 3 or target.ndim != 3:
+            raise ValueError("source and target must be preprocessed to shape (B, C, N).")
+
+        # initialize u and v, and expand dimensions for batch and channel
+        u = torch.ones_like(source, device=self.device)
+        v = torch.ones_like(target, device=self.device)
+        k = torch.exp(-cost_matrix * self.reg)
+
+        B, C = source.shape[:2]
+        k = k.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
+
+        if self.reg <= 0:
+            raise ValueError("Regularization parameter must be positive.")
+
+        # generate sinkhorn iteration judged by threshold or iteration times
+        if self.itrstep < 0 or self.threshold < 0:
+            raise ValueError("iterstep and threshold must be non-negative.")
+        elif self.itrstep == 0 and self.threshold == 0:
+            raise ValueError("Invalid iteration or step settings. ")
+
+        itrstep = self.itrstep
+        if self.threshold > 0 and self.itrstep == 0:
+            itrstep = 500
+
+        # iteration
+        convergence_flag = False
+        convergence_check = 10
+        for i in range(itrstep):
+            u_old = u
+            u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
+            v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
+            if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < self.threshold:
+                convergence_flag = True
+                break
+
+        # use broadcasting to calculate the optimal transport plan P from u, v and K
+        P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
+
+        if verbose:
+            if convergence_flag:
+                print(
+                    f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}"
+                )
+            else:
+                print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete.")
+
+        # calculate gradient after iteration covergence ,return results
+        if returnuv:
+            return {"plan": P, "u": u, "v": v}
         else:
-            source = source.view(-1)
-            target = target.view(-1)
-    elif source.ndim > 2:
-        # cut and reshape source and target to coordinate tensors, then calculate cost matrix
-        B, C = source.shape[:2]  # get batch and channel dimensions
-        if cost_matrix is None:
-            spatial_shape = source.shape[2:]
-            cost_matrix = build_cost_matrix(spatial_shape, device=device, p=p)
-        source = source.view(B, C, -1)
-        target = target.view(B, C, -1)
-
-    # move to specified device(cuda) and ensure dtype float32
-    source = source.to(device=device, dtype=dtype)
-    target = target.to(device=device, dtype=dtype)
-    cost_matrix = cost_matrix.to(device=device, dtype=dtype)
-
-    # ensure non-negativity
-    source = torch.clamp(source, min=0)
-    target = torch.clamp(target, min=0)
-
-    # normalize source and target to make them valid probability distributions
-    source /= torch.clamp(source.sum(dim=-1, keepdim=True).to(device), min=1e-12)
-    target /= torch.clamp(target.sum(dim=-1, keepdim=True).to(device), min=1e-12)
-
-    return source, target, cost_matrix
-
-
-def sinkhorn_balanced(
-    source: Tensor,
-    target: Tensor,
-    cost_matrix: Optional[Tensor] = None,
-    p: int = 2,
-    reg: float = 1.0,
-    itrstep: int = 0,
-    threshold: float = 0,
-    returngrad: bool = False,
-    device: str = 'cuda',
-    verbose: bool = False,
-):
-    source, target, cost_matrix = data_preprocess(
-        source,
-        target,
-        cost_matrix=cost_matrix,
-        p=p,
-        device=device,
-        force_batched=True,
-    )
-    result = sinkhorn_balanced_batch(
-        source,
-        target,
-        cost_matrix=cost_matrix,
-        reg=reg,
-        itrstep=itrstep,
-        threshold=threshold,
-        returnuv=returngrad,
-        device=device,
-        verbose=verbose,
-    )
-
-    if returngrad:
-        return result
-    return result["plan"]
-
-
-def sinkhorn_balanced_batch(
-    source: Tensor,  # (B,C,*Spatial)
-    target: Tensor,  # (B,C,*Spatial)
-    cost_matrix: Tensor,
-    reg: float = 1.0,
-    itrstep: int = 0,
-    threshold: float = 0,
-    returnuv: bool = False,  # whether to return the gradient of the dual potential
-    device: str = 'cuda',
-    verbose: bool = False,
-):
-    """Compute a balanced entropic optimal transport plan with Sinkhorn iterations.
-
-    Builds K = exp(-reg * C), iteratively updates scaling vectors u and v,
-    and returns the transport plan P.
-
-    Args:
-        source (Tensor): Preprocessed source distribution with shape (B, C, N).
-        target (Tensor): Preprocessed target distribution with shape (B, C, N).
-        cost_matrix (Tensor): Pairwise cost matrix with shape (N, N).
-        reg (float): Inverse regularization parameter.
-        itrstep (int): Number of Sinkhorn iterations.
-        threshold (float): Optional convergence threshold.
-        returnuv (bool): If True, also return u and v.
-
-    Returns:
-        dict: Transport plan, and optionally scaling vectors u and v.
-    """
-    # Device check
-    if not torch.cuda.is_available():
-        raise ValueError("CUDA device not available.")
-
-    if source.ndim != 3 or target.ndim != 3:
-        raise ValueError("source and target must be preprocessed to shape (B, C, N).")
-
-    # initialize u and v, and expand dimensions for batch and channel
-    u = torch.ones_like(source, device=device)
-    v = torch.ones_like(target, device=device)
-    k = torch.exp(-cost_matrix * reg)
-
-    B, C = source.shape[:2]
-    k = k.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
-
-    if reg <= 0:
-        raise ValueError("Regularization parameter must be positive.")
-
-    # generate sinkhorn iteration judged by threshold or iteration times
-    if itrstep < 0 or threshold < 0:
-        raise ValueError("iterstep and threshold must be non-negative.")
-    elif itrstep == 0 and threshold == 0:
-        raise ValueError("Invalid iteration or step settings. ")
-
-    if threshold > 0 and itrstep == 0:
-        itrstep = 500
-
-    # iteration
-    convergence_flag = False
-    convergence_check = 10
-    for i in range(itrstep):
-        u_old = u
-        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
-        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
-        if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < threshold:
-            convergence_flag = True
-            break
-
-    # use broadcasting to calculate the optimal transport plan P from u, v and K
-    P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
-
-    if verbose:
-        if convergence_flag:
-            print(f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}")
+            return {"plan": P}
+
+    def sinkhorn__nobatch(
+        self,
+        source: Tensor,  # (N,)
+        target: Tensor,  # (N,)
+        cost_matrix: Tensor,
+        returnuv: bool = False,  # whether to return the gradient of the dual potential
+    ):
+        """Compute Sinkhorn iterations for one flat distribution pair.
+
+        Source and target should be one-dimensional probability vectors.
+        This path is a simple PyTorch reference for the non-batched case.
+        """
+
+        # initialize u and v, and expand dimensions for batch and channel
+        u = torch.ones_like(source, device=self.device)
+        v = torch.ones_like(target, device=self.device)
+        k = torch.exp(-cost_matrix * self.reg)
+
+        if self.reg <= 0:
+            raise ValueError("Regularization parameter must be positive.")
+
+        # generate sinkhorn iteration judged by threshold or iteration times
+        if self.itrstep < 0 or self.threshold < 0:
+            raise ValueError("iterstep and threshold must be non-negative.")
+        elif self.itrstep == 0 and self.threshold == 0:
+            raise ValueError("Invalid iteration or step settings. ")
+
+        itrstep = self.itrstep
+        if self.threshold > 0 and self.itrstep == 0:
+            itrstep = 500
+
+        # iteration
+        convergence_check = 10
+        for i in range(itrstep):
+            u_old = u
+            u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
+            v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
+            if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < self.threshold:
+                break
+
+        # use broadcasting to calculate the optimal transport plan P from u, v and K
+        P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
+
+        # calculate gradient after iteration covergence ,return results
+        if returnuv:
+            return {"plan": P, "u": u, "v": v}
+        else:
+            return {"plan": P}
+
+    def gradient(self, u: Tensor, v: Tensor):
+        """Convert Sinkhorn scaling vectors to source and target gradients.
+
+        The scaling vectors are logged, centered to remove additive constants,
+        and rescaled according to this module's regularization convention.
+        """
+        u = torch.log(torch.clamp(u, min=1e-12))
+        v = torch.log(torch.clamp(v, min=1e-12))
+
+        u = u - u.mean(dim=-1, keepdim=True)  # remove constants
+        v = v - v.mean(dim=-1, keepdim=True)
+
+        grad_source = u / self.reg
+        grad_target = v / self.reg
+
+        return grad_source, grad_target
+
+    def sinkhorn_cuda(
+        self,
+        source: Tensor,
+        target: Tensor,
+        cost_matrix: Tensor,
+        returnuv: bool = False,
+    ):
+        """Run the CUDA Sinkhorn solver in standard scaling form.
+
+        The CUDA kernel updates u and v, after which the Python wrapper
+        reconstructs the transport plan from u, K, and v.
+        """
+        from torchmorph import _C
+
+        N = source.shape[-1]
+        K = torch.exp(-cost_matrix * self.reg)
+        u = torch.ones_like(source)
+        v = torch.ones_like(target)
+
+        u, v = _C.sinkhorn_fastiter(source, target, K, u, v, int(self.itrstep), int(N))
+
+        P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
+
+        if returnuv:
+            return {"plan": P, "u": u, "v": v}
         else:
-            print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete.")
-
-    # calculate gradient after iteration covergence ,return results
-    if returnuv:
-        return {"plan": P, "u": u, "v": v}
-    else:
-        return {"plan": P}
-
-
-def sinkhorn_balanced_nobatch(
-    source: Tensor,  # (N,)
-    target: Tensor,  # (N,)
-    cost_matrix: Tensor,
-    reg: float = 1.0,
-    itrstep: int = 0,
-    threshold: float = 0,
-    returnuv: bool = False,  # whether to return the gradient of the dual potential
-    device: str = 'cuda',
-):
-    device = device if torch.cuda.is_available() else 'cpu'
-
-    # initialize u and v, and expand dimensions for batch and channel
-    u = torch.ones_like(source, device=device)
-    v = torch.ones_like(target, device=device)
-    k = torch.exp(-cost_matrix * reg)
-
-    if reg <= 0:
-        raise ValueError("Regularization parameter must be positive.")
-
-    # generate sinkhorn iteration judged by threshold or iteration times
-    if itrstep < 0 or threshold < 0:
-        raise ValueError("iterstep and threshold must be non-negative.")
-    elif itrstep == 0 and threshold == 0:
-        raise ValueError("Invalid iteration or step settings. ")
-
-    if threshold > 0 and itrstep == 0:
-        itrstep = 500
-
-    # iteration
-    convergence_check = 10
-    for i in range(itrstep):
-        u_old = u
-        u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
-        v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
-        if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < threshold:
-            break
-
-    # use broadcasting to calculate the optimal transport plan P from u, v and K
-    P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
-
-    # calculate gradient after iteration covergence ,return results
-    if returnuv:
-        return {"plan": P, "u": u, "v": v}
-    else:
-        return {"plan": P}
-
-
-def sinkhorn_gradient(u: Tensor, v: Tensor, reg: float):
-    u = torch.log(torch.clamp(u, min=1e-12))
-    v = torch.log(torch.clamp(v, min=1e-12))
-
-    u = u - u.mean(dim=-1, keepdim=True)  # remove constants
-    v = v - v.mean(dim=-1, keepdim=True)
-
-    grad_source = u / reg
-    grad_target = v / reg
-
-    return grad_source, grad_target
-
-
-def sinkhorn_balanced_cuda(
-    source: Tensor,
-    target: Tensor,
-    cost_matrix: Tensor,
-    reg: float = 1.0,
-    itrstep: int = 100,
-    returnuv: bool = False,
-):
-    from torchmorph import _C
-
-    N = source.shape[-1]
-    K = torch.exp(-cost_matrix * reg)
-    u = torch.ones_like(source)
-    v = torch.ones_like(target)
-
-    u, v = _C.sinkhorn_fastiter(source, target, K, u, v, int(itrstep), int(N))
-
-    P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
-
-    if returnuv:
-        return {"plan": P, "u": u, "v": v}
-    else:
-        return {"plan": P}
-
-
-def sinkhorn_balanced_log(
-    source: Tensor,
-    target: Tensor,
-    cost_matrix: Tensor,
-    reg: float = 1.0,
-    itrstep: int = 100,
-    dtype=torch.float32,
-):
-    from torchmorph import _C
-
-    source = source.to(dtype=dtype).log()
-    target = target.to(dtype=dtype).log()
-    N = source.shape[-1]
-    u = torch.zeros_like(source)
-    v = torch.zeros_like(target)
-
-    u, v = _C.sinkhorn_logiter(source, target, cost_matrix, u, v, int(itrstep), int(N), float(reg))
-
-    u = u.exp()
-    v = v.exp()
-
-    return u, v
+            return {"plan": P}
+
+    def sinkhorn_log_cuda(
+        self,
+        source: Tensor,
+        target: Tensor,
+        cost_matrix: Tensor,
+        cost_matrix_T: Tensor,
+        dtype=torch.float32,
+    ):
+        """Run the CUDA Sinkhorn solver in log-domain form.
+
+        Inputs are converted to log probabilities for numerical stability.
+        The returned u and v are exponentiated scaling vectors.
+        """
+        from torchmorph import _C
+
+        with torch.no_grad():
+            source = source.to(dtype=dtype).log()
+            target = target.to(dtype=dtype).log()
+            N = source.shape[-1]
+            u = torch.zeros_like(source)
+            v = torch.zeros_like(target)
+
+            u, v = _C.sinkhorn_logiter(
+                source,
+                target,
+                cost_matrix,
+                cost_matrix_T,
+                u,
+                v,
+                int(self.itrstep),
+                int(N),
+                float(self.reg),
+            )
+
+            u = u.exp()
+            v = v.exp()
+
+        return u, v

From 549410a7c735d6450a5ef78659c8788db9ec6e8f Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Tue, 2 Jun 2026 15:35:00 +0800
Subject: [PATCH 10/20] more test

---
 test/test_optimal_transport.py  | 282 +++++++++++++++++++-------------
 torchmorph/optimal_transport.py |  92 +++++++++++
 2 files changed, 257 insertions(+), 117 deletions(-)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index ec5d82c..cd37949 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -1,8 +1,15 @@
 import importlib.util
 from pathlib import Path
 
+import numpy as np
+import pytest
 import torch
 
+try:
+    import ot
+except ImportError:  # pragma: no cover - exercised only when POT is absent
+    ot = None
+
 # Project root directory (parent of test)
 project_root = Path(__file__).parent.parent
 
@@ -15,6 +22,44 @@
 spec.loader.exec_module(tr)
 
 
+def _require_cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for Sinkhorn tests")
+
+
+def _require_pot():
+    if ot is None:
+        pytest.skip("POT is required for baseline Sinkhorn tests")
+
+
+def _make_positive(shape, device, seed):
+    generator = torch.Generator(device=device).manual_seed(seed)
+    return torch.rand(shape, generator=generator, device=device) + 0.1
+
+
+def _expected_batch_channel_n(shape):
+    if len(shape) <= 2:
+        n = int(np.prod(shape))
+        return 1, 1, n
+    b, c = shape[:2]
+    n = int(np.prod(shape[2:]))
+    return b, c, n
+
+
+def _pot_plan(source, target, cost_matrix, solver):
+    source_np = source.detach().cpu().numpy()
+    target_np = target.detach().cpu().numpy()
+    cost_np = cost_matrix.detach().cpu().numpy()
+    return ot.sinkhorn(
+        source_np,
+        target_np,
+        cost_np,
+        reg=1.0 / solver.reg,
+        numItermax=solver.itrstep,
+        stopThr=1e-9,
+    )
+
+
 def test_build_cost_matrix():
     """Simple test for the correctness of the build_cost_matrix function."""
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -71,125 +116,136 @@ def test_build_cost_matrix():
     assert torch.allclose(cost_l2, expected), "L2 distance matrix incorrect"
 
 
-def test_sinkhorn_balanced_print():
-
-    print("=" * 50)
-    print("Testing sinkhorn_balanced function")
-    print("Input shape: ")
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (4,),
+        (2, 3),
+        (2, 2, 2),
+        (1, 1, 2, 3),
+        (2, 3, 2, 2),
+        (1, 2, 2, 2, 2),
+    ],
+)
+def test_run_once_accepts_multiple_raw_dimensions(shape):
+    _require_cuda()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=300, threshold=0, device=device)
+
+    source = _make_positive(shape, device=device, seed=11)
+    target = _make_positive(shape, device=device, seed=23)
+    out = solver.run_once(
+        source,
+        target,
+        return_plan=True,
+        return_uv=True,
+        return_distance=True,
+    )
 
-    # Check CUDA availability
-    if not torch.cuda.is_available():
-        raise ValueError("Warning: CUDA not available, tests may fail.")
-    else:
-        device = 'cuda'
-        print("CUDA available, running on GPU.")
-
-    # Create two random distributions
-    source = torch.rand(1, 1, 2, 2, device=device)
-    target = torch.rand(1, 1, 2, 2, device=device)
-
-    print("\n--- Input source ---")
-    print(source)
-    print("\n--- Input target ---")
-    print(target)
-
-    print("\n--- Calling sinkhorn_balanced ---")
-    try:
-        solver = tr.SinkhornSolver(
-            reg=0.1,
-            itrstep=100,
-            threshold=1e-5,
-            device=device,
-            p=2,
-        )
-        P = solver.sinkhorn(
-            source,
-            target,
-            cost_matrix=None,
-            returngrad=False,
-        )
-        print("\n--- Output P (transport plan) shape ---")
-        print(P.shape)
-        print("\n--- Output P content (partial) ---")
-        print(P)
-    except Exception as e:
-        print("\n--- Error occurred ---")
-        print(f"Error type: {type(e).__name__}")
-        print(f"Error message: {e}")
-        import traceback
-
-        traceback.print_exc()
-
-
-def test_sinkhorn_cuda():
-    if not torch.cuda.is_available():
-        raise ValueError("CUDA not available, skipping CUDA test.")
+    batch, channel, n = _expected_batch_channel_n(shape)
+    assert out["source"].shape == (batch, channel, n)
+    assert out["target"].shape == (batch, channel, n)
+    assert out["cost_matrix"].shape == (n, n)
+    assert out["plan"].shape == (batch, channel, n, n)
+    assert out["u"].shape == (batch, channel, n)
+    assert out["v"].shape == (batch, channel, n)
+    assert out["distance"].shape == (batch, channel)
+
+    distance_from_plan = (out["plan"] * out["cost_matrix"]).sum(dim=(-2, -1))
+    assert torch.allclose(out["distance"], distance_from_plan, rtol=1e-5, atol=1e-6)
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (4,),
+        (2, 3),
+        (2, 2, 2),
+        (2, 2, 2, 2),
+    ],
+)
+def test_run_once_distance_matches_pot_sinkhorn(shape):
+    _require_cuda()
+    _require_pot()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+
+    source = _make_positive(shape, device=device, seed=101)
+    target = _make_positive(shape, device=device, seed=211)
+    out = solver.run_once(source, target, return_plan=True, return_distance=True)
+
+    expected_distances = []
+    for b in range(out["source"].shape[0]):
+        channel_distances = []
+        for c in range(out["source"].shape[1]):
+            pot_transport = _pot_plan(
+                out["source"][b, c],
+                out["target"][b, c],
+                out["cost_matrix"],
+                solver,
+            )
+            pot_distance = np.sum(pot_transport * out["cost_matrix"].detach().cpu().numpy())
+            channel_distances.append(pot_distance)
+        expected_distances.append(channel_distances)
 
-    device = 'cuda'
-    # Create two random distributions
-    torch.manual_seed(42)
-    source = torch.rand(2, 2, device=device)
-    target = torch.rand(2, 2, device=device)
-
-    print("\n--- Input source ---")
-    print(source)
-    print("\n--- Input target ---")
-    print(target)
-
-    solver = tr.SinkhornSolver(reg=0.02, itrstep=100, device=device)
-    source, target, cost_matrix, _ = solver.data_preprocess(source, target)
-
-    torch_result = solver.sinkhorn(
-        source=source,
-        target=target,
-        cost_matrix=cost_matrix,
-        returngrad=True,
-    )
+    expected = torch.tensor(expected_distances, device=device, dtype=out["distance"].dtype)
+    assert torch.allclose(out["distance"], expected, rtol=5e-3, atol=1e-5)
 
-    cuda_result = solver.sinkhorn_cuda(
-        source=source,
-        target=target,
-        cost_matrix=cost_matrix,
-    )
 
-    torch_plan = torch_result["plan"].squeeze(0).squeeze(0)
-    relative_error = torch.linalg.norm(cuda_result["plan"] - torch_plan) / torch.clamp(
-        torch.linalg.norm(torch_plan), min=1e-12
-    )
+def test_run_once_plan_matches_pot_sinkhorn_small():
+    _require_cuda()
+    _require_pot()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
 
-    print(f"CUDA relative error: {relative_error.item():.4e}")
-    assert relative_error < 1e-4
+    source = _make_positive((2, 2), device=device, seed=307)
+    target = _make_positive((2, 2), device=device, seed=401)
+    out = solver.run_once(source, target, return_plan=True)
 
+    pot_transport = _pot_plan(out["source"][0, 0], out["target"][0, 0], out["cost_matrix"], solver)
+    expected = torch.tensor(pot_transport, device=device, dtype=out["plan"].dtype)
+    assert torch.allclose(out["plan"][0, 0], expected, rtol=5e-3, atol=1e-5)
 
-def test_sinkhorn_log():
-    if not torch.cuda.is_available():
-        raise ValueError("CUDA not available, skipping CUDA test.")
 
-    device = 'cuda'
-    # Create two random distributions
-    torch.manual_seed(42)
-    source = torch.rand(2, 2, device=device).abs()
-    target = torch.rand(2, 2, device=device).abs()
-
-    solver = tr.SinkhornSolver(reg=0.02, itrstep=100, device=device)
-    source, target, cost_matrix, cost_matrix_T = solver.data_preprocess(source, target)
-
-    print("\n--- Input source ---")
-    print(source)
-    print("\n--- Input target ---")
-    print(target)
-
-    u, v = solver.sinkhorn_log_cuda(
-        source=source,
-        target=target,
-        cost_matrix=cost_matrix,
-        cost_matrix_T=cost_matrix_T,
+def test_gradient_matches_numeric_directional_derivative():
+    _require_cuda()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+
+    source = torch.tensor([0.20, 0.30, 0.25, 0.25], device=device)
+    target = torch.tensor([0.15, 0.35, 0.20, 0.30], device=device)
+    direction = torch.tensor([0.20, -0.10, 0.05, -0.15], device=device)
+    eps = 1e-3
+
+    def regularized_objective(x):
+        out = solver.run_once(x, target, return_plan=True, return_distance=True)
+        plan = out["plan"]
+        entropy_term = (plan * torch.log(torch.clamp(plan, min=1e-12))).sum()
+        return out["distance"].sum() + entropy_term / solver.reg
+
+    out = solver.run_once(
+        source,
+        target,
+        return_plan=False,
+        return_uv=True,
+        return_grad=True,
+        return_distance=True,
     )
+    analytic_grad = out["grad_source"].reshape(-1)
+    analytic_directional = torch.sum(analytic_grad * direction)
+
+    plus = source + eps * direction
+    minus = source - eps * direction
+    loss_plus = regularized_objective(plus)
+    loss_minus = regularized_objective(minus)
+    numeric_directional = (loss_plus - loss_minus) / (2 * eps)
 
-    print("\n--- Output u  ---")
-    print(u)
-    print("\n--- Output v  ---")
-    print(v)
+    assert torch.allclose(
+        analytic_directional,
+        numeric_directional,
+        rtol=2e-2,
+        atol=2e-3,
+    )
 
 
 def loss_test():
@@ -234,11 +290,3 @@ def loss_test():
         print(f"step {step}, loss = {loss_value.item():.6f}")
         print("source:", source.detach())
         print("target:", target)
-
-
-if __name__ == "__main__":
-    test_sinkhorn_balanced_print()
-    test_build_cost_matrix()
-    test_sinkhorn_log()
-    test_sinkhorn_cuda()
-    loss_test()
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 5efb1c6..c09fd79 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -108,6 +108,98 @@ def sinkhorn(
             return result
         return result["plan"]
 
+    def run_once(
+        self,
+        source: Tensor,
+        target: Tensor,
+        cost_matrix: Optional[Tensor] = None,
+        *,
+        use_cuda: bool = False,
+        log_domain: bool = False,
+        return_plan: bool = True,
+        return_uv: bool = False,
+        return_grad: bool = False,
+        return_distance: bool = False,
+        verbose: bool = False,
+        dtype=torch.float32,
+    ):
+        """Run the full Sinkhorn workflow from raw inputs.
+
+        This wrapper preprocesses raw source and target tensors, selects one
+        Sinkhorn backend, optionally reconstructs the transport plan, and
+        optionally returns the OT distance, dual-scaling vectors, or gradients.
+        """
+        source, target, cost_matrix, cost_matrix_T = self.data_preprocess(
+            source,
+            target,
+            cost_matrix=cost_matrix,
+            force_batched=True,
+            dtype=dtype,
+        )
+
+        need_uv = return_uv or return_grad or log_domain
+
+        if log_domain:
+            u, v = self.sinkhorn_log_cuda(
+                source=source,
+                target=target,
+                cost_matrix=cost_matrix,
+                cost_matrix_T=cost_matrix_T,
+                dtype=dtype,
+            )
+            backend = "cuda_log"
+            result = {"u": u, "v": v}
+        elif use_cuda:
+            result = self.sinkhorn_cuda(
+                source=source,
+                target=target,
+                cost_matrix=cost_matrix,
+                returnuv=need_uv,
+            )
+            backend = "cuda"
+        else:
+            result = self.sinkhorn_batch(
+                source=source,
+                target=target,
+                cost_matrix=cost_matrix,
+                returnuv=need_uv,
+                verbose=verbose,
+            )
+            backend = "torch"
+
+        output = {
+            "source": source,
+            "target": target,
+            "cost_matrix": cost_matrix,
+            "cost_matrix_T": cost_matrix_T,
+            "backend": backend,
+        }
+
+        plan = None
+        if return_plan or return_distance:
+            if "plan" in result:
+                plan = result["plan"]
+            else:
+                k = torch.exp(-cost_matrix * self.reg)
+                plan = result["u"].unsqueeze(-1) * k * result["v"].unsqueeze(-2)
+
+        if return_plan:
+            output["plan"] = plan
+
+        if return_distance:
+            output["distance"] = (plan * cost_matrix).sum(dim=(-2, -1))
+
+        if return_uv or return_grad:
+            output["u"] = result["u"]
+            output["v"] = result["v"]
+
+        if return_grad:
+            grad_source, grad_target = self.gradient(result["u"], result["v"])
+            output["grad_source"] = grad_source
+            output["grad_target"] = grad_target
+
+        return output
+
     def sinkhorn_batch(
         self,
         source: Tensor,  # (B,C,*Spatial)

From c0ee5c4835382bb331c40a1aa42bb07e730a9cb7 Mon Sep 17 00:00:00 2001
From: maokovski <“1498743918@qq.com”>
Date: Tue, 2 Jun 2026 16:51:49 +0800
Subject: [PATCH 11/20] try test

---
 .github/workflows/test.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 6b6efb6..e3d92f5 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,7 +17,9 @@ jobs:
               export CUDA_HOME=/usr/local/cuda
               export PATH="$CUDA_HOME/bin:$PATH"
               export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}"
+              export TORCH_CUDA_ARCH_LIST="8.9"
               echo "CUDA_HOME=$CUDA_HOME"
+              echo "TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
               pip install -r requirements-dev.txt --break-system-packages --user
               pip uninstall torchmorph --yes --break-system-packages
               python3 setup.py install --user

From 0442882e877b4256c66f73c20d1b2c2ff9d8941d Mon Sep 17 00:00:00 2001
From: maokovski <“1498743918@qq.com”>
Date: Tue, 2 Jun 2026 17:24:47 +0800
Subject: [PATCH 12/20] test on cuda kernal

---
 test/test_optimal_transport.py  | 57 +++++++++++++++++++++++++++++++++
 torchmorph/optimal_transport.py |  8 +++++
 2 files changed, 65 insertions(+)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index cd37949..7b90224 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -207,6 +207,63 @@ def test_run_once_plan_matches_pot_sinkhorn_small():
     assert torch.allclose(out["plan"][0, 0], expected, rtol=5e-3, atol=1e-5)
 
 
+def test_run_once_cuda_matches_torch_backend_small():
+    _require_cuda()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+
+    source = _make_positive((2, 2), device=device, seed=503)
+    target = _make_positive((2, 2), device=device, seed=607)
+
+    torch_out = solver.run_once(source, target, return_plan=True, return_distance=True)
+    cuda_out = solver.run_once(
+        source,
+        target,
+        use_cuda=True,
+        return_plan=True,
+        return_distance=True,
+    )
+
+    assert cuda_out["backend"] == "cuda"
+    assert torch.allclose(cuda_out["plan"], torch_out["plan"], rtol=1e-4, atol=1e-5)
+    assert torch.allclose(cuda_out["distance"], torch_out["distance"], rtol=1e-4, atol=1e-5)
+
+
+def test_run_once_log_domain_cuda_matches_torch_backend_small():
+    _require_cuda()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+
+    source = _make_positive((2, 2), device=device, seed=701)
+    target = _make_positive((2, 2), device=device, seed=809)
+
+    torch_out = solver.run_once(source, target, return_plan=True, return_distance=True)
+    log_out = solver.run_once(
+        source,
+        target,
+        log_domain=True,
+        return_plan=True,
+        return_distance=True,
+    )
+
+    assert log_out["backend"] == "cuda_log"
+    assert torch.allclose(log_out["plan"], torch_out["plan"], rtol=1e-4, atol=1e-5)
+    assert torch.allclose(log_out["distance"], torch_out["distance"], rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize("kwargs", [{"use_cuda": True}, {"log_domain": True}])
+def test_cuda_backends_reject_batched_inputs(kwargs):
+    _require_cuda()
+    device = "cuda"
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=300, threshold=0, device=device)
+
+    source = _make_positive((2, 1, 2, 2), device=device, seed=911)
+    target = _make_positive((2, 1, 2, 2), device=device, seed=919)
+
+    with pytest.raises(ValueError, match="B=C=1"):
+        solver.run_once(source, target, **kwargs)
+
+
 def test_gradient_matches_numeric_directional_derivative():
     _require_cuda()
     device = "cuda"
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index c09fd79..0f6e847 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -335,6 +335,12 @@ def gradient(self, u: Tensor, v: Tensor):
 
         return grad_source, grad_target
 
+    def _validate_single_cuda_distribution(self, source: Tensor, target: Tensor):
+        if source.ndim != 3 or target.ndim != 3:
+            raise ValueError("CUDA Sinkhorn inputs must be preprocessed to shape (B, C, N).")
+        if source.shape[:2] != (1, 1) or target.shape[:2] != (1, 1):
+            raise ValueError("CUDA Sinkhorn kernels currently support only B=C=1 inputs.")
+
     def sinkhorn_cuda(
         self,
         source: Tensor,
@@ -347,6 +353,7 @@ def sinkhorn_cuda(
         The CUDA kernel updates u and v, after which the Python wrapper
         reconstructs the transport plan from u, K, and v.
         """
+        self._validate_single_cuda_distribution(source, target)
         from torchmorph import _C
 
         N = source.shape[-1]
@@ -376,6 +383,7 @@ def sinkhorn_log_cuda(
         Inputs are converted to log probabilities for numerical stability.
         The returned u and v are exponentiated scaling vectors.
         """
+        self._validate_single_cuda_distribution(source, target)
         from torchmorph import _C
 
         with torch.no_grad():

From aab88576f3b79e1885662f4b3f70548d16109544 Mon Sep 17 00:00:00 2001
From: maokovski <1498743918@qq.com>
Date: Wed, 10 Jun 2026 14:47:39 +0800
Subject: [PATCH 13/20] tighten numeric gradient test and support more
 dimensions

---
 test/test_optimal_transport.py | 58 ++++++++++++++++++++++++++++------
 1 file changed, 48 insertions(+), 10 deletions(-)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 7b90224..53eb541 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -264,18 +264,43 @@ def test_cuda_backends_reject_batched_inputs(kwargs):
         solver.run_once(source, target, **kwargs)
 
 
-def test_gradient_matches_numeric_directional_derivative():
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (4,),
+        (2, 3),
+        (2, 2, 2),
+        (1, 1, 2, 3),
+        (2, 2, 2, 2),
+    ],
+)
+def test_gradient_matches_numeric_directional_derivative(shape):
     _require_cuda()
     device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+    solver = tr.SinkhornSolver(reg=0.7, itrstep=3000, threshold=0, device=device)
+
+    source_raw = _make_positive(shape, device=device, seed=1009)
+    target_raw = _make_positive(shape, device=device, seed=2003)
+    source, target, cost_matrix, _ = solver.data_preprocess(
+        source_raw,
+        target_raw,
+        force_batched=True,
+        dtype=torch.float64,
+    )
 
-    source = torch.tensor([0.20, 0.30, 0.25, 0.25], device=device)
-    target = torch.tensor([0.15, 0.35, 0.20, 0.30], device=device)
-    direction = torch.tensor([0.20, -0.10, 0.05, -0.15], device=device)
-    eps = 1e-3
+    direction = _make_positive(source.shape, device=device, seed=3001).double() - 0.6
+    direction = direction - direction.mean(dim=-1, keepdim=True)
+    eps = 1e-4
 
     def regularized_objective(x):
-        out = solver.run_once(x, target, return_plan=True, return_distance=True)
+        out = solver.run_once(
+            x,
+            target,
+            cost_matrix=cost_matrix,
+            return_plan=True,
+            return_distance=True,
+            dtype=torch.float64,
+        )
         plan = out["plan"]
         entropy_term = (plan * torch.log(torch.clamp(plan, min=1e-12))).sum()
         return out["distance"].sum() + entropy_term / solver.reg
@@ -283,13 +308,15 @@ def regularized_objective(x):
     out = solver.run_once(
         source,
         target,
+        cost_matrix=cost_matrix,
         return_plan=False,
         return_uv=True,
         return_grad=True,
         return_distance=True,
+        dtype=torch.float64,
     )
     analytic_grad = out["grad_source"].reshape(-1)
-    analytic_directional = torch.sum(analytic_grad * direction)
+    analytic_directional = torch.sum(analytic_grad * direction.reshape(-1))
 
     plus = source + eps * direction
     minus = source - eps * direction
@@ -300,8 +327,8 @@ def regularized_objective(x):
     assert torch.allclose(
         analytic_directional,
         numeric_directional,
-        rtol=2e-2,
-        atol=2e-3,
+        rtol=1e-7,
+        atol=1e-7,
     )
 
 
@@ -347,3 +374,14 @@ def loss_test():
         print(f"step {step}, loss = {loss_value.item():.6f}")
         print("source:", source.detach())
         print("target:", target)
+
+
+if __name__ == "__main__":
+    pytest.main(
+        [
+            __file__,
+            "-k",
+            "not loss_test",
+            "-q",
+        ]
+    )

From a7d80b4600895307057eafa963f48b50a8b77aa8 Mon Sep 17 00:00:00 2001
From: Yiping Mao <1498743918@qq.com>
Date: Wed, 10 Jun 2026 14:54:15 +0800
Subject: [PATCH 14/20] cancel unnecessary change in test.yml

Removed TORCH_CUDA_ARCH_LIST environment variable from the workflow.
---
 .github/workflows/test.yml | 2 --
 1 file changed, 2 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index e3d92f5..6b6efb6 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,9 +17,7 @@ jobs:
               export CUDA_HOME=/usr/local/cuda
               export PATH="$CUDA_HOME/bin:$PATH"
               export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}"
-              export TORCH_CUDA_ARCH_LIST="8.9"
               echo "CUDA_HOME=$CUDA_HOME"
-              echo "TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
               pip install -r requirements-dev.txt --break-system-packages --user
               pip uninstall torchmorph --yes --break-system-packages
               python3 setup.py install --user

From 9b55ac451e37005a5c7c92313cb0220db51b63c1 Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Wed, 8 Jul 2026 13:02:27 +0800
Subject: [PATCH 15/20] Rework Sinkhorn OT: standard epsilon convention,
 log-space end-to-end, batched (n, d) API

- Use the standard convention K = exp(-C / epsilon); POT baselines need no
  parameter inversion anymore.
- Keep log potentials in log space end-to-end: the log-domain CUDA solver
  returns (log_u, log_v), gradients are computed as f = eps * log_u, and the
  transport plan is reconstructed as exp(log_u - C/eps + log_v) for all
  backends so it cannot underflow at small epsilon.
- Every backend takes strictly (n, d) batched histograms sharing one (d, d)
  cost matrix; the CUDA kernels are batched over a 2-D launch grid, removing
  the previous B=C=1 restriction.
- Merge the four kernels into two (scaling_update, log_scaling_update), fixing
  the uncoalesced K^T reads; collapse entry points into solve(); delete the
  dead sinkhorn__nobatch and loss_test.
- CUDA hygiene: current-stream launches, CUDAGuard, TORCH_CHECK input
  validation, kernel-launch checks, and a NaN guard for all -inf LSE rows.
- Export SinkhornSolver / build_cost_matrix from the package; tests import it
  directly, run the torch backend on CPU, and add a small-epsilon stability
  test against POT's sinkhorn_log.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 benchmark/optimal_transport.py  | 454 ++++++------------------------
 test/test_optimal_transport.py  | 454 +++++++++++-------------------
 torchmorph/__init__.py          |   3 +
 torchmorph/csrc/ot_kernel.cu    | 380 +++++++++----------------
 torchmorph/csrc/torchmorph.cpp  |  46 ++--
 torchmorph/optimal_transport.py | 473 ++++++++------------------------
 6 files changed, 514 insertions(+), 1296 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index 44fd210..4d1086f 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -1,425 +1,145 @@
-import importlib.util
-from pathlib import Path
-
 import ot
 import torch
 from torch.utils import benchmark
 
-# Project root directory (parent of test)
-project_root = Path(__file__).parent.parent
-
-# Target module path and Load module
-module_path = project_root / "torchmorph" / "optimal_transport.py"
-spec = importlib.util.spec_from_file_location("optimal_transport", module_path)
-tr = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(tr)
+from torchmorph import SinkhornSolver, build_cost_matrix
 
 
-def run_sinkhorn_balanced_benchmark(
-    B: int = 1,
-    C: int = 1,
-    H: int = 32,
-    W: int = 32,
-    myreg: float = 0.02,
-    itrstep: int = 100,
-    threshold: float = 0,
-):
-    torch.cuda.reset_peak_memory_stats()
-    # Set parameters
+def _grid_problem(n, H, W, epsilon, max_iter, seed=42):
+    """Random (n, H*W) marginals with a 2-D grid cost matrix, preprocessed."""
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    torch.manual_seed(42)
-
-    # Prepare input data
-    source = torch.rand(B, C, H, W, device=device, dtype=torch.float32)
-    target = torch.rand(B, C, H, W, device=device, dtype=torch.float32)
-    source = source / source.sum()
-    target = target / target.sum()
-    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
-
-    solver = tr.SinkhornSolver(
-        reg=myreg,
-        itrstep=itrstep,
-        threshold=threshold,
-        p=2,
-        device=device,
+    torch.manual_seed(seed)
+    source = torch.rand(n, H * W, device=device)
+    target = torch.rand(n, H * W, device=device)
+    solver = SinkhornSolver(epsilon=epsilon, max_iter=max_iter, device=device)
+    source, target, cost_matrix = solver.data_preprocess(
+        source, target, cost_matrix=build_cost_matrix((H, W), device=device)
     )
-    source, target, cost_matrix, _ = solver.data_preprocess(source, target, cost_matrix=cost_matrix)
-
-    globals_dict = {
-        "solver": solver,
-        "source": source,
-        "target": target,
-        "cost_matrix": cost_matrix,
-    }
-
-    stmt = "solver.sinkhorn_batch(source, target, cost_matrix=cost_matrix)"
-
-    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
-    # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
-    # number_of_runs = 100
-    result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=3)
-
-    peak_allocated = torch.cuda.max_memory_allocated()
-    peak_reserved = torch.cuda.max_memory_reserved()
-
-    print(f"peak allocated: {peak_allocated / 1024**2:.2f} MB")
-    print(f"peak reserved:   {peak_reserved / 1024**2:.2f} MB")
+    return solver, source, target, cost_matrix
 
-    return result_sinkhorn_balanced
 
+def _time_backend(method_name, solver, source, target, cost_matrix, min_run_time=3):
+    torch.cuda.reset_peak_memory_stats()
+    timer = benchmark.Timer(
+        stmt=f"solver.{method_name}(source, target, cost_matrix)",
+        globals={
+            "solver": solver,
+            "source": source,
+            "target": target,
+            "cost_matrix": cost_matrix,
+        },
+    )
+    result = timer.blocked_autorange(min_run_time=min_run_time)
+    print(f"peak allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+    print(f"peak reserved:  {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
+    return result
 
-def run_ot_sinkhorn_benchmark(H: int = 32, W: int = 32, myreg: float = 0.02):
-    # Set parameters
-    import ot
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype = torch.float32
-    torch.manual_seed(42)
-
-    # Prepare input data
-    source = torch.rand(H * W, device=device, dtype=dtype)
-    target = torch.rand(H * W, device=device, dtype=dtype)
-    source = source / source.sum()
-    target = target / target.sum()
-    M = tr.build_cost_matrix((H, W), device=device, p=2)
 
-    globals_dict = {
-        "ot": ot,
-        "source": source,
-        "target": target,
-        "M": M,
-        "myreg": myreg,
-    }
+def run_sinkhorn_torch_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100):
+    args = _grid_problem(n, H, W, epsilon, max_iter)
+    return _time_backend("sinkhorn_torch", *args)
 
-    stmt = "ot.sinkhorn(source, target, M, reg=1/myreg, numItermax=100, stopThr=1e-5)"
 
-    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
-    # timer.blocked_autorange(min_run_time=1)
-    # number_of_runs = 100
-    result_ot = timer.blocked_autorange(min_run_time=2)
-    return result_ot
+def run_sinkhorn_cuda_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100):
+    args = _grid_problem(n, H, W, epsilon, max_iter)
+    return _time_backend("sinkhorn_cuda", *args)
 
 
-@torch.no_grad()
-def run_sinkhorn_relative_error(H: int = 32, W: int = 32, myreg: float = 0.02):
-    # Set parameters
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype = torch.float32
-    torch.manual_seed(42)
+def run_sinkhorn_log_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100):
+    args = _grid_problem(n, H, W, epsilon, max_iter)
+    return _time_backend("sinkhorn_log_cuda", *args)
 
-    # Prepare the same input data for both implementations.
-    source = torch.rand(1, 1, H, W, device=device, dtype=dtype)
-    target = torch.rand(1, 1, H, W, device=device, dtype=dtype)
-    source = source / source.sum()
-    target = target / target.sum()
-    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
 
-    solver = tr.SinkhornSolver(
-        reg=myreg,
-        itrstep=100,
-        threshold=1e-5,
-        p=2,
-        device=device,
-    )
-    source_preprocessed, target_preprocessed, cost_matrix, _ = solver.data_preprocess(
-        source, target, cost_matrix=cost_matrix
+def run_pot_sinkhorn_benchmark(H=32, W=32, epsilon=1.0, max_iter=100):
+    solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter)
+    timer = benchmark.Timer(
+        stmt="ot.sinkhorn(source, target, M, reg=epsilon, numItermax=max_iter, stopThr=1e-5)",
+        globals={
+            "ot": ot,
+            "source": source[0],
+            "target": target[0],
+            "M": cost_matrix,
+            "epsilon": epsilon,
+            "max_iter": max_iter,
+        },
     )
+    return timer.blocked_autorange(min_run_time=2)
 
-    sinkhorn_plan = (
-        solver.sinkhorn_batch(
-            source_preprocessed,
-            target_preprocessed,
-            cost_matrix=cost_matrix,
-        )["plan"]
-        .squeeze(0)
-        .squeeze(0)
-    )
 
+@torch.no_grad()
+def run_sinkhorn_relative_error(H=32, W=32, epsilon=1.0, max_iter=100):
+    solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter)
+    out = solver.solve(source, target, cost_matrix=cost_matrix, return_plan=True)
     ot_plan = ot.sinkhorn(
-        source_preprocessed.squeeze(0).squeeze(0),
-        target_preprocessed.squeeze(0).squeeze(0),
-        cost_matrix,
-        reg=1 / myreg,
-        numItermax=100,
-        stopThr=1e-5,
+        source[0], target[0], cost_matrix, reg=epsilon, numItermax=max_iter, stopThr=1e-5
     )
-
-    return torch.linalg.norm(sinkhorn_plan - ot_plan) / torch.clamp(
+    return torch.linalg.norm(out["plan"][0] - ot_plan) / torch.clamp(
         torch.linalg.norm(ot_plan), min=1e-12
     )
 
 
-def large_scale_sinkhorn_balanced_benchmark(B: int = 1, C: int = 1, H: int = 128, W: int = 128):
-    # Set parameters
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype = torch.float32
-
-    # Prepare input data
-    source = torch.rand(B, C, H, W, device=device, dtype=dtype)
-    target = torch.rand(B, C, H, W, device=device, dtype=dtype)
-    source = source / source.sum()
-    target = target / target.sum()
-    solver = tr.SinkhornSolver(
-        reg=25.0,
-        itrstep=100,
-        threshold=1e-5,
-        device=device,
-    )
-
-    globals_dict = {
-        "solver": solver,
-        "source": source,
-        "target": target,
-    }
-
-    stmt = "solver.sinkhorn(source, target)"
-
-    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
-    # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
-    number_of_runs = 10
-    result_sinkhorn_balanced = timer.timeit(number_of_runs)
-    return result_sinkhorn_balanced
-
-
-def batch_channel_sinkhorn_balanced_benchmark():
-    # Set parameters
-    B, C, H, W = 32, 32, 32, 32
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    dtype = torch.float32
-
-    # Prepare input data
-    source = torch.rand(B, C, H, W, device=device, dtype=dtype)
-    target = torch.rand(B, C, H, W, device=device, dtype=dtype)
-    source = source / source.sum()
-    target = target / target.sum()
-    solver = tr.SinkhornSolver(
-        reg=25.0,
-        itrstep=100,
-        threshold=1e-5,
-        device=device,
-    )
-
-    globals_dict = {
-        "solver": solver,
-        "source": source,
-        "target": target,
-    }
-
-    stmt = "solver.sinkhorn(source, target)"
-
-    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
-    # result_sinkhorn_balanced = timer.blocked_autorange(min_run_time=1)
-    number_of_runs = 10
-    result_sinkhorn_balanced = timer.timeit(number_of_runs)
-    return result_sinkhorn_balanced
-
-
-def run_sinkhorn_log_benchmark(
-    H: int = 32,
-    W: int = 32,
-    myreg: float = 0.02,
-    itrstep: int = 100,
-):
-    torch.cuda.reset_peak_memory_stats()
-    # Set parameters
-    if not torch.cuda.is_available():
-        raise ValueError("CUDA device not available.")
-
-    device = "cuda"
-    torch.manual_seed(42)
-
-    # Prepare input data
-    source = torch.rand(H, W, device=device, dtype=torch.float32)
-    target = torch.rand(H, W, device=device, dtype=torch.float32)
-    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
-
-    solver = tr.SinkhornSolver(reg=myreg, itrstep=itrstep, p=2, device=device)
-    source, target, cost_matrix, cost_matrix_T = solver.data_preprocess(
-        source, target, cost_matrix=cost_matrix
-    )
-
-    globals_dict = {
-        "solver": solver,
-        "source": source,
-        "target": target,
-        "cost_matrix": cost_matrix,
-        "cost_matrix_T": cost_matrix_T,
-    }
-
-    stmt = "solver.sinkhorn_log_cuda(source, target, cost_matrix, cost_matrix_T)"
-
-    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
-    result_sinkhorn_log = timer.blocked_autorange(min_run_time=3)
-
-    peak_allocated = torch.cuda.max_memory_allocated()
-    peak_reserved = torch.cuda.max_memory_reserved()
-
-    print(f"peak allocated: {peak_allocated / 1024**2:.2f} MB")
-    print(f"peak reserved:   {peak_reserved / 1024**2:.2f} MB")
-
-    return result_sinkhorn_log
-
-
-def run_sinkhorn_cuda_benchmark(
-    H: int = 32,
-    W: int = 32,
-    myreg: float = 0.02,
-    itrstep: int = 100,
-):
-    torch.cuda.reset_peak_memory_stats()
-    # Set parameters
-    if not torch.cuda.is_available():
-        raise ValueError("CUDA device not available.")
-
-    device = "cuda"
-    torch.manual_seed(42)
-
-    # Prepare input data
-    source = torch.rand(H, W, device=device, dtype=torch.float32)
-    target = torch.rand(H, W, device=device, dtype=torch.float32)
-    cost_matrix = tr.build_cost_matrix((H, W), device=device, p=2)
-
-    solver = tr.SinkhornSolver(reg=myreg, itrstep=itrstep, p=2, device=device)
-    source, target, cost_matrix, _ = solver.data_preprocess(source, target, cost_matrix=cost_matrix)
-
-    globals_dict = {
-        "solver": solver,
-        "source": source,
-        "target": target,
-        "cost_matrix": cost_matrix,
-    }
-
-    stmt = "solver.sinkhorn_cuda(source, target, cost_matrix=cost_matrix)"
-
-    timer = benchmark.Timer(stmt=stmt, globals=globals_dict)
-    result_sinkhorn_cuda = timer.blocked_autorange(min_run_time=3)
-
-    peak_allocated = torch.cuda.max_memory_allocated()
-    peak_reserved = torch.cuda.max_memory_reserved()
-
-    print(f"peak allocated: {peak_allocated / 1024**2:.2f} MB")
-    print(f"peak reserved:   {peak_reserved / 1024**2:.2f} MB")
-
-    return result_sinkhorn_cuda
-
-
-def reference_error_check(
-    size: int = 32, myreg: float = 0.1, itrstep: int = 200, verbose: bool = False
-):
+def reference_error_check(size=32, epsilon=10.0, max_iter=200, verbose=False):
+    """Compare log-domain CUDA gradients against POT's sinkhorn_log potentials."""
     if not torch.cuda.is_available():
-        raise ValueError("CUDA device not available.")
-
-    device = "cuda"
-    torch.manual_seed(2)
+        raise RuntimeError("CUDA device not available.")
+    solver, source, target, cost_matrix = _grid_problem(1, size, size, epsilon, max_iter, seed=2)
 
-    # Prepare input data
-    source = torch.rand(size, size, device=device, dtype=torch.float32)
-    target = torch.rand(size, size, device=device, dtype=torch.float32)
-    cost_matrix = tr.build_cost_matrix((size, size), device=device, p=2)
-
-    solver = tr.SinkhornSolver(reg=myreg, itrstep=itrstep, p=2, device=device)
-    source, target, cost_matrix, cost_matrix_T = solver.data_preprocess(
-        source, target, cost_matrix=cost_matrix
-    )
+    log_u, log_v = solver.sinkhorn_log_cuda(source, target, cost_matrix)
+    grad_f, grad_g = solver.gradient(log_u, log_v)
+    grad_f, grad_g = grad_f[0], grad_g[0]
 
     _, log = ot.bregman.sinkhorn_log(
-        source,
-        target,
+        source[0],
+        target[0],
         cost_matrix,
-        reg=1 / myreg,
-        numItermax=itrstep,
+        reg=epsilon,
+        numItermax=max_iter,
         stopThr=0,
         log=True,
     )
-
-    u, v = solver.sinkhorn_log_cuda(
-        source,
-        target,
-        cost_matrix,
-        cost_matrix_T,
-    )
-
-    grad_f, grad_g = solver.gradient(u, v)
-
-    log_u = log["log_u"]
-    log_v = log["log_v"]
-
-    f = log_u / myreg
-    g = log_v / myreg
+    f = epsilon * log["log_u"]
+    g = epsilon * log["log_v"]
     f = f - f.mean()
     g = g - g.mean()
 
     if verbose:
-        print(f)
-        print(g)
-        print(grad_f)
-        print(grad_g)
+        print(f, g, grad_f, grad_g, sep="\n")
 
-    f_abs_error = (f - grad_f).abs()
-    g_abs_error = (g - grad_g).abs()
     f_rel_l2 = torch.linalg.norm(f - grad_f) / torch.clamp(torch.linalg.norm(f), min=1e-12)
     g_rel_l2 = torch.linalg.norm(g - grad_g) / torch.clamp(torch.linalg.norm(g), min=1e-12)
-
     return {
         "size": size,
         "points": size * size,
-        "reg": myreg,
-        "itrstep": itrstep,
-        "f_grad_max_abs": f_abs_error.max().item(),
-        "f_grad_mean_abs": f_abs_error.mean().item(),
+        "epsilon": epsilon,
+        "max_iter": max_iter,
+        "f_grad_max_abs": (f - grad_f).abs().max().item(),
         "f_grad_rel_l2": f_rel_l2.item(),
-        "g_grad_max_abs": g_abs_error.max().item(),
-        "g_grad_mean_abs": g_abs_error.mean().item(),
+        "g_grad_max_abs": (g - grad_g).abs().max().item(),
         "g_grad_rel_l2": g_rel_l2.item(),
         "f_allclose": torch.allclose(f, grad_f, rtol=1e-4, atol=1e-4),
         "g_allclose": torch.allclose(g, grad_g, rtol=1e-4, atol=1e-4),
     }
 
 
-def print_reference_error_table(
-    sizes=(2, 4, 8, 16, 32, 64),
-    itrsteps=(100, 200, 500),
-    myreg: float = 0.1,
-):
+def print_reference_error_table(sizes=(2, 4, 8, 16, 32, 64), max_iters=(100, 200, 500), epsilon=10.0):
     """Print a size and iteration sweep for POT-vs-CUDA gradient error."""
-    rows = [
-        reference_error_check(size=size, myreg=myreg, itrstep=itrstep)
-        for size in sizes
-        for itrstep in itrsteps
-    ]
-
-    headers = (
-        "size",
-        "N",
-        "reg",
-        "itr",
-        "f_max",
-        "f_mean",
-        "f_rel_l2",
-        "g_max",
-        "g_mean",
-        "g_rel_l2",
-        "close",
-    )
+    headers = ("size", "N", "eps", "itr", "f_max", "f_rel_l2", "g_max", "g_rel_l2", "close")
     print("| " + " | ".join(headers) + " |")
     print("| " + " | ".join(["---"] * len(headers)) + " |")
-    for row in rows:
-        close = f"{row['f_allclose']}/{row['g_allclose']}"
-        print(
-            f"| {row['size']}x{row['size']} "
-            f"| {row['points']} "
-            f"| {row['reg']:.3g} "
-            f"| {row['itrstep']} "
-            f"| {row['f_grad_max_abs']:.3e} "
-            f"| {row['f_grad_mean_abs']:.3e} "
-            f"| {row['f_grad_rel_l2']:.3e} "
-            f"| {row['g_grad_max_abs']:.3e} "
-            f"| {row['g_grad_mean_abs']:.3e} "
-            f"| {row['g_grad_rel_l2']:.3e} "
-            f"| {close} |"
-        )
-
-    return rows
+    for size in sizes:
+        for max_iter in max_iters:
+            row = reference_error_check(size=size, epsilon=epsilon, max_iter=max_iter)
+            print(
+                f"| {row['size']}x{row['size']} "
+                f"| {row['points']} "
+                f"| {row['epsilon']:.3g} "
+                f"| {row['max_iter']} "
+                f"| {row['f_grad_max_abs']:.3e} "
+                f"| {row['f_grad_rel_l2']:.3e} "
+                f"| {row['g_grad_max_abs']:.3e} "
+                f"| {row['g_grad_rel_l2']:.3e} "
+                f"| {row['f_allclose']}/{row['g_allclose']} |"
+            )
 
 
 if __name__ == "__main__":
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 53eb541..4e962ac 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -1,30 +1,22 @@
-import importlib.util
-from pathlib import Path
-
 import numpy as np
 import pytest
 import torch
 
+from torchmorph import SinkhornSolver, build_cost_matrix
+
 try:
     import ot
 except ImportError:  # pragma: no cover - exercised only when POT is absent
     ot = None
 
-# Project root directory (parent of test)
-project_root = Path(__file__).parent.parent
 
-# Target module path
-module_path = project_root / "torchmorph" / "optimal_transport.py"
-
-# Load module
-spec = importlib.util.spec_from_file_location("optimal_transport", module_path)
-tr = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(tr)
+def _device():
+    return "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def _require_cuda():
     if not torch.cuda.is_available():
-        pytest.skip("CUDA is required for Sinkhorn tests")
+        pytest.skip("CUDA is required for the CUDA Sinkhorn backends")
 
 
 def _require_pot():
@@ -37,351 +29,211 @@ def _make_positive(shape, device, seed):
     return torch.rand(shape, generator=generator, device=device) + 0.1
 
 
-def _expected_batch_channel_n(shape):
-    if len(shape) <= 2:
-        n = int(np.prod(shape))
-        return 1, 1, n
-    b, c = shape[:2]
-    n = int(np.prod(shape[2:]))
-    return b, c, n
-
-
 def _pot_plan(source, target, cost_matrix, solver):
-    source_np = source.detach().cpu().numpy()
-    target_np = target.detach().cpu().numpy()
-    cost_np = cost_matrix.detach().cpu().numpy()
     return ot.sinkhorn(
-        source_np,
-        target_np,
-        cost_np,
-        reg=1.0 / solver.reg,
-        numItermax=solver.itrstep,
+        source.detach().cpu().numpy().astype(np.float64),
+        target.detach().cpu().numpy().astype(np.float64),
+        cost_matrix.detach().cpu().numpy().astype(np.float64),
+        reg=solver.epsilon,
+        numItermax=solver.max_iter,
         stopThr=1e-9,
     )
 
 
 def test_build_cost_matrix():
-    """Simple test for the correctness of the build_cost_matrix function."""
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-    # Test 1D shape
-    shape_1d = (5,)
-    cost_1d = tr.build_cost_matrix(shape_1d, device=device, p=2)
-    N_1d = shape_1d[0]
-    assert cost_1d.shape == (
-        N_1d,
-        N_1d,
-    ), f"Expected shape {(N_1d, N_1d)}, got {cost_1d.shape}"
-    # Diagonal should be 0
-    assert torch.allclose(cost_1d.diag(), torch.zeros(N_1d, device=device)), "Diagonal not zero"
-    # Symmetry
-    assert torch.allclose(cost_1d, cost_1d.T), "Matrix not symmetric"
-    # Distance between adjacent elements should be 1
-    assert torch.allclose(
-        cost_1d[0, 1], torch.tensor(1.0, device=device)
-    ), "Adjacent distance should be 1"
-
-    # ---------- 2. Test 2D shape ----------
-    shape_2d = (3, 4)  # H=3, W=4
-    cost_2d = tr.build_cost_matrix(shape_2d, device=device, p=2)
-    N_2d = shape_2d[0] * shape_2d[1]
-    assert cost_2d.shape == (N_2d, N_2d), f"Expected shape {(N_2d, N_2d)}, got {cost_2d.shape}"
-    # Diagonal should be 0
-    assert torch.allclose(cost_2d.diag(), torch.zeros(N_2d, device=device)), "Diagonal not zero"
-    # Symmetry
-    assert torch.allclose(cost_2d, cost_2d.T), "Matrix not symmetric"
-
-    idx_00 = 0 * shape_2d[1] + 0  # Row 0, Column 0
-    idx_01 = 0 * shape_2d[1] + 1  # Row 0, Column 1
-    assert torch.allclose(
-        cost_2d[idx_00, idx_01], torch.tensor(1.0, device=device)
-    ), f"Expected distance 1.0 between adjacent columns, got {cost_2d[idx_00, idx_01]}"
-
-    # ---------- 3. Test 3D shape ----------
-    shape_3d = (2, 2, 2)  # A small 2x2x2 volume
-    cost_3d = tr.build_cost_matrix(shape_3d, device=device, p=2)
-    N_3d = 8
-    assert cost_3d.shape == (N_3d, N_3d), f"Expected shape {(N_3d, N_3d)}, got {cost_3d.shape}"
-    assert torch.allclose(cost_3d.diag(), torch.zeros(N_3d, device=device))
-    assert torch.allclose(cost_3d, cost_3d.T)
-
-    # ---------- 4. Test different p values (Manhattan distance) ----------
-    shape = (1, 1, 1, 3)  # Correctly define shape before usage
-    print(f"Input shape: {shape}")
-    cost_l1 = tr.build_cost_matrix(shape, device=device, p=1)
-    cost_l2 = tr.build_cost_matrix(shape, device=device, p=2)
-
-    expected = torch.tensor([[0.0, 1.0, 2.0], [1.0, 0.0, 1.0], [2.0, 1.0, 0.0]], device=device)
-    assert torch.allclose(cost_l1, expected), "L1 distance matrix incorrect"
-    assert torch.allclose(cost_l2, expected), "L2 distance matrix incorrect"
-
-
-@pytest.mark.parametrize(
-    "shape",
-    [
-        (4,),
-        (2, 3),
-        (2, 2, 2),
-        (1, 1, 2, 3),
-        (2, 3, 2, 2),
-        (1, 2, 2, 2, 2),
-    ],
-)
-def test_run_once_accepts_multiple_raw_dimensions(shape):
-    _require_cuda()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=300, threshold=0, device=device)
-
-    source = _make_positive(shape, device=device, seed=11)
-    target = _make_positive(shape, device=device, seed=23)
-    out = solver.run_once(
+    # 1-D grid: zero diagonal, symmetry, unit spacing
+    cost_1d = build_cost_matrix((5,))
+    assert cost_1d.shape == (5, 5)
+    assert torch.allclose(cost_1d.diag(), torch.zeros(5))
+    assert torch.allclose(cost_1d, cost_1d.T)
+    assert cost_1d[0, 1] == pytest.approx(1.0)
+
+    # 2-D grid: adjacent columns are distance 1 apart
+    cost_2d = build_cost_matrix((3, 4))
+    assert cost_2d.shape == (12, 12)
+    assert torch.allclose(cost_2d, cost_2d.T)
+    assert cost_2d[0, 1] == pytest.approx(1.0)
+
+    # p-norms differ on diagonal moves: grid points (0, 0) and (1, 1)
+    assert build_cost_matrix((2, 2), p=1)[0, 3] == pytest.approx(2.0)
+    assert build_cost_matrix((2, 2), p=2)[0, 3] == pytest.approx(2.0**0.5)
+
+
+@pytest.mark.parametrize("n,d", [(1, 4), (3, 6), (2, 16)])
+def test_solve_shapes(n, d):
+    device = _device()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=300, device=device)
+
+    source = _make_positive((n, d), device=device, seed=11)
+    target = _make_positive((n, d), device=device, seed=23)
+    out = solver.solve(
         source,
         target,
         return_plan=True,
-        return_uv=True,
+        return_potentials=True,
         return_distance=True,
     )
 
-    batch, channel, n = _expected_batch_channel_n(shape)
-    assert out["source"].shape == (batch, channel, n)
-    assert out["target"].shape == (batch, channel, n)
-    assert out["cost_matrix"].shape == (n, n)
-    assert out["plan"].shape == (batch, channel, n, n)
-    assert out["u"].shape == (batch, channel, n)
-    assert out["v"].shape == (batch, channel, n)
-    assert out["distance"].shape == (batch, channel)
+    assert out["source"].shape == (n, d)
+    assert out["target"].shape == (n, d)
+    assert out["cost_matrix"].shape == (d, d)
+    assert out["plan"].shape == (n, d, d)
+    assert out["log_u"].shape == (n, d)
+    assert out["log_v"].shape == (n, d)
+    assert out["distance"].shape == (n,)
 
     distance_from_plan = (out["plan"] * out["cost_matrix"]).sum(dim=(-2, -1))
     assert torch.allclose(out["distance"], distance_from_plan, rtol=1e-5, atol=1e-6)
+    # rows sum to the normalized marginals
+    assert torch.allclose(out["source"].sum(dim=-1), torch.ones(n, device=device))
 
 
-@pytest.mark.parametrize(
-    "shape",
-    [
-        (4,),
-        (2, 3),
-        (2, 2, 2),
-        (2, 2, 2, 2),
-    ],
-)
-def test_run_once_distance_matches_pot_sinkhorn(shape):
-    _require_cuda()
+@pytest.mark.parametrize("shape", [(4,), (2, 2, 2)])
+def test_solve_rejects_non_2d_inputs(shape):
+    solver = SinkhornSolver(epsilon=1.5, max_iter=10)
+    source = _make_positive(shape, device="cpu", seed=911)
+    target = _make_positive(shape, device="cpu", seed=919)
+
+    with pytest.raises(ValueError, match=r"\(n, d\)"):
+        solver.solve(source, target)
+
+
+def test_solve_rejects_mismatched_cost_matrix():
+    solver = SinkhornSolver(epsilon=1.5, max_iter=10)
+    source = _make_positive((2, 4), device="cpu", seed=5)
+    target = _make_positive((2, 4), device="cpu", seed=7)
+
+    with pytest.raises(ValueError, match="cost_matrix"):
+        solver.solve(source, target, cost_matrix=torch.zeros(3, 3))
+
+
+def test_solve_distance_matches_pot_sinkhorn():
     _require_pot()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
-
-    source = _make_positive(shape, device=device, seed=101)
-    target = _make_positive(shape, device=device, seed=211)
-    out = solver.run_once(source, target, return_plan=True, return_distance=True)
-
-    expected_distances = []
-    for b in range(out["source"].shape[0]):
-        channel_distances = []
-        for c in range(out["source"].shape[1]):
-            pot_transport = _pot_plan(
-                out["source"][b, c],
-                out["target"][b, c],
-                out["cost_matrix"],
-                solver,
-            )
-            pot_distance = np.sum(pot_transport * out["cost_matrix"].detach().cpu().numpy())
-            channel_distances.append(pot_distance)
-        expected_distances.append(channel_distances)
-
-    expected = torch.tensor(expected_distances, device=device, dtype=out["distance"].dtype)
+    device = _device()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, device=device)
+
+    source = _make_positive((3, 8), device=device, seed=101)
+    target = _make_positive((3, 8), device=device, seed=211)
+    out = solver.solve(source, target, return_distance=True)
+
+    cost_np = out["cost_matrix"].detach().cpu().numpy()
+    expected = torch.tensor(
+        [
+            float(np.sum(_pot_plan(out["source"][i], out["target"][i], out["cost_matrix"], solver) * cost_np))
+            for i in range(out["source"].shape[0])
+        ],
+        device=device,
+        dtype=out["distance"].dtype,
+    )
     assert torch.allclose(out["distance"], expected, rtol=5e-3, atol=1e-5)
 
 
-def test_run_once_plan_matches_pot_sinkhorn_small():
-    _require_cuda()
+def test_solve_plan_matches_pot_sinkhorn_grid_cost():
+    """A flattened 2-D image with an explicit grid cost matrix matches POT."""
     _require_pot()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+    device = _device()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, device=device)
+    cost_matrix = build_cost_matrix((2, 3), device=device)
 
-    source = _make_positive((2, 2), device=device, seed=307)
-    target = _make_positive((2, 2), device=device, seed=401)
-    out = solver.run_once(source, target, return_plan=True)
+    source = _make_positive((1, 6), device=device, seed=307)
+    target = _make_positive((1, 6), device=device, seed=401)
+    out = solver.solve(source, target, cost_matrix=cost_matrix, return_plan=True)
 
-    pot_transport = _pot_plan(out["source"][0, 0], out["target"][0, 0], out["cost_matrix"], solver)
+    pot_transport = _pot_plan(out["source"][0], out["target"][0], out["cost_matrix"], solver)
     expected = torch.tensor(pot_transport, device=device, dtype=out["plan"].dtype)
-    assert torch.allclose(out["plan"][0, 0], expected, rtol=5e-3, atol=1e-5)
+    assert torch.allclose(out["plan"][0], expected, rtol=5e-3, atol=1e-5)
 
 
-def test_run_once_cuda_matches_torch_backend_small():
+@pytest.mark.parametrize("backend", ["cuda", "cuda_log"])
+def test_cuda_backends_match_torch_backend_batched(backend):
     _require_cuda()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, device="cuda")
 
-    source = _make_positive((2, 2), device=device, seed=503)
-    target = _make_positive((2, 2), device=device, seed=607)
+    source = _make_positive((3, 6), device="cuda", seed=503)
+    target = _make_positive((3, 6), device="cuda", seed=607)
 
-    torch_out = solver.run_once(source, target, return_plan=True, return_distance=True)
-    cuda_out = solver.run_once(
-        source,
-        target,
-        use_cuda=True,
-        return_plan=True,
-        return_distance=True,
-    )
+    torch_out = solver.solve(source, target, return_plan=True, return_distance=True)
+    cuda_out = solver.solve(source, target, backend=backend, return_plan=True, return_distance=True)
 
-    assert cuda_out["backend"] == "cuda"
+    assert cuda_out["backend"] == backend
     assert torch.allclose(cuda_out["plan"], torch_out["plan"], rtol=1e-4, atol=1e-5)
     assert torch.allclose(cuda_out["distance"], torch_out["distance"], rtol=1e-4, atol=1e-5)
 
 
-def test_run_once_log_domain_cuda_matches_torch_backend_small():
+def test_log_domain_survives_small_epsilon():
+    """At this epsilon exp(-C/eps) underflows float32, so only the log backend is usable."""
     _require_cuda()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=1000, threshold=0, device=device)
-
-    source = _make_positive((2, 2), device=device, seed=701)
-    target = _make_positive((2, 2), device=device, seed=809)
+    _require_pot()
+    solver = SinkhornSolver(epsilon=0.1, max_iter=5000, device="cuda")
+    cost_matrix = build_cost_matrix((8, 8), device="cuda")
 
-    torch_out = solver.run_once(source, target, return_plan=True, return_distance=True)
-    log_out = solver.run_once(
+    source = _make_positive((1, 64), device="cuda", seed=31)
+    target = _make_positive((1, 64), device="cuda", seed=37)
+    out = solver.solve(
         source,
         target,
-        log_domain=True,
+        cost_matrix=cost_matrix,
+        backend="cuda_log",
         return_plan=True,
+        return_potentials=True,
         return_distance=True,
     )
 
-    assert log_out["backend"] == "cuda_log"
-    assert torch.allclose(log_out["plan"], torch_out["plan"], rtol=1e-4, atol=1e-5)
-    assert torch.allclose(log_out["distance"], torch_out["distance"], rtol=1e-4, atol=1e-5)
+    assert torch.isfinite(out["log_u"]).all()
+    assert torch.isfinite(out["log_v"]).all()
+    plan = out["plan"][0]
+    assert torch.isfinite(plan).all()
+    assert (plan >= 0).all()
+    assert torch.allclose(plan.sum(dim=0), out["target"][0], atol=1e-5)
+    assert torch.allclose(plan.sum(dim=1), out["source"][0], atol=5e-3)
+
+    cost_np = out["cost_matrix"].cpu().numpy().astype(np.float64)
+    pot_plan = ot.bregman.sinkhorn_log(
+        out["source"][0].cpu().numpy().astype(np.float64),
+        out["target"][0].cpu().numpy().astype(np.float64),
+        cost_np,
+        reg=solver.epsilon,
+        numItermax=solver.max_iter,
+        stopThr=1e-12,
+    )
+    pot_distance = float(np.sum(pot_plan * cost_np))
+    assert out["distance"][0].item() == pytest.approx(pot_distance, rel=2e-2)
 
 
-@pytest.mark.parametrize("kwargs", [{"use_cuda": True}, {"log_domain": True}])
-def test_cuda_backends_reject_batched_inputs(kwargs):
-    _require_cuda()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=300, threshold=0, device=device)
-
-    source = _make_positive((2, 1, 2, 2), device=device, seed=911)
-    target = _make_positive((2, 1, 2, 2), device=device, seed=919)
-
-    with pytest.raises(ValueError, match="B=C=1"):
-        solver.run_once(source, target, **kwargs)
-
-
-@pytest.mark.parametrize(
-    "shape",
-    [
-        (4,),
-        (2, 3),
-        (2, 2, 2),
-        (1, 1, 2, 3),
-        (2, 2, 2, 2),
-    ],
-)
-def test_gradient_matches_numeric_directional_derivative(shape):
-    _require_cuda()
-    device = "cuda"
-    solver = tr.SinkhornSolver(reg=0.7, itrstep=3000, threshold=0, device=device)
-
-    source_raw = _make_positive(shape, device=device, seed=1009)
-    target_raw = _make_positive(shape, device=device, seed=2003)
-    source, target, cost_matrix, _ = solver.data_preprocess(
-        source_raw,
-        target_raw,
-        force_batched=True,
-        dtype=torch.float64,
+@pytest.mark.parametrize("n,d", [(1, 4), (3, 8)])
+def test_gradient_matches_numeric_directional_derivative(n, d):
+    device = _device()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=3000, device=device)
+
+    source_raw = _make_positive((n, d), device=device, seed=1009)
+    target_raw = _make_positive((n, d), device=device, seed=2003)
+    source, target, cost_matrix = solver.data_preprocess(
+        source_raw, target_raw, dtype=torch.float64
     )
 
-    direction = _make_positive(source.shape, device=device, seed=3001).double() - 0.6
+    direction = _make_positive((n, d), device=device, seed=3001).double() - 0.6
     direction = direction - direction.mean(dim=-1, keepdim=True)
-    eps = 1e-4
-
-    def regularized_objective(x):
-        out = solver.run_once(
-            x,
-            target,
-            cost_matrix=cost_matrix,
-            return_plan=True,
-            return_distance=True,
-            dtype=torch.float64,
+    eps_fd = 1e-4
+
+    def regularized_cost(x):
+        out = solver.solve(
+            x, target, cost_matrix=cost_matrix, return_distance=True, dtype=torch.float64
         )
-        plan = out["plan"]
-        entropy_term = (plan * torch.log(torch.clamp(plan, min=1e-12))).sum()
-        return out["distance"].sum() + entropy_term / solver.reg
+        entropy = (out["plan"] * torch.log(out["plan"].clamp(min=1e-300))).sum()
+        return out["distance"].sum() + solver.epsilon * entropy
 
-    out = solver.run_once(
+    out = solver.solve(
         source,
         target,
         cost_matrix=cost_matrix,
         return_plan=False,
-        return_uv=True,
         return_grad=True,
-        return_distance=True,
         dtype=torch.float64,
     )
-    analytic_grad = out["grad_source"].reshape(-1)
-    analytic_directional = torch.sum(analytic_grad * direction.reshape(-1))
-
-    plus = source + eps * direction
-    minus = source - eps * direction
-    loss_plus = regularized_objective(plus)
-    loss_minus = regularized_objective(minus)
-    numeric_directional = (loss_plus - loss_minus) / (2 * eps)
-
-    assert torch.allclose(
-        analytic_directional,
-        numeric_directional,
-        rtol=1e-7,
-        atol=1e-7,
-    )
-
-
-def loss_test():
-    if not torch.cuda.is_available():
-        raise ValueError("CUDA not available, skipping loss test.")
-
-    device = "cuda"
-    torch.manual_seed(42)
+    analytic_directional = (out["grad_source"] * direction).sum()
 
-    reg = 5
-    itrstep = 100
-    solver = tr.SinkhornSolver(reg=reg, itrstep=itrstep, device=device)
+    loss_plus = regularized_cost(source + eps_fd * direction)
+    loss_minus = regularized_cost(source - eps_fd * direction)
+    numeric_directional = (loss_plus - loss_minus) / (2 * eps_fd)
 
-    source = torch.nn.Parameter(torch.randn(4, device=device).abs())
-    optimizer = torch.optim.Adam([source], lr=1e-2)
-
-    target = torch.rand(4, device=device).abs()
-
-    for step in range(500):
-        optimizer.zero_grad()
-
-        source_prob, target, cost_matrix, cost_matrix_T = solver.data_preprocess(source, target)
-
-        with torch.no_grad():
-            u, v = solver.sinkhorn_log_cuda(
-                source=source_prob.detach(),
-                target=target.detach(),
-                cost_matrix=cost_matrix,
-                cost_matrix_T=cost_matrix_T,
-            )
-
-            K = torch.exp(-cost_matrix * reg)
-            P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
-            loss_value = (P * cost_matrix).sum()
-
-            grad_source, _ = solver.gradient(u, v)
-
-        source_prob.backward(grad_source)
-
-        optimizer.step()
-
-        print(f"step {step}, loss = {loss_value.item():.6f}")
-        print("source:", source.detach())
-        print("target:", target)
-
-
-if __name__ == "__main__":
-    pytest.main(
-        [
-            __file__,
-            "-k",
-            "not loss_test",
-            "-q",
-        ]
-    )
+    assert torch.allclose(analytic_directional, numeric_directional, rtol=1e-5, atol=1e-7)
diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py
index 096562b..de23f1c 100644
--- a/torchmorph/__init__.py
+++ b/torchmorph/__init__.py
@@ -19,6 +19,7 @@
     morphological_laplace,
     white_tophat,
 )
+from .optimal_transport import SinkhornSolver, build_cost_matrix
 
 __all__ = [
     "euclidean_distance_transform",
@@ -38,4 +39,6 @@
     "morphological_laplace",
     "white_tophat",
     "black_tophat",
+    "SinkhornSolver",
+    "build_cost_matrix",
 ]
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
index 13c4ec5..5481268 100644
--- a/torchmorph/csrc/ot_kernel.cu
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -1,299 +1,179 @@
 #include <torch/extension.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <tuple>
-#include <cmath>
-#include <cfloat>
 
+namespace {
 
-__global__ void u_update(
-    const float* a,
-    const float* b,
-    const float* K,
-    float* u,
-    float* v,
-    int N
-){
-    int bidx = blockIdx.x;
-    int blk = blockDim.x;
-    int tid = threadIdx.x;
-    int total = N;
-    float local_sum = 0.0f;
-
-    if(bidx >= total){
-        return;
-    }
-
-    extern __shared__ float sdata[];
-
-    //get ready for sinkhorn:u= a / (k v)
-    for (int j = tid; j < N; j += blk){
-        local_sum += K[bidx*N+j] * v[j];
-    }//calculate one column
+constexpr int kThreads = 256;  // power of two, assumed by the reductions below
 
-    sdata[tid]=local_sum;//get one element of one column
-    __syncthreads();
-
-    
-    for (int stride=blk/2; stride > 0; stride /= 2){
-        if (tid < stride){
-        sdata[tid]+=sdata[tid+stride];
-        }
-        __syncthreads();
-    }
-
-    if (tid == 0){
-        u[bidx] = a[bidx] / (sdata[0]+1e-12f);
-    }
-}
-
-
-__global__ void v_update(
-    const float* a,
-    const float* b,
-    const float* K,
-    float* u,
-    float* v,
-    int N
+// One Sinkhorn scaling update. Vectors are batched (n, d); the (d, d) matrix
+// is shared across the batch. One block per (row, batch) pair computes:
+//   out[b, i] = num[b, i] / (sum_j mat[i, j] * scale[b, j] + 1e-12)
+// The u-update passes (a, K, v); the v-update passes (b, K^T, u).
+__global__ void scaling_update(
+    const float* __restrict__ num,
+    const float* __restrict__ mat,
+    const float* __restrict__ scale,
+    float* __restrict__ out,
+    int d
 ){
-    int bidx = blockIdx.x;
-    int blk = blockDim.x;
-    int tid = threadIdx.x;
-    int total = N;
-    float local_sum = 0.0f;
-
-    if(bidx >= total){
-        return;
-    }
-
-    extern __shared__ float sdata[];
-
-    //get ready for sinkhorn:u= b / (k^t u)
-    for (int j = tid; j < N; j += blk){
-        local_sum += K[bidx+j*N] * u[j];
-    }//calculate one column
-
-    sdata[tid]=local_sum;//get one element of one column
-    __syncthreads();
-
-    
-    for (int stride=blk/2; stride > 0; stride /= 2){
-        if (tid < stride){
-            sdata[tid]+=sdata[tid+stride];
-        }
-        __syncthreads();
-    }
-
-    if (tid == 0){
-        v[bidx] = b[bidx] / (sdata[0]+1e-12f);
-    }
-}
-
-
-// Update one log-domain Sinkhorn scaling vector.
-//
-// Each block owns one row i and computes:
-//   log_u[i] = log_a[i] - logsumexp_j(-reg * M[i, j] + log_v[j])
-// The logsumexp is evaluated with a two-pass block reduction:
-// first max(z_j), then sum(exp(z_j - max)).
-__global__ void u_update_log(
-    const float* log_a,
-    const float* M,
-    float* log_u,
-    const float* log_v,
-    int N,
-    float reg
-){
-    //sinkhorn for log domain iteration
-    int blk = blockDim.x;
-    int bidx = blockIdx.x;
-    int tid = threadIdx.x;
-    int total = N;
-    float local_max = -INFINITY;
-
-    if(bidx >= total){
-        return;
-    }
-
     extern __shared__ float sdata[];
+    const int row = blockIdx.x;
+    const int tid = threadIdx.x;
+    const size_t offset = (size_t)blockIdx.y * d;
 
-    //generate each log_kv
-    for (int j = tid; j < N; j += blk){
-        float log_k = - reg * M[bidx*N+j];
-        if(local_max <= (log_k + log_v[j])){
-            local_max = log_k + log_v[j];
-        }     
-    }
-
-    sdata[tid] = local_max;
-    __syncthreads();
-
-    //first reduction to find max(log_k+log_v)
-    for (int stride=blk/2; stride > 0; stride /= 2){
-        if (tid < stride){
-            sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);//get larger num
-        }
-        __syncthreads();
-    }
-
-    float row_max = sdata[0];
-    __syncthreads();
     float local_sum = 0.0f;
-
-    //log-domain iteration log_u = log_a - LSE(log_k+log_v)
-    for (int j = tid; j < N; j += blk){
-        float log_k = - reg * M[bidx*N+j];
-        local_sum += expf(log_k + log_v[j] - row_max);
+    for (int j = tid; j < d; j += blockDim.x){
+        local_sum += mat[(size_t)row * d + j] * scale[offset + j];
     }
-
     sdata[tid] = local_sum;
     __syncthreads();
-
-    for (int stride=blk/2; stride > 0; stride /= 2){
-        if (tid < stride){
-            sdata[tid] += sdata[tid + stride];
-        }
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2){
+        if (tid < stride) sdata[tid] += sdata[tid + stride];
         __syncthreads();
     }
 
-    if (tid == 0){
-        log_u[bidx] = log_a[bidx] - row_max - logf(sdata[0]);
-    }
-
+    if (tid == 0) out[offset + row] = num[offset + row] / (sdata[0] + 1e-12f);
 }
 
-__global__ void v_update_log(
-    const float* log_b,
-    const float* M_T,
-    const float* log_u,
-    float* log_v,
-    int N,
-    float reg
+// One log-domain Sinkhorn update, one block per (row, batch) pair:
+//   log_out[b, i] = log_num[b, i] - logsumexp_j(-cost[i, j] / eps + log_scale[b, j])
+// evaluated with a two-pass block reduction (max, then sum of shifted exps).
+// The u-update passes (log_a, M, log_v); the v-update passes (log_b, M^T, log_u).
+__global__ void log_scaling_update(
+    const float* __restrict__ log_num,
+    const float* __restrict__ cost,
+    const float* __restrict__ log_scale,
+    float* __restrict__ log_out,
+    int d,
+    float inv_eps
 ){
-    //sinkhorn for log domain iteration
-    int blk = blockDim.x;
-    int bidx = blockIdx.x;
-    int tid = threadIdx.x;
-    int total = N;
-    float local_max = -INFINITY;
-
-    if(bidx >= total){
-        return;
-    }
-
     extern __shared__ float sdata[];
+    const int row = blockIdx.x;
+    const int tid = threadIdx.x;
+    const size_t offset = (size_t)blockIdx.y * d;
 
-    //generate each log_kTu
-    for (int j = tid; j < N; j += blk){
-        float log_k = - reg * M_T[bidx*N+j];
-        if(local_max <= (log_k + log_u[j])){
-            local_max = log_k + log_u[j];
-        }     
+    float local_max = -INFINITY;
+    for (int j = tid; j < d; j += blockDim.x){
+        local_max = fmaxf(local_max, log_scale[offset + j] - cost[(size_t)row * d + j] * inv_eps);
     }
-
     sdata[tid] = local_max;
     __syncthreads();
-
-    //first reduction to find max(log_k+log_v)
-    for (int stride=blk/2; stride > 0; stride /= 2){
-        if (tid < stride){
-            sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);//get larger num
-        }
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2){
+        if (tid < stride) sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);
         __syncthreads();
     }
-
-    float row_max = sdata[0];
+    const float row_max = sdata[0];
     __syncthreads();
-    float local_sum = 0.0f;
 
-    //log-domain iteration log_v = log_b - LSE(log_k+log_u)
-    for (int j = tid; j < N; j += blk){
-        float log_k = - reg * M_T[bidx*N+j];
-        local_sum += expf(log_k + log_u[j] - row_max);
+    float local_sum = 0.0f;
+    for (int j = tid; j < d; j += blockDim.x){
+        local_sum += expf(log_scale[offset + j] - cost[(size_t)row * d + j] * inv_eps - row_max);
     }
-
     sdata[tid] = local_sum;
     __syncthreads();
-
-    for (int stride=blk/2; stride > 0; stride /= 2){
-        if (tid < stride){
-            sdata[tid] += sdata[tid + stride];
-        }
+    for (int stride = blockDim.x / 2; stride > 0; stride /= 2){
+        if (tid < stride) sdata[tid] += sdata[tid + stride];
         __syncthreads();
     }
 
+    // row_max == -inf means the opposite marginal is all zero; pin the
+    // potential instead of propagating -inf - (-inf) = NaN.
     if (tid == 0){
-        log_v[bidx] = log_b[bidx] - row_max - logf(sdata[0]);
+        log_out[offset + row] = row_max == -INFINITY
+            ? -INFINITY
+            : log_num[offset + row] - row_max - logf(sdata[0]);
     }
+}
 
+void check_input(const torch::Tensor& t, int64_t numel, const char* name){
+    TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+    TORCH_CHECK(t.scalar_type() == torch::kFloat32, name, " must be float32");
+    TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+    TORCH_CHECK(t.numel() == numel, name, " must have ", numel, " elements, got ", t.numel());
 }
 
+// Validate batched (n, d) vectors against a shared (d, d) matrix and return
+// the launch grid: one block per (row, batch) pair.
+dim3 make_grid(const torch::Tensor& u, const torch::Tensor& mat){
+    TORCH_CHECK(mat.dim() == 2 && mat.size(0) == mat.size(1), "cost/kernel matrix must be square");
+    TORCH_CHECK(u.dim() == 2 && u.size(1) == mat.size(0),
+                "scaling vectors must be (n, d) with d matching the matrix");
+    return dim3(static_cast<unsigned>(u.size(1)), static_cast<unsigned>(u.size(0)));
+}
+
+}  // namespace
+
 std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
-    const torch::Tensor source,
-    const torch::Tensor target,
-    const torch::Tensor k,
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& K,
+    const torch::Tensor& K_T,
     torch::Tensor u,
     torch::Tensor v,
-    int itrstep,
-    int N
+    int64_t n_iter
 ){
-    int threads = 256;
-    int blocks = N;
-    int shared_memory = threads * sizeof(float);
-
-    for (int i = 0; i < itrstep; i++){
-        u_update<<<blocks, threads, shared_memory>>>(
-            source.data_ptr<float>(),
-            target.data_ptr<float>(),
-            k.data_ptr<float>(),
-            u.data_ptr<float>(),
-            v.data_ptr<float>(),
-            N
-        );
-        v_update<<<blocks, threads, shared_memory>>>(
-            source.data_ptr<float>(),
-            target.data_ptr<float>(),
-            k.data_ptr<float>(),
-            u.data_ptr<float>(),
-            v.data_ptr<float>(),
-            N
-        );
-    }
-return std::make_tuple(u, v);
+    const dim3 grid = make_grid(u, K);
+    const int64_t d = K.size(0);
+    check_input(a, u.numel(), "a");
+    check_input(b, u.numel(), "b");
+    check_input(K, d * d, "K");
+    check_input(K_T, d * d, "K_T");
+    check_input(u, u.numel(), "u");
+    check_input(v, u.numel(), "v");
+
+    const at::cuda::CUDAGuard guard(u.device());
+    const auto stream = at::cuda::getCurrentCUDAStream();
+    const int shmem = kThreads * sizeof(float);
+
+    for (int64_t i = 0; i < n_iter; i++){
+        scaling_update<<<grid, kThreads, shmem, stream>>>(
+            a.data_ptr<float>(), K.data_ptr<float>(), v.data_ptr<float>(),
+            u.data_ptr<float>(), static_cast<int>(d));
+        scaling_update<<<grid, kThreads, shmem, stream>>>(
+            b.data_ptr<float>(), K_T.data_ptr<float>(), u.data_ptr<float>(),
+            v.data_ptr<float>(), static_cast<int>(d));
+    }
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    return std::make_tuple(u, v);
 }
 
 std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
-    const torch::Tensor log_a,
-    const torch::Tensor log_b,
-    const torch::Tensor M,
-    const torch::Tensor M_T,
+    const torch::Tensor& log_a,
+    const torch::Tensor& log_b,
+    const torch::Tensor& M,
+    const torch::Tensor& M_T,
     torch::Tensor log_u,
     torch::Tensor log_v,
-    int itrstep,
-    int N,
-    float reg
+    int64_t n_iter,
+    double epsilon
 ){
-    int threads = 256;
-    int blocks = N;
-    int shared_memory = threads * sizeof(float);
-
-    for (int i = 0; i < itrstep; i++){
-        u_update_log<<<blocks, threads, shared_memory>>>(
-            log_a.data_ptr<float>(),
-            M.data_ptr<float>(),
-            log_u.data_ptr<float>(),
-            log_v.data_ptr<float>(),
-            N,
-            reg
-        );
-        v_update_log<<<blocks, threads, shared_memory>>>(
-            log_b.data_ptr<float>(),
-            M_T.data_ptr<float>(),
-            log_u.data_ptr<float>(),
-            log_v.data_ptr<float>(),
-            N,
-            reg
-        );
-    }
-return std::make_tuple(log_u, log_v);
+    TORCH_CHECK(epsilon > 0, "epsilon must be positive, got ", epsilon);
+    const dim3 grid = make_grid(log_u, M);
+    const int64_t d = M.size(0);
+    check_input(log_a, log_u.numel(), "log_a");
+    check_input(log_b, log_u.numel(), "log_b");
+    check_input(M, d * d, "M");
+    check_input(M_T, d * d, "M_T");
+    check_input(log_u, log_u.numel(), "log_u");
+    check_input(log_v, log_u.numel(), "log_v");
+
+    const at::cuda::CUDAGuard guard(log_u.device());
+    const auto stream = at::cuda::getCurrentCUDAStream();
+    const int shmem = kThreads * sizeof(float);
+    const float inv_eps = static_cast<float>(1.0 / epsilon);
+
+    for (int64_t i = 0; i < n_iter; i++){
+        log_scaling_update<<<grid, kThreads, shmem, stream>>>(
+            log_a.data_ptr<float>(), M.data_ptr<float>(), log_v.data_ptr<float>(),
+            log_u.data_ptr<float>(), static_cast<int>(d), inv_eps);
+        log_scaling_update<<<grid, kThreads, shmem, stream>>>(
+            log_b.data_ptr<float>(), M_T.data_ptr<float>(), log_u.data_ptr<float>(),
+            log_v.data_ptr<float>(), static_cast<int>(d), inv_eps);
+    }
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    return std::make_tuple(log_u, log_v);
 }
diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp
index 318fda7..669c56d 100644
--- a/torchmorph/csrc/torchmorph.cpp
+++ b/torchmorph/csrc/torchmorph.cpp
@@ -43,30 +43,26 @@ std::tuple<torch::Tensor, torch::Tensor> bfdt_cuda(
 
 // Optimal Transport functions
 std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
-    const torch::Tensor source,
-    const torch::Tensor target,
-    const torch::Tensor k,
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& K,
+    const torch::Tensor& K_T,
     torch::Tensor u,
     torch::Tensor v,
-    int itrstep,
-    int N
+    int64_t n_iter
 );
 
 std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
-    const torch::Tensor log_a,
-    const torch::Tensor log_b,
-    const torch::Tensor M,
-    const torch::Tensor M_T,
+    const torch::Tensor& log_a,
+    const torch::Tensor& log_b,
+    const torch::Tensor& M,
+    const torch::Tensor& M_T,
     torch::Tensor log_u,
     torch::Tensor log_v,
-    int itrstep,
-    int N,
-    float reg
+    int64_t n_iter,
+    double epsilon
 );
 
-
-
-
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("grey_erosion_cuda", &grey_erosion_cuda,
           "N-dimensional fused grey erosion",
@@ -110,26 +106,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
 
     // Optimal Transport
     m.def("sinkhorn_fastiter", &sinkhorn_fastiter,
-          "Sinkhorn Fast Iteration CUDA",
-          py::arg("source"),
-          py::arg("target"),
-          py::arg("k"),
+          "Sinkhorn scaling-form iterations (CUDA)",
+          py::arg("a"),
+          py::arg("b"),
+          py::arg("K"),
+          py::arg("K_T"),
           py::arg("u"),
           py::arg("v"),
-          py::arg("itrstep"),
-          py::arg("N"));
+          py::arg("n_iter"));
 
     m.def("sinkhorn_logiter", &sinkhorn_logiter,
-          "Sinkhorn Log-Domain Iteration CUDA",
+          "Sinkhorn log-domain iterations (CUDA)",
           py::arg("log_a"),
           py::arg("log_b"),
           py::arg("M"),
           py::arg("M_T"),
           py::arg("log_u"),
           py::arg("log_v"),
-          py::arg("itrstep"),
-          py::arg("N"),
-          py::arg("reg")
-    );
+          py::arg("n_iter"),
+          py::arg("epsilon"));
 }
 
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 0f6e847..080f5e1 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -4,408 +4,177 @@
 from torch import Tensor
 
 
-def build_cost_matrix(shape, p=2, device='cuda') -> torch.Tensor:  # Euclidean distance by default
-    """Build the pairwise ground-cost matrix for a spatial grid.
+def build_cost_matrix(shape, p=2, device=None) -> Tensor:
+    """Pairwise L^p distance matrix between the points of a spatial grid.
 
-    The shape defines grid coordinates, and p selects the distance norm.
-    Returns an (N, N) matrix for N flattened grid points.
+    Returns a (d, d) cost matrix for the d flattened grid points of `shape`.
+    Flatten grid-shaped distributions to rows of an (n, d) tensor and pass the
+    result as `cost_matrix` to :class:`SinkhornSolver`.
     """
-    coords = torch.stack(torch.meshgrid([torch.arange(s) for s in shape], indexing="ij"), dim=-1)
-    coords_flat = coords.view(-1, coords.size(-1)).float()
-
-    cost_matrix = torch.cdist(coords_flat, coords_flat, p=p).to(device)
-    return cost_matrix
+    coords = torch.stack(
+        torch.meshgrid([torch.arange(s, device=device) for s in shape], indexing="ij"), dim=-1
+    )
+    coords = coords.reshape(-1, len(shape)).float()
+    return torch.cdist(coords, coords, p=p)
 
 
 class SinkhornSolver:
-    def __init__(self, reg=0.02, itrstep=100, threshold=0, p=2, device='cuda'):
-        self.reg = reg
-        self.itrstep = itrstep
+    """Entropy-regularized balanced optimal transport between batched histograms.
+
+    Every method takes source and target as (n, d) tensors: n distributions of
+    d bins each, sharing one (d, d) cost matrix. The standard convention
+    K = exp(-C / epsilon) is used: smaller epsilon means weaker blurring and a
+    sharper transport plan. All backends return the log scaling vectors
+    (log_u, log_v); the dual potentials are f = epsilon * log_u and
+    g = epsilon * log_v.
+    """
+
+    def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, device=None):
+        if epsilon <= 0:
+            raise ValueError("epsilon must be positive.")
+        if max_iter <= 0:
+            raise ValueError("max_iter must be positive.")
+        if threshold < 0:
+            raise ValueError("threshold must be non-negative.")
+        self.epsilon = epsilon
+        self.max_iter = max_iter
         self.threshold = threshold
-        self.device = device
         self.p = p
+        self.device = device
 
     def data_preprocess(
         self,
         source: Tensor,
         target: Tensor,
         cost_matrix: Optional[Tensor] = None,
-        force_batched: bool = False,
         dtype=torch.float32,
     ):
-        """Preprocess source and target distributions for Sinkhorn solvers.
+        """Validate (n, d) inputs, clamp negatives, and normalize each row to unit mass.
 
-        This checks matching shapes, flattens spatial axes, clamps negatives,
-        normalizes mass, and prepares the cost matrix transpose.
+        When `cost_matrix` is None the d bins are treated as points on a line.
         """
-        # check dimensions,reshape and choose strategy
-        if not source.ndim == target.ndim or not source.shape == target.shape:
-            raise ValueError(
-                "Source and target must have the same number of dimensions and the same shape."
-            )
-        elif source.ndim <= 2:
-            if cost_matrix is None:
-                cost_matrix = build_cost_matrix(source.shape, self.p, self.device)
-            if force_batched:
-                source = source.reshape(1, 1, -1)
-                target = target.reshape(1, 1, -1)
-            else:
-                source = source.reshape(-1)
-                target = target.reshape(-1)
-        elif source.ndim > 2:
-            # cut and reshape source and target to coordinate tensors, then calculate cost matrix
-            B, C = source.shape[:2]  # get batch and channel dimensions
-            if cost_matrix is None:
-                spatial_shape = source.shape[2:]
-                cost_matrix = build_cost_matrix(spatial_shape, self.p, self.device)
-            source = source.reshape(B, C, -1)
-            target = target.reshape(B, C, -1)
-
-        # move to specified device(cuda) and ensure dtype float32
-        source = source.to(device=self.device, dtype=dtype)
-        target = target.to(device=self.device, dtype=dtype)
-        cost_matrix = cost_matrix.to(device=self.device, dtype=dtype)
-
-        # ensure non-negativity
-        source = torch.clamp(source, min=0)
-        target = torch.clamp(target, min=0)
-
-        # normalize source and target to make them valid probability distributions
-        source /= torch.clamp(source.sum(dim=-1, keepdim=True).to(self.device), min=1e-12)
-        target /= torch.clamp(target.sum(dim=-1, keepdim=True).to(self.device), min=1e-12)
-
-        if not torch.allclose(cost_matrix, cost_matrix.transpose(-2, -1)):
-            cost_matrix_T = cost_matrix.transpose(-2, -1).contiguous()
-        else:
-            cost_matrix_T = cost_matrix
-
-        return source, target, cost_matrix, cost_matrix_T
-
-    def sinkhorn(
-        self,
-        source: Tensor,
-        target: Tensor,
-        cost_matrix: Optional[Tensor] = None,
-        returngrad: bool = False,
-        verbose: bool = False,
-    ):
-        """Compute a balanced Sinkhorn transport plan from raw inputs.
-
-        The function preprocesses inputs, runs the batched solver, and returns
-        either the plan or auxiliary variables when gradients are requested.
-        """
-        source, target, cost_matrix, _ = self.data_preprocess(
-            source,
-            target,
-            cost_matrix=cost_matrix,
-            force_batched=True,
-        )
-        result = self.sinkhorn_batch(
-            source, target, cost_matrix=cost_matrix, returnuv=returngrad, verbose=verbose
-        )
-
-        if returngrad:
-            return result
-        return result["plan"]
-
-    def run_once(
+        if source.ndim != 2 or source.shape != target.shape:
+            raise ValueError("source and target must be (n, d) tensors of the same shape.")
+        device = torch.device(self.device) if self.device is not None else source.device
+        d = source.shape[1]
+        if cost_matrix is None:
+            cost_matrix = build_cost_matrix((d,), self.p, device)
+        elif cost_matrix.shape != (d, d):
+            raise ValueError(f"cost_matrix must have shape ({d}, {d}), got {tuple(cost_matrix.shape)}.")
+
+        source = source.to(device=device, dtype=dtype).clamp(min=0)
+        target = target.to(device=device, dtype=dtype).clamp(min=0)
+        cost_matrix = cost_matrix.to(device=device, dtype=dtype).contiguous()
+        source = source / source.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+        target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+        return source, target, cost_matrix
+
+    def solve(
         self,
         source: Tensor,
         target: Tensor,
         cost_matrix: Optional[Tensor] = None,
         *,
-        use_cuda: bool = False,
-        log_domain: bool = False,
+        backend: str = "torch",
         return_plan: bool = True,
-        return_uv: bool = False,
+        return_potentials: bool = False,
         return_grad: bool = False,
         return_distance: bool = False,
-        verbose: bool = False,
         dtype=torch.float32,
     ):
-        """Run the full Sinkhorn workflow from raw inputs.
+        """Run Sinkhorn end-to-end on (n, d) source and target distributions.
 
-        This wrapper preprocesses raw source and target tensors, selects one
-        Sinkhorn backend, optionally reconstructs the transport plan, and
-        optionally returns the OT distance, dual-scaling vectors, or gradients.
+        backend is "torch" (batched matmuls, any device), "cuda" (scaling-form
+        kernel), or "cuda_log" (log-domain kernel, stable for small epsilon).
+        The (n, d, d) plan is always reconstructed in log space.
         """
-        source, target, cost_matrix, cost_matrix_T = self.data_preprocess(
-            source,
-            target,
-            cost_matrix=cost_matrix,
-            force_batched=True,
-            dtype=dtype,
-        )
-
-        need_uv = return_uv or return_grad or log_domain
-
-        if log_domain:
-            u, v = self.sinkhorn_log_cuda(
-                source=source,
-                target=target,
-                cost_matrix=cost_matrix,
-                cost_matrix_T=cost_matrix_T,
-                dtype=dtype,
-            )
-            backend = "cuda_log"
-            result = {"u": u, "v": v}
-        elif use_cuda:
-            result = self.sinkhorn_cuda(
-                source=source,
-                target=target,
-                cost_matrix=cost_matrix,
-                returnuv=need_uv,
-            )
-            backend = "cuda"
-        else:
-            result = self.sinkhorn_batch(
-                source=source,
-                target=target,
-                cost_matrix=cost_matrix,
-                returnuv=need_uv,
-                verbose=verbose,
-            )
-            backend = "torch"
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix, dtype)
+        backends = {
+            "torch": self.sinkhorn_torch,
+            "cuda": self.sinkhorn_cuda,
+            "cuda_log": self.sinkhorn_log_cuda,
+        }
+        if backend not in backends:
+            raise ValueError(f"Unknown backend {backend!r}; expected one of {sorted(backends)}.")
+        log_u, log_v = backends[backend](source, target, cost_matrix)
 
         output = {
             "source": source,
             "target": target,
             "cost_matrix": cost_matrix,
-            "cost_matrix_T": cost_matrix_T,
             "backend": backend,
         }
-
-        plan = None
+        if return_potentials:
+            output["log_u"], output["log_v"] = log_u, log_v
         if return_plan or return_distance:
-            if "plan" in result:
-                plan = result["plan"]
-            else:
-                k = torch.exp(-cost_matrix * self.reg)
-                plan = result["u"].unsqueeze(-1) * k * result["v"].unsqueeze(-2)
-
-        if return_plan:
-            output["plan"] = plan
-
-        if return_distance:
-            output["distance"] = (plan * cost_matrix).sum(dim=(-2, -1))
-
-        if return_uv or return_grad:
-            output["u"] = result["u"]
-            output["v"] = result["v"]
-
+            plan = torch.exp(
+                log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2)
+            )
+            if return_plan:
+                output["plan"] = plan
+            if return_distance:
+                output["distance"] = (plan * cost_matrix).sum(dim=(-2, -1))
         if return_grad:
-            grad_source, grad_target = self.gradient(result["u"], result["v"])
-            output["grad_source"] = grad_source
-            output["grad_target"] = grad_target
-
+            output["grad_source"], output["grad_target"] = self.gradient(log_u, log_v)
         return output
 
-    def sinkhorn_batch(
-        self,
-        source: Tensor,  # (B,C,*Spatial)
-        target: Tensor,  # (B,C,*Spatial)
-        cost_matrix: Tensor,
-        returnuv: bool = False,  # whether to return the gradient of the dual potential
-        verbose: bool = False,
-    ):
-        """Compute batched balanced Sinkhorn iterations in PyTorch.
-
-        Inputs should be preprocessed to shape (B, C, N). The solver builds
-        K = exp(-reg * C), updates u and v, and returns the transport plan.
-        """
-        # Device check
-        if not torch.cuda.is_available():
-            raise ValueError("CUDA device not available.")
-
-        if source.ndim != 3 or target.ndim != 3:
-            raise ValueError("source and target must be preprocessed to shape (B, C, N).")
-
-        # initialize u and v, and expand dimensions for batch and channel
-        u = torch.ones_like(source, device=self.device)
-        v = torch.ones_like(target, device=self.device)
-        k = torch.exp(-cost_matrix * self.reg)
-
-        B, C = source.shape[:2]
-        k = k.unsqueeze(0).unsqueeze(0).expand(B, C, -1, -1)
-
-        if self.reg <= 0:
-            raise ValueError("Regularization parameter must be positive.")
+    def sinkhorn_torch(self, source: Tensor, target: Tensor, cost_matrix: Tensor):
+        """Batched Sinkhorn in scaling form with dense matmuls.
 
-        # generate sinkhorn iteration judged by threshold or iteration times
-        if self.itrstep < 0 or self.threshold < 0:
-            raise ValueError("iterstep and threshold must be non-negative.")
-        elif self.itrstep == 0 and self.threshold == 0:
-            raise ValueError("Invalid iteration or step settings. ")
-
-        itrstep = self.itrstep
-        if self.threshold > 0 and self.itrstep == 0:
-            itrstep = 500
-
-        # iteration
-        convergence_flag = False
-        convergence_check = 10
-        for i in range(itrstep):
-            u_old = u
-            u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
-            v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
-            if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < self.threshold:
-                convergence_flag = True
-                break
-
-        # use broadcasting to calculate the optimal transport plan P from u, v and K
-        P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
-
-        if verbose:
-            if convergence_flag:
-                print(
-                    f"Sinkhorn iteration completed in {i+1} steps.Convergence: {convergence_flag}"
-                )
-            else:
-                print(f"Sinkhorn iteration completed in {i+1} steps,Covergence didn't complete.")
-
-        # calculate gradient after iteration covergence ,return results
-        if returnuv:
-            return {"plan": P, "u": u, "v": v}
-        else:
-            return {"plan": P}
-
-    def sinkhorn__nobatch(
-        self,
-        source: Tensor,  # (N,)
-        target: Tensor,  # (N,)
-        cost_matrix: Tensor,
-        returnuv: bool = False,  # whether to return the gradient of the dual potential
-    ):
-        """Compute Sinkhorn iterations for one flat distribution pair.
-
-        Source and target should be one-dimensional probability vectors.
-        This path is a simple PyTorch reference for the non-batched case.
+        Expects preprocessed (n, d) inputs; returns (log_u, log_v).
         """
-
-        # initialize u and v, and expand dimensions for batch and channel
-        u = torch.ones_like(source, device=self.device)
-        v = torch.ones_like(target, device=self.device)
-        k = torch.exp(-cost_matrix * self.reg)
-
-        if self.reg <= 0:
-            raise ValueError("Regularization parameter must be positive.")
-
-        # generate sinkhorn iteration judged by threshold or iteration times
-        if self.itrstep < 0 or self.threshold < 0:
-            raise ValueError("iterstep and threshold must be non-negative.")
-        elif self.itrstep == 0 and self.threshold == 0:
-            raise ValueError("Invalid iteration or step settings. ")
-
-        itrstep = self.itrstep
-        if self.threshold > 0 and self.itrstep == 0:
-            itrstep = 500
-
-        # iteration
-        convergence_check = 10
-        for i in range(itrstep):
-            u_old = u
-            u = source / (torch.matmul(k, v.unsqueeze(-1)).squeeze(-1) + 1e-12)
-            v = target / (torch.matmul(k.transpose(-2, -1), u.unsqueeze(-1)).squeeze(-1) + 1e-12)
-            if i % convergence_check == 0 and (u_old - u).abs().sum().mean() < self.threshold:
-                break
-
-        # use broadcasting to calculate the optimal transport plan P from u, v and K
-        P = u.unsqueeze(-1) * k * v.unsqueeze(-2)
-
-        # calculate gradient after iteration covergence ,return results
-        if returnuv:
-            return {"plan": P, "u": u, "v": v}
-        else:
-            return {"plan": P}
-
-    def gradient(self, u: Tensor, v: Tensor):
-        """Convert Sinkhorn scaling vectors to source and target gradients.
-
-        The scaling vectors are logged, centered to remove additive constants,
-        and rescaled according to this module's regularization convention.
-        """
-        u = torch.log(torch.clamp(u, min=1e-12))
-        v = torch.log(torch.clamp(v, min=1e-12))
-
-        u = u - u.mean(dim=-1, keepdim=True)  # remove constants
-        v = v - v.mean(dim=-1, keepdim=True)
-
-        grad_source = u / self.reg
-        grad_target = v / self.reg
-
-        return grad_source, grad_target
-
-    def _validate_single_cuda_distribution(self, source: Tensor, target: Tensor):
-        if source.ndim != 3 or target.ndim != 3:
-            raise ValueError("CUDA Sinkhorn inputs must be preprocessed to shape (B, C, N).")
-        if source.shape[:2] != (1, 1) or target.shape[:2] != (1, 1):
-            raise ValueError("CUDA Sinkhorn kernels currently support only B=C=1 inputs.")
-
-    def sinkhorn_cuda(
-        self,
-        source: Tensor,
-        target: Tensor,
-        cost_matrix: Tensor,
-        returnuv: bool = False,
-    ):
-        """Run the CUDA Sinkhorn solver in standard scaling form.
-
-        The CUDA kernel updates u and v, after which the Python wrapper
-        reconstructs the transport plan from u, K, and v.
-        """
-        self._validate_single_cuda_distribution(source, target)
+        K = torch.exp(-cost_matrix / self.epsilon)
+        u = torch.ones_like(source)
+        v = torch.ones_like(target)
+        for i in range(self.max_iter):
+            u_prev = u
+            u = source / (v @ K.mT + 1e-12)
+            v = target / (u @ K + 1e-12)
+            if self.threshold > 0 and i % 10 == 0:
+                if (u - u_prev).abs().sum(dim=-1).max() <= self.threshold:
+                    break
+        return u.log(), v.log()
+
+    @torch.no_grad()
+    def sinkhorn_cuda(self, source: Tensor, target: Tensor, cost_matrix: Tensor):
+        """Scaling-form Sinkhorn on the fused CUDA kernel; returns (log_u, log_v)."""
         from torchmorph import _C
 
-        N = source.shape[-1]
-        K = torch.exp(-cost_matrix * self.reg)
+        K = torch.exp(-cost_matrix / self.epsilon)
         u = torch.ones_like(source)
         v = torch.ones_like(target)
+        u, v = _C.sinkhorn_fastiter(
+            source.contiguous(), target.contiguous(), K, K.mT.contiguous(), u, v, self.max_iter
+        )
+        return u.log(), v.log()
 
-        u, v = _C.sinkhorn_fastiter(source, target, K, u, v, int(self.itrstep), int(N))
-
-        P = u.unsqueeze(-1) * K * v.unsqueeze(-2)
-
-        if returnuv:
-            return {"plan": P, "u": u, "v": v}
-        else:
-            return {"plan": P}
+    @torch.no_grad()
+    def sinkhorn_log_cuda(self, source: Tensor, target: Tensor, cost_matrix: Tensor):
+        """Log-domain Sinkhorn on the fused CUDA kernel; stable for small epsilon.
 
-    def sinkhorn_log_cuda(
-        self,
-        source: Tensor,
-        target: Tensor,
-        cost_matrix: Tensor,
-        cost_matrix_T: Tensor,
-        dtype=torch.float32,
-    ):
-        """Run the CUDA Sinkhorn solver in log-domain form.
-
-        Inputs are converted to log probabilities for numerical stability.
-        The returned u and v are exponentiated scaling vectors.
+        Returns (log_u, log_v) without ever leaving log space.
         """
-        self._validate_single_cuda_distribution(source, target)
         from torchmorph import _C
 
-        with torch.no_grad():
-            source = source.to(dtype=dtype).log()
-            target = target.to(dtype=dtype).log()
-            N = source.shape[-1]
-            u = torch.zeros_like(source)
-            v = torch.zeros_like(target)
-
-            u, v = _C.sinkhorn_logiter(
-                source,
-                target,
-                cost_matrix,
-                cost_matrix_T,
-                u,
-                v,
-                int(self.itrstep),
-                int(N),
-                float(self.reg),
-            )
+        log_u = torch.zeros_like(source)
+        log_v = torch.zeros_like(target)
+        return _C.sinkhorn_logiter(
+            source.log().contiguous(),
+            target.log().contiguous(),
+            cost_matrix,
+            cost_matrix.mT.contiguous(),
+            log_u,
+            log_v,
+            self.max_iter,
+            self.epsilon,
+        )
 
-            u = u.exp()
-            v = v.exp()
+    def gradient(self, log_u: Tensor, log_v: Tensor):
+        """Gradients of the entropic OT cost w.r.t. the marginals.
 
-        return u, v
+        These are the dual potentials f = epsilon * log_u and g = epsilon * log_v,
+        centered per distribution to fix the additive gauge freedom.
+        """
+        f = self.epsilon * log_u
+        g = self.epsilon * log_v
+        return f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True)

From 49d19a08f7997109516c8f03978a3008674b910b Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Wed, 8 Jul 2026 14:05:32 +0800
Subject: [PATCH 16/20] Apply black formatting to optimal transport files

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 benchmark/optimal_transport.py  | 4 +++-
 test/test_optimal_transport.py  | 7 ++++++-
 torchmorph/optimal_transport.py | 8 ++++----
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index 4d1086f..b9afa33 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -121,7 +121,9 @@ def reference_error_check(size=32, epsilon=10.0, max_iter=200, verbose=False):
     }
 
 
-def print_reference_error_table(sizes=(2, 4, 8, 16, 32, 64), max_iters=(100, 200, 500), epsilon=10.0):
+def print_reference_error_table(
+    sizes=(2, 4, 8, 16, 32, 64), max_iters=(100, 200, 500), epsilon=10.0
+):
     """Print a size and iteration sweep for POT-vs-CUDA gradient error."""
     headers = ("size", "N", "eps", "itr", "f_max", "f_rel_l2", "g_max", "g_rel_l2", "close")
     print("| " + " | ".join(headers) + " |")
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 4e962ac..5363d58 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -119,7 +119,12 @@ def test_solve_distance_matches_pot_sinkhorn():
     cost_np = out["cost_matrix"].detach().cpu().numpy()
     expected = torch.tensor(
         [
-            float(np.sum(_pot_plan(out["source"][i], out["target"][i], out["cost_matrix"], solver) * cost_np))
+            float(
+                np.sum(
+                    _pot_plan(out["source"][i], out["target"][i], out["cost_matrix"], solver)
+                    * cost_np
+                )
+            )
             for i in range(out["source"].shape[0])
         ],
         device=device,
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 080f5e1..e249a8f 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -60,7 +60,9 @@ def data_preprocess(
         if cost_matrix is None:
             cost_matrix = build_cost_matrix((d,), self.p, device)
         elif cost_matrix.shape != (d, d):
-            raise ValueError(f"cost_matrix must have shape ({d}, {d}), got {tuple(cost_matrix.shape)}.")
+            raise ValueError(
+                f"cost_matrix must have shape ({d}, {d}), got {tuple(cost_matrix.shape)}."
+            )
 
         source = source.to(device=device, dtype=dtype).clamp(min=0)
         target = target.to(device=device, dtype=dtype).clamp(min=0)
@@ -107,9 +109,7 @@ def solve(
         if return_potentials:
             output["log_u"], output["log_v"] = log_u, log_v
         if return_plan or return_distance:
-            plan = torch.exp(
-                log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2)
-            )
+            plan = torch.exp(log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2))
             if return_plan:
                 output["plan"] = plan
             if return_distance:

From ec0119eb00c76187fb01662048515b38aa3fa069 Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Wed, 8 Jul 2026 22:38:14 +0800
Subject: [PATCH 17/20] Make SinkhornSolver a differentiable nn.Module with
 automatic device dispatch

- SinkhornSolver now inherits torch.nn.Module: forward(source, target,
  cost_matrix) returns the (n,) transport distances, and a custom
  autograd.Function implements backward via the centered dual potentials
  (envelope theorem), chained through the differentiable normalization.
- Replace the confused torch/cuda/cuda_log backend strings with one
  orthogonal option: log_space=True/False selects the algorithm, while the
  implementation is picked from the input automatically (fused CUDA kernels
  for float32 CUDA tensors, pure torch ops elsewhere - including a
  logsumexp-based log-space path that also runs on CPU and float64).
- Convenience methods plan() and potentials() expose the transport plan and
  dual potentials; extra_repr() reports the module configuration.
- Tests cover both log_space modes across forward shapes, marginals, POT
  baselines, fused-vs-CPU consistency, small-epsilon stability, and
  autograd-vs-numeric gradients; benchmark updated to the module API.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 benchmark/optimal_transport.py  |  66 +++-----
 test/test_optimal_transport.py  | 187 ++++++++++------------
 torchmorph/optimal_transport.py | 271 +++++++++++++++++---------------
 3 files changed, 256 insertions(+), 268 deletions(-)

diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
index b9afa33..b800933 100644
--- a/benchmark/optimal_transport.py
+++ b/benchmark/optimal_transport.py
@@ -5,23 +5,23 @@
 from torchmorph import SinkhornSolver, build_cost_matrix
 
 
-def _grid_problem(n, H, W, epsilon, max_iter, seed=42):
-    """Random (n, H*W) marginals with a 2-D grid cost matrix, preprocessed."""
+def _grid_problem(n, H, W, epsilon, max_iter, log_space=False, seed=42):
+    """Random (n, H*W) marginals with a 2-D grid cost matrix."""
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.manual_seed(seed)
     source = torch.rand(n, H * W, device=device)
     target = torch.rand(n, H * W, device=device)
-    solver = SinkhornSolver(epsilon=epsilon, max_iter=max_iter, device=device)
-    source, target, cost_matrix = solver.data_preprocess(
-        source, target, cost_matrix=build_cost_matrix((H, W), device=device)
-    )
+    cost_matrix = build_cost_matrix((H, W), device=device)
+    solver = SinkhornSolver(epsilon=epsilon, max_iter=max_iter, log_space=log_space)
     return solver, source, target, cost_matrix
 
 
-def _time_backend(method_name, solver, source, target, cost_matrix, min_run_time=3):
-    torch.cuda.reset_peak_memory_stats()
+def run_forward_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100, log_space=False):
+    solver, source, target, cost_matrix = _grid_problem(n, H, W, epsilon, max_iter, log_space)
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
     timer = benchmark.Timer(
-        stmt=f"solver.{method_name}(source, target, cost_matrix)",
+        stmt="solver(source, target, cost_matrix)",
         globals={
             "solver": solver,
             "source": source,
@@ -29,29 +29,16 @@ def _time_backend(method_name, solver, source, target, cost_matrix, min_run_time
             "cost_matrix": cost_matrix,
         },
     )
-    result = timer.blocked_autorange(min_run_time=min_run_time)
-    print(f"peak allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
-    print(f"peak reserved:  {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
+    result = timer.blocked_autorange(min_run_time=3)
+    if torch.cuda.is_available():
+        print(f"peak allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+        print(f"peak reserved:  {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
     return result
 
 
-def run_sinkhorn_torch_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100):
-    args = _grid_problem(n, H, W, epsilon, max_iter)
-    return _time_backend("sinkhorn_torch", *args)
-
-
-def run_sinkhorn_cuda_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100):
-    args = _grid_problem(n, H, W, epsilon, max_iter)
-    return _time_backend("sinkhorn_cuda", *args)
-
-
-def run_sinkhorn_log_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100):
-    args = _grid_problem(n, H, W, epsilon, max_iter)
-    return _time_backend("sinkhorn_log_cuda", *args)
-
-
 def run_pot_sinkhorn_benchmark(H=32, W=32, epsilon=1.0, max_iter=100):
     solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter)
+    source, target, cost_matrix = solver.data_preprocess(source, target, cost_matrix)
     timer = benchmark.Timer(
         stmt="ot.sinkhorn(source, target, M, reg=epsilon, numItermax=max_iter, stopThr=1e-5)",
         globals={
@@ -69,28 +56,27 @@ def run_pot_sinkhorn_benchmark(H=32, W=32, epsilon=1.0, max_iter=100):
 @torch.no_grad()
 def run_sinkhorn_relative_error(H=32, W=32, epsilon=1.0, max_iter=100):
     solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter)
-    out = solver.solve(source, target, cost_matrix=cost_matrix, return_plan=True)
-    ot_plan = ot.sinkhorn(
-        source[0], target[0], cost_matrix, reg=epsilon, numItermax=max_iter, stopThr=1e-5
-    )
-    return torch.linalg.norm(out["plan"][0] - ot_plan) / torch.clamp(
-        torch.linalg.norm(ot_plan), min=1e-12
-    )
+    plan = solver.plan(source, target, cost_matrix)
+    a, b, cost_matrix = solver.data_preprocess(source, target, cost_matrix)
+    ot_plan = ot.sinkhorn(a[0], b[0], cost_matrix, reg=epsilon, numItermax=max_iter, stopThr=1e-5)
+    return torch.linalg.norm(plan[0] - ot_plan) / torch.clamp(torch.linalg.norm(ot_plan), min=1e-12)
 
 
 def reference_error_check(size=32, epsilon=10.0, max_iter=200, verbose=False):
-    """Compare log-domain CUDA gradients against POT's sinkhorn_log potentials."""
+    """Compare log-space fused-kernel potentials against POT's sinkhorn_log."""
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA device not available.")
-    solver, source, target, cost_matrix = _grid_problem(1, size, size, epsilon, max_iter, seed=2)
+    solver, source, target, cost_matrix = _grid_problem(
+        1, size, size, epsilon, max_iter, log_space=True, seed=2
+    )
 
-    log_u, log_v = solver.sinkhorn_log_cuda(source, target, cost_matrix)
-    grad_f, grad_g = solver.gradient(log_u, log_v)
+    grad_f, grad_g = solver.potentials(source, target, cost_matrix)
     grad_f, grad_g = grad_f[0], grad_g[0]
 
+    a, b, cost_matrix = solver.data_preprocess(source, target, cost_matrix)
     _, log = ot.bregman.sinkhorn_log(
-        source[0],
-        target[0],
+        a[0],
+        b[0],
         cost_matrix,
         reg=epsilon,
         numItermax=max_iter,
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 5363d58..2785fc6 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -16,7 +16,7 @@ def _device():
 
 def _require_cuda():
     if not torch.cuda.is_available():
-        pytest.skip("CUDA is required for the CUDA Sinkhorn backends")
+        pytest.skip("CUDA is required for the fused Sinkhorn kernels")
 
 
 def _require_pot():
@@ -59,186 +59,163 @@ def test_build_cost_matrix():
     assert build_cost_matrix((2, 2), p=2)[0, 3] == pytest.approx(2.0**0.5)
 
 
+@pytest.mark.parametrize("log_space", [False, True])
 @pytest.mark.parametrize("n,d", [(1, 4), (3, 6), (2, 16)])
-def test_solve_shapes(n, d):
+def test_forward_shapes_and_marginals(n, d, log_space):
     device = _device()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=300, device=device)
+    solver = SinkhornSolver(epsilon=1.5, max_iter=300, log_space=log_space)
 
     source = _make_positive((n, d), device=device, seed=11)
     target = _make_positive((n, d), device=device, seed=23)
-    out = solver.solve(
-        source,
-        target,
-        return_plan=True,
-        return_potentials=True,
-        return_distance=True,
-    )
 
-    assert out["source"].shape == (n, d)
-    assert out["target"].shape == (n, d)
-    assert out["cost_matrix"].shape == (d, d)
-    assert out["plan"].shape == (n, d, d)
-    assert out["log_u"].shape == (n, d)
-    assert out["log_v"].shape == (n, d)
-    assert out["distance"].shape == (n,)
+    distance = solver(source, target)
+    plan = solver.plan(source, target)
+    f, g = solver.potentials(source, target)
+    assert distance.shape == (n,)
+    assert plan.shape == (n, d, d)
+    assert f.shape == (n, d) and g.shape == (n, d)
 
-    distance_from_plan = (out["plan"] * out["cost_matrix"]).sum(dim=(-2, -1))
-    assert torch.allclose(out["distance"], distance_from_plan, rtol=1e-5, atol=1e-6)
-    # rows sum to the normalized marginals
-    assert torch.allclose(out["source"].sum(dim=-1), torch.ones(n, device=device))
+    a, b, cost_matrix = solver.data_preprocess(source, target)
+    assert torch.allclose(distance, (plan * cost_matrix).sum(dim=(-2, -1)), rtol=1e-5, atol=1e-6)
+    assert torch.allclose(plan.sum(dim=-1), a, atol=1e-4)
+    assert torch.allclose(plan.sum(dim=-2), b, atol=1e-4)
 
 
 @pytest.mark.parametrize("shape", [(4,), (2, 2, 2)])
-def test_solve_rejects_non_2d_inputs(shape):
+def test_forward_rejects_non_2d_inputs(shape):
     solver = SinkhornSolver(epsilon=1.5, max_iter=10)
     source = _make_positive(shape, device="cpu", seed=911)
     target = _make_positive(shape, device="cpu", seed=919)
 
     with pytest.raises(ValueError, match=r"\(n, d\)"):
-        solver.solve(source, target)
+        solver(source, target)
 
 
-def test_solve_rejects_mismatched_cost_matrix():
+def test_forward_rejects_mismatched_cost_matrix():
     solver = SinkhornSolver(epsilon=1.5, max_iter=10)
     source = _make_positive((2, 4), device="cpu", seed=5)
     target = _make_positive((2, 4), device="cpu", seed=7)
 
     with pytest.raises(ValueError, match="cost_matrix"):
-        solver.solve(source, target, cost_matrix=torch.zeros(3, 3))
+        solver(source, target, cost_matrix=torch.zeros(3, 3))
 
 
-def test_solve_distance_matches_pot_sinkhorn():
+def test_distance_matches_pot_sinkhorn():
     _require_pot()
     device = _device()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, device=device)
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000)
 
     source = _make_positive((3, 8), device=device, seed=101)
     target = _make_positive((3, 8), device=device, seed=211)
-    out = solver.solve(source, target, return_distance=True)
+    distance = solver(source, target)
 
-    cost_np = out["cost_matrix"].detach().cpu().numpy()
+    a, b, cost_matrix = solver.data_preprocess(source, target)
+    cost_np = cost_matrix.cpu().numpy()
     expected = torch.tensor(
         [
-            float(
-                np.sum(
-                    _pot_plan(out["source"][i], out["target"][i], out["cost_matrix"], solver)
-                    * cost_np
-                )
-            )
-            for i in range(out["source"].shape[0])
+            float(np.sum(_pot_plan(a[i], b[i], cost_matrix, solver) * cost_np))
+            for i in range(a.shape[0])
         ],
         device=device,
-        dtype=out["distance"].dtype,
+        dtype=distance.dtype,
     )
-    assert torch.allclose(out["distance"], expected, rtol=5e-3, atol=1e-5)
+    assert torch.allclose(distance, expected, rtol=5e-3, atol=1e-5)
 
 
-def test_solve_plan_matches_pot_sinkhorn_grid_cost():
+def test_plan_matches_pot_sinkhorn_grid_cost():
     """A flattened 2-D image with an explicit grid cost matrix matches POT."""
     _require_pot()
     device = _device()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, device=device)
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000)
     cost_matrix = build_cost_matrix((2, 3), device=device)
 
     source = _make_positive((1, 6), device=device, seed=307)
     target = _make_positive((1, 6), device=device, seed=401)
-    out = solver.solve(source, target, cost_matrix=cost_matrix, return_plan=True)
+    plan = solver.plan(source, target, cost_matrix)
 
-    pot_transport = _pot_plan(out["source"][0], out["target"][0], out["cost_matrix"], solver)
-    expected = torch.tensor(pot_transport, device=device, dtype=out["plan"].dtype)
-    assert torch.allclose(out["plan"][0], expected, rtol=5e-3, atol=1e-5)
+    a, b, _ = solver.data_preprocess(source, target, cost_matrix)
+    expected = torch.tensor(
+        _pot_plan(a[0], b[0], cost_matrix, solver), device=device, dtype=plan.dtype
+    )
+    assert torch.allclose(plan[0], expected, rtol=5e-3, atol=1e-5)
 
 
-@pytest.mark.parametrize("backend", ["cuda", "cuda_log"])
-def test_cuda_backends_match_torch_backend_batched(backend):
+@pytest.mark.parametrize("log_space", [False, True])
+def test_fused_cuda_kernels_match_cpu(log_space):
+    """CUDA float32 inputs dispatch to the fused kernels; results must match CPU."""
     _require_cuda()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, device="cuda")
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
 
-    source = _make_positive((3, 6), device="cuda", seed=503)
-    target = _make_positive((3, 6), device="cuda", seed=607)
+    source = _make_positive((3, 6), device="cpu", seed=503)
+    target = _make_positive((3, 6), device="cpu", seed=607)
 
-    torch_out = solver.solve(source, target, return_plan=True, return_distance=True)
-    cuda_out = solver.solve(source, target, backend=backend, return_plan=True, return_distance=True)
+    cpu_distance = solver(source, target)
+    cpu_plan = solver.plan(source, target)
+    cuda_distance = solver(source.cuda(), target.cuda())
+    cuda_plan = solver.plan(source.cuda(), target.cuda())
 
-    assert cuda_out["backend"] == backend
-    assert torch.allclose(cuda_out["plan"], torch_out["plan"], rtol=1e-4, atol=1e-5)
-    assert torch.allclose(cuda_out["distance"], torch_out["distance"], rtol=1e-4, atol=1e-5)
+    assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5)
+    assert torch.allclose(cuda_plan.cpu(), cpu_plan, rtol=1e-4, atol=1e-5)
 
 
-def test_log_domain_survives_small_epsilon():
-    """At this epsilon exp(-C/eps) underflows float32, so only the log backend is usable."""
+def test_log_space_survives_small_epsilon():
+    """At this epsilon exp(-C/eps) underflows float32, so only log_space=True is usable."""
     _require_cuda()
     _require_pot()
-    solver = SinkhornSolver(epsilon=0.1, max_iter=5000, device="cuda")
+    solver = SinkhornSolver(epsilon=0.1, max_iter=5000, log_space=True)
     cost_matrix = build_cost_matrix((8, 8), device="cuda")
 
     source = _make_positive((1, 64), device="cuda", seed=31)
     target = _make_positive((1, 64), device="cuda", seed=37)
-    out = solver.solve(
-        source,
-        target,
-        cost_matrix=cost_matrix,
-        backend="cuda_log",
-        return_plan=True,
-        return_potentials=True,
-        return_distance=True,
-    )
 
-    assert torch.isfinite(out["log_u"]).all()
-    assert torch.isfinite(out["log_v"]).all()
-    plan = out["plan"][0]
+    f, g = solver.potentials(source, target, cost_matrix)
+    assert torch.isfinite(f).all()
+    assert torch.isfinite(g).all()
+
+    plan = solver.plan(source, target, cost_matrix)[0]
+    a, b, _ = solver.data_preprocess(source, target, cost_matrix)
     assert torch.isfinite(plan).all()
     assert (plan >= 0).all()
-    assert torch.allclose(plan.sum(dim=0), out["target"][0], atol=1e-5)
-    assert torch.allclose(plan.sum(dim=1), out["source"][0], atol=5e-3)
+    assert torch.allclose(plan.sum(dim=0), b[0], atol=1e-5)
+    assert torch.allclose(plan.sum(dim=1), a[0], atol=5e-3)
 
-    cost_np = out["cost_matrix"].cpu().numpy().astype(np.float64)
+    cost_np = cost_matrix.cpu().numpy().astype(np.float64)
     pot_plan = ot.bregman.sinkhorn_log(
-        out["source"][0].cpu().numpy().astype(np.float64),
-        out["target"][0].cpu().numpy().astype(np.float64),
+        a[0].cpu().numpy().astype(np.float64),
+        b[0].cpu().numpy().astype(np.float64),
         cost_np,
         reg=solver.epsilon,
         numItermax=solver.max_iter,
         stopThr=1e-12,
     )
     pot_distance = float(np.sum(pot_plan * cost_np))
-    assert out["distance"][0].item() == pytest.approx(pot_distance, rel=2e-2)
+    distance = solver(source, target, cost_matrix)
+    assert distance[0].item() == pytest.approx(pot_distance, rel=2e-2)
 
 
+@pytest.mark.parametrize("log_space", [False, True])
 @pytest.mark.parametrize("n,d", [(1, 4), (3, 8)])
-def test_gradient_matches_numeric_directional_derivative(n, d):
+def test_backward_matches_numeric_directional_derivative(n, d, log_space):
+    """Autograd through forward() must match a numeric derivative of the entropic cost."""
     device = _device()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=3000, device=device)
-
-    source_raw = _make_positive((n, d), device=device, seed=1009)
-    target_raw = _make_positive((n, d), device=device, seed=2003)
-    source, target, cost_matrix = solver.data_preprocess(
-        source_raw, target_raw, dtype=torch.float64
-    )
+    solver = SinkhornSolver(epsilon=1.5, max_iter=3000, log_space=log_space)
+    cost_matrix = build_cost_matrix((d,), device=device).double()
 
+    source = _make_positive((n, d), device=device, seed=1009).double().requires_grad_(True)
+    target = _make_positive((n, d), device=device, seed=2003).double()
     direction = _make_positive((n, d), device=device, seed=3001).double() - 0.6
-    direction = direction - direction.mean(dim=-1, keepdim=True)
-    eps_fd = 1e-4
 
-    def regularized_cost(x):
-        out = solver.solve(
-            x, target, cost_matrix=cost_matrix, return_distance=True, dtype=torch.float64
-        )
-        entropy = (out["plan"] * torch.log(out["plan"].clamp(min=1e-300))).sum()
-        return out["distance"].sum() + solver.epsilon * entropy
-
-    out = solver.solve(
-        source,
-        target,
-        cost_matrix=cost_matrix,
-        return_plan=False,
-        return_grad=True,
-        dtype=torch.float64,
-    )
-    analytic_directional = (out["grad_source"] * direction).sum()
+    solver(source, target, cost_matrix).sum().backward()
+    analytic = (source.grad * direction).sum()
 
-    loss_plus = regularized_cost(source + eps_fd * direction)
-    loss_minus = regularized_cost(source - eps_fd * direction)
-    numeric_directional = (loss_plus - loss_minus) / (2 * eps_fd)
+    def entropic_cost(x):
+        plan = solver.plan(x, target, cost_matrix)
+        entropy = (plan * plan.clamp(min=1e-300).log()).sum()
+        return (plan * cost_matrix).sum() + solver.epsilon * entropy
 
-    assert torch.allclose(analytic_directional, numeric_directional, rtol=1e-5, atol=1e-7)
+    eps_fd = 1e-4
+    raw = source.detach()
+    numeric = (
+        entropic_cost(raw + eps_fd * direction) - entropic_cost(raw - eps_fd * direction)
+    ) / (2 * eps_fd)
+    assert torch.allclose(analytic, numeric, rtol=1e-5, atol=1e-7)
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index e249a8f..3415283 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -1,7 +1,7 @@
 from typing import Optional
 
 import torch
-from torch import Tensor
+from torch import Tensor, nn
 
 
 def build_cost_matrix(shape, p=2, device=None) -> Tensor:
@@ -18,18 +18,112 @@ def build_cost_matrix(shape, p=2, device=None) -> Tensor:
     return torch.cdist(coords, coords, p=p)
 
 
-class SinkhornSolver:
-    """Entropy-regularized balanced optimal transport between batched histograms.
+def _use_fused_kernels(t: Tensor) -> bool:
+    return t.is_cuda and t.dtype == torch.float32
 
-    Every method takes source and target as (n, d) tensors: n distributions of
-    d bins each, sharing one (d, d) cost matrix. The standard convention
-    K = exp(-C / epsilon) is used: smaller epsilon means weaker blurring and a
-    sharper transport plan. All backends return the log scaling vectors
-    (log_u, log_v); the dual potentials are f = epsilon * log_u and
-    g = epsilon * log_v.
+
+def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold):
+    """Plain scaling-form Sinkhorn on (n, d) marginals; returns (log_u, log_v)."""
+    K = torch.exp(-cost_matrix / epsilon)
+    if _use_fused_kernels(a):
+        from torchmorph import _C
+
+        u, v = _C.sinkhorn_fastiter(
+            a.contiguous(),
+            b.contiguous(),
+            K,
+            K.mT.contiguous(),
+            torch.ones_like(a),
+            torch.ones_like(b),
+            max_iter,
+        )
+        return u.log(), v.log()
+
+    u = torch.ones_like(a)
+    v = torch.ones_like(b)
+    for i in range(max_iter):
+        u_prev = u
+        u = a / (v @ K.mT + 1e-12)
+        v = b / (u @ K + 1e-12)
+        if threshold > 0 and i % 10 == 0:
+            if (u - u_prev).abs().sum(dim=-1).max() <= threshold:
+                break
+    return u.log(), v.log()
+
+
+def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold):
+    """Log-domain Sinkhorn on (n, d) marginals; returns (log_u, log_v)."""
+    log_a, log_b = a.log(), b.log()
+    if _use_fused_kernels(a):
+        from torchmorph import _C
+
+        return _C.sinkhorn_logiter(
+            log_a.contiguous(),
+            log_b.contiguous(),
+            cost_matrix,
+            cost_matrix.mT.contiguous(),
+            torch.zeros_like(a),
+            torch.zeros_like(b),
+            max_iter,
+            epsilon,
+        )
+
+    log_K = -cost_matrix / epsilon
+    log_u = torch.zeros_like(a)
+    log_v = torch.zeros_like(b)
+    for i in range(max_iter):
+        log_u_prev = log_u
+        log_u = log_a - torch.logsumexp(log_K + log_v.unsqueeze(-2), dim=-1)
+        log_v = log_b - torch.logsumexp(log_K.mT + log_u.unsqueeze(-2), dim=-1)
+        if threshold > 0 and i % 10 == 0:
+            if (log_u - log_u_prev).abs().sum(dim=-1).max() <= threshold:
+                break
+    return log_u, log_v
+
+
+class _SinkhornDistance(torch.autograd.Function):
+    """Transport distance <P, C> with envelope-theorem gradients.
+
+    The backward pass returns the centered dual potentials, i.e. the exact
+    gradients of the entropic OT cost w.r.t. the (normalized) marginals.
     """
 
-    def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, device=None):
+    @staticmethod
+    def forward(ctx, a, b, cost_matrix, epsilon, max_iter, threshold, log_space):
+        iterate = _log_iterations if log_space else _scaling_iterations
+        log_u, log_v = iterate(a, b, cost_matrix, epsilon, max_iter, threshold)
+        plan = torch.exp(log_u.unsqueeze(-1) - cost_matrix / epsilon + log_v.unsqueeze(-2))
+        distance = (plan * cost_matrix).sum(dim=(-2, -1))
+        f = epsilon * log_u
+        g = epsilon * log_v
+        ctx.save_for_backward(f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True))
+        return distance
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        f, g = ctx.saved_tensors
+        grad = grad_output.unsqueeze(-1)
+        return grad * f, grad * g, None, None, None, None, None
+
+
+class SinkhornSolver(nn.Module):
+    """Entropy-regularized balanced optimal transport as a differentiable module.
+
+    forward(source, target, cost_matrix=None) takes two (n, d) batches of
+    histograms sharing one (d, d) cost matrix and returns the (n,) transport
+    distances <P, C>, with K = exp(-C / epsilon). Gradients w.r.t. source and
+    target are the centered dual potentials (envelope theorem), so the module
+    can be used directly as a loss.
+
+    log_space=True runs the iterations entirely in the log domain, which is
+    numerically stable for small epsilon. The implementation is picked from
+    the input automatically: CUDA float32 tensors use the fused kernels,
+    everything else uses pure torch ops (where `threshold` > 0 enables early
+    stopping; the fused kernels always run `max_iter` iterations).
+    """
+
+    def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False):
+        super().__init__()
         if epsilon <= 0:
             raise ValueError("epsilon must be positive.")
         if max_iter <= 0:
@@ -40,141 +134,72 @@ def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, device=None):
         self.max_iter = max_iter
         self.threshold = threshold
         self.p = p
-        self.device = device
-
-    def data_preprocess(
-        self,
-        source: Tensor,
-        target: Tensor,
-        cost_matrix: Optional[Tensor] = None,
-        dtype=torch.float32,
-    ):
+        self.log_space = log_space
+
+    def extra_repr(self):
+        return (
+            f"epsilon={self.epsilon}, max_iter={self.max_iter}, "
+            f"threshold={self.threshold}, log_space={self.log_space}"
+        )
+
+    def data_preprocess(self, source, target, cost_matrix: Optional[Tensor] = None):
         """Validate (n, d) inputs, clamp negatives, and normalize each row to unit mass.
 
-        When `cost_matrix` is None the d bins are treated as points on a line.
+        All ops are differentiable, so forward gradients flow through the
+        normalization. When `cost_matrix` is None the d bins are treated as
+        points on a line.
         """
         if source.ndim != 2 or source.shape != target.shape:
             raise ValueError("source and target must be (n, d) tensors of the same shape.")
-        device = torch.device(self.device) if self.device is not None else source.device
+        if not source.is_floating_point():
+            raise ValueError("source and target must be floating-point tensors.")
         d = source.shape[1]
         if cost_matrix is None:
-            cost_matrix = build_cost_matrix((d,), self.p, device)
+            cost_matrix = build_cost_matrix((d,), self.p, source.device)
         elif cost_matrix.shape != (d, d):
             raise ValueError(
                 f"cost_matrix must have shape ({d}, {d}), got {tuple(cost_matrix.shape)}."
             )
+        cost_matrix = cost_matrix.to(device=source.device, dtype=source.dtype).contiguous()
 
-        source = source.to(device=device, dtype=dtype).clamp(min=0)
-        target = target.to(device=device, dtype=dtype).clamp(min=0)
-        cost_matrix = cost_matrix.to(device=device, dtype=dtype).contiguous()
+        source = source.clamp(min=0)
+        target = target.clamp(min=0)
         source = source / source.sum(dim=-1, keepdim=True).clamp(min=1e-12)
         target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-12)
         return source, target, cost_matrix
 
-    def solve(
-        self,
-        source: Tensor,
-        target: Tensor,
-        cost_matrix: Optional[Tensor] = None,
-        *,
-        backend: str = "torch",
-        return_plan: bool = True,
-        return_potentials: bool = False,
-        return_grad: bool = False,
-        return_distance: bool = False,
-        dtype=torch.float32,
-    ):
-        """Run Sinkhorn end-to-end on (n, d) source and target distributions.
-
-        backend is "torch" (batched matmuls, any device), "cuda" (scaling-form
-        kernel), or "cuda_log" (log-domain kernel, stable for small epsilon).
-        The (n, d, d) plan is always reconstructed in log space.
-        """
-        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix, dtype)
-        backends = {
-            "torch": self.sinkhorn_torch,
-            "cuda": self.sinkhorn_cuda,
-            "cuda_log": self.sinkhorn_log_cuda,
-        }
-        if backend not in backends:
-            raise ValueError(f"Unknown backend {backend!r}; expected one of {sorted(backends)}.")
-        log_u, log_v = backends[backend](source, target, cost_matrix)
-
-        output = {
-            "source": source,
-            "target": target,
-            "cost_matrix": cost_matrix,
-            "backend": backend,
-        }
-        if return_potentials:
-            output["log_u"], output["log_v"] = log_u, log_v
-        if return_plan or return_distance:
-            plan = torch.exp(log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2))
-            if return_plan:
-                output["plan"] = plan
-            if return_distance:
-                output["distance"] = (plan * cost_matrix).sum(dim=(-2, -1))
-        if return_grad:
-            output["grad_source"], output["grad_target"] = self.gradient(log_u, log_v)
-        return output
-
-    def sinkhorn_torch(self, source: Tensor, target: Tensor, cost_matrix: Tensor):
-        """Batched Sinkhorn in scaling form with dense matmuls.
-
-        Expects preprocessed (n, d) inputs; returns (log_u, log_v).
-        """
-        K = torch.exp(-cost_matrix / self.epsilon)
-        u = torch.ones_like(source)
-        v = torch.ones_like(target)
-        for i in range(self.max_iter):
-            u_prev = u
-            u = source / (v @ K.mT + 1e-12)
-            v = target / (u @ K + 1e-12)
-            if self.threshold > 0 and i % 10 == 0:
-                if (u - u_prev).abs().sum(dim=-1).max() <= self.threshold:
-                    break
-        return u.log(), v.log()
-
-    @torch.no_grad()
-    def sinkhorn_cuda(self, source: Tensor, target: Tensor, cost_matrix: Tensor):
-        """Scaling-form Sinkhorn on the fused CUDA kernel; returns (log_u, log_v)."""
-        from torchmorph import _C
-
-        K = torch.exp(-cost_matrix / self.epsilon)
-        u = torch.ones_like(source)
-        v = torch.ones_like(target)
-        u, v = _C.sinkhorn_fastiter(
-            source.contiguous(), target.contiguous(), K, K.mT.contiguous(), u, v, self.max_iter
-        )
-        return u.log(), v.log()
-
-    @torch.no_grad()
-    def sinkhorn_log_cuda(self, source: Tensor, target: Tensor, cost_matrix: Tensor):
-        """Log-domain Sinkhorn on the fused CUDA kernel; stable for small epsilon.
-
-        Returns (log_u, log_v) without ever leaving log space.
-        """
-        from torchmorph import _C
-
-        log_u = torch.zeros_like(source)
-        log_v = torch.zeros_like(target)
-        return _C.sinkhorn_logiter(
-            source.log().contiguous(),
-            target.log().contiguous(),
+    def forward(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None):
+        """(n,) transport distances between the rows of source and target."""
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix)
+        return _SinkhornDistance.apply(
+            source,
+            target,
             cost_matrix,
-            cost_matrix.mT.contiguous(),
-            log_u,
-            log_v,
-            self.max_iter,
             self.epsilon,
+            self.max_iter,
+            self.threshold,
+            self.log_space,
         )
 
-    def gradient(self, log_u: Tensor, log_v: Tensor):
-        """Gradients of the entropic OT cost w.r.t. the marginals.
+    @torch.no_grad()
+    def plan(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None):
+        """(n, d, d) transport plans, reconstructed in log space."""
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix)
+        log_u, log_v = self._iterate(source, target, cost_matrix)
+        return torch.exp(log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2))
 
-        These are the dual potentials f = epsilon * log_u and g = epsilon * log_v,
-        centered per distribution to fix the additive gauge freedom.
+    @torch.no_grad()
+    def potentials(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None):
+        """Centered dual potentials (f, g), each (n, d).
+
+        These are the gradients of the entropic OT cost w.r.t. the marginals.
         """
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix)
+        log_u, log_v = self._iterate(source, target, cost_matrix)
         f = self.epsilon * log_u
         g = self.epsilon * log_v
         return f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True)
+
+    def _iterate(self, a, b, cost_matrix):
+        iterate = _log_iterations if self.log_space else _scaling_iterations
+        return iterate(a, b, cost_matrix, self.epsilon, self.max_iter, self.threshold)

From 2d4da6c909ec3cdd1c18e0a6fb317ed1b6d7cefb Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Wed, 8 Jul 2026 22:44:11 +0800
Subject: [PATCH 18/20] Expand Sinkhorn tests to a full device x log_space x
 dtype matrix

Previously most tests picked a single device (CUDA when available), so the
CPU paths went untested on GPU machines, POT baselines only covered the
plain iteration, and backward never ran through the fused kernels. Now:

- Shared DEVICES parametrization runs every applicable test on both CPU and
  CUDA (skipping cleanly when CUDA is absent).
- POT distance/plan baselines and the small-epsilon stability test cover
  both log_space modes and both devices (torch logsumexp path on CPU, fused
  log kernel on CUDA).
- New tests: dtype/device consistency (float64 CUDA exercises the torch
  fallback on GPU), backward through the fused float32 kernels vs CPU,
  fused-vs-CPU potentials, and threshold early stopping reaching the same
  solution on both algorithms.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 test/test_optimal_transport.py | 149 +++++++++++++++++++++++++--------
 1 file changed, 115 insertions(+), 34 deletions(-)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 2785fc6..6ceb62f 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -9,9 +9,16 @@
 except ImportError:  # pragma: no cover - exercised only when POT is absent
     ot = None
 
-
-def _device():
-    return "cuda" if torch.cuda.is_available() else "cpu"
+# Every test parametrized over DEVICES runs on CPU and, when available, CUDA,
+# so both the pure-torch and the fused-kernel implementations are covered.
+DEVICES = [
+    "cpu",
+    pytest.param(
+        "cuda",
+        marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"),
+    ),
+]
+LOG_SPACE = [False, True]
 
 
 def _require_cuda():
@@ -40,29 +47,31 @@ def _pot_plan(source, target, cost_matrix, solver):
     )
 
 
-def test_build_cost_matrix():
+@pytest.mark.parametrize("device", DEVICES)
+def test_build_cost_matrix(device):
     # 1-D grid: zero diagonal, symmetry, unit spacing
-    cost_1d = build_cost_matrix((5,))
+    cost_1d = build_cost_matrix((5,), device=device)
     assert cost_1d.shape == (5, 5)
-    assert torch.allclose(cost_1d.diag(), torch.zeros(5))
+    assert cost_1d.device.type == device
+    assert torch.allclose(cost_1d.diag(), torch.zeros(5, device=device))
     assert torch.allclose(cost_1d, cost_1d.T)
-    assert cost_1d[0, 1] == pytest.approx(1.0)
+    assert cost_1d[0, 1].item() == pytest.approx(1.0)
 
     # 2-D grid: adjacent columns are distance 1 apart
-    cost_2d = build_cost_matrix((3, 4))
+    cost_2d = build_cost_matrix((3, 4), device=device)
     assert cost_2d.shape == (12, 12)
     assert torch.allclose(cost_2d, cost_2d.T)
-    assert cost_2d[0, 1] == pytest.approx(1.0)
+    assert cost_2d[0, 1].item() == pytest.approx(1.0)
 
     # p-norms differ on diagonal moves: grid points (0, 0) and (1, 1)
-    assert build_cost_matrix((2, 2), p=1)[0, 3] == pytest.approx(2.0)
-    assert build_cost_matrix((2, 2), p=2)[0, 3] == pytest.approx(2.0**0.5)
+    assert build_cost_matrix((2, 2), p=1, device=device)[0, 3].item() == pytest.approx(2.0)
+    assert build_cost_matrix((2, 2), p=2, device=device)[0, 3].item() == pytest.approx(2.0**0.5)
 
 
-@pytest.mark.parametrize("log_space", [False, True])
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
 @pytest.mark.parametrize("n,d", [(1, 4), (3, 6), (2, 16)])
-def test_forward_shapes_and_marginals(n, d, log_space):
-    device = _device()
+def test_forward_shapes_and_marginals(n, d, log_space, device):
     solver = SinkhornSolver(epsilon=1.5, max_iter=300, log_space=log_space)
 
     source = _make_positive((n, d), device=device, seed=11)
@@ -72,6 +81,7 @@ def test_forward_shapes_and_marginals(n, d, log_space):
     plan = solver.plan(source, target)
     f, g = solver.potentials(source, target)
     assert distance.shape == (n,)
+    assert distance.device.type == device
     assert plan.shape == (n, d, d)
     assert f.shape == (n, d) and g.shape == (n, d)
 
@@ -100,10 +110,11 @@ def test_forward_rejects_mismatched_cost_matrix():
         solver(source, target, cost_matrix=torch.zeros(3, 3))
 
 
-def test_distance_matches_pot_sinkhorn():
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_distance_matches_pot_sinkhorn(log_space, device):
     _require_pot()
-    device = _device()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=1000)
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
 
     source = _make_positive((3, 8), device=device, seed=101)
     target = _make_positive((3, 8), device=device, seed=211)
@@ -122,11 +133,12 @@ def test_distance_matches_pot_sinkhorn():
     assert torch.allclose(distance, expected, rtol=5e-3, atol=1e-5)
 
 
-def test_plan_matches_pot_sinkhorn_grid_cost():
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_plan_matches_pot_sinkhorn_grid_cost(log_space, device):
     """A flattened 2-D image with an explicit grid cost matrix matches POT."""
     _require_pot()
-    device = _device()
-    solver = SinkhornSolver(epsilon=1.5, max_iter=1000)
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
     cost_matrix = build_cost_matrix((2, 3), device=device)
 
     source = _make_positive((1, 6), device=device, seed=307)
@@ -140,7 +152,7 @@ def test_plan_matches_pot_sinkhorn_grid_cost():
     assert torch.allclose(plan[0], expected, rtol=5e-3, atol=1e-5)
 
 
-@pytest.mark.parametrize("log_space", [False, True])
+@pytest.mark.parametrize("log_space", LOG_SPACE)
 def test_fused_cuda_kernels_match_cpu(log_space):
     """CUDA float32 inputs dispatch to the fused kernels; results must match CPU."""
     _require_cuda()
@@ -151,22 +163,53 @@ def test_fused_cuda_kernels_match_cpu(log_space):
 
     cpu_distance = solver(source, target)
     cpu_plan = solver.plan(source, target)
+    cpu_f, cpu_g = solver.potentials(source, target)
     cuda_distance = solver(source.cuda(), target.cuda())
     cuda_plan = solver.plan(source.cuda(), target.cuda())
+    cuda_f, cuda_g = solver.potentials(source.cuda(), target.cuda())
 
     assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5)
     assert torch.allclose(cuda_plan.cpu(), cpu_plan, rtol=1e-4, atol=1e-5)
-
-
-def test_log_space_survives_small_epsilon():
-    """At this epsilon exp(-C/eps) underflows float32, so only log_space=True is usable."""
-    _require_cuda()
+    assert torch.allclose(cuda_f.cpu(), cpu_f, rtol=1e-4, atol=1e-4)
+    assert torch.allclose(cuda_g.cpu(), cpu_g, rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_dtype_device_combinations_agree(log_space):
+    """float32/float64 on CPU/CUDA must all reach the same solution.
+
+    float64 CUDA inputs exercise the torch fallback on the GPU, float32 CUDA
+    inputs the fused kernels, and CPU inputs the pure-torch paths.
+    """
+    solver = SinkhornSolver(epsilon=1.5, max_iter=2000, log_space=log_space)
+    source = _make_positive((2, 8), device="cpu", seed=41)
+    target = _make_positive((2, 8), device="cpu", seed=43)
+    reference = solver(source.double(), target.double())
+
+    combos = [(source, target)]  # cpu float32
+    if torch.cuda.is_available():
+        combos += [
+            (source.cuda(), target.cuda()),  # cuda float32: fused kernels
+            (source.cuda().double(), target.cuda().double()),  # cuda float64: torch ops
+        ]
+    for src, tgt in combos:
+        distance = solver(src, tgt)
+        tol = 1e-4 if distance.dtype == torch.float32 else 1e-8
+        assert torch.allclose(distance.cpu().double(), reference, rtol=tol, atol=tol)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_log_space_survives_small_epsilon(device):
+    """At this epsilon exp(-C/eps) underflows float32, so only log_space=True is usable.
+
+    On CPU this exercises the torch logsumexp path, on CUDA the fused log kernel.
+    """
     _require_pot()
     solver = SinkhornSolver(epsilon=0.1, max_iter=5000, log_space=True)
-    cost_matrix = build_cost_matrix((8, 8), device="cuda")
+    cost_matrix = build_cost_matrix((8, 8), device=device)
 
-    source = _make_positive((1, 64), device="cuda", seed=31)
-    target = _make_positive((1, 64), device="cuda", seed=37)
+    source = _make_positive((1, 64), device=device, seed=31)
+    target = _make_positive((1, 64), device=device, seed=37)
 
     f, g = solver.potentials(source, target, cost_matrix)
     assert torch.isfinite(f).all()
@@ -193,11 +236,14 @@ def test_log_space_survives_small_epsilon():
     assert distance[0].item() == pytest.approx(pot_distance, rel=2e-2)
 
 
-@pytest.mark.parametrize("log_space", [False, True])
-@pytest.mark.parametrize("n,d", [(1, 4), (3, 8)])
-def test_backward_matches_numeric_directional_derivative(n, d, log_space):
-    """Autograd through forward() must match a numeric derivative of the entropic cost."""
-    device = _device()
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_backward_matches_numeric_directional_derivative(log_space, device):
+    """Autograd through forward() must match a numeric derivative of the entropic cost.
+
+    Runs in float64, i.e. through the torch paths on both CPU and CUDA.
+    """
+    n, d = 3, 8
     solver = SinkhornSolver(epsilon=1.5, max_iter=3000, log_space=log_space)
     cost_matrix = build_cost_matrix((d,), device=device).double()
 
@@ -219,3 +265,38 @@ def entropic_cost(x):
         entropic_cost(raw + eps_fd * direction) - entropic_cost(raw - eps_fd * direction)
     ) / (2 * eps_fd)
     assert torch.allclose(analytic, numeric, rtol=1e-5, atol=1e-7)
+
+
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_backward_through_fused_kernels_matches_cpu(log_space):
+    """float32 CUDA backward runs through the fused kernels; grads must match CPU."""
+    _require_cuda()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
+
+    source = _make_positive((3, 6), device="cpu", seed=1213)
+    target = _make_positive((3, 6), device="cpu", seed=1217)
+
+    cpu_source = source.clone().requires_grad_(True)
+    solver(cpu_source, target).sum().backward()
+    cuda_source = source.cuda().requires_grad_(True)
+    solver(cuda_source, target.cuda()).sum().backward()
+
+    assert torch.isfinite(cuda_source.grad).all()
+    assert torch.allclose(cuda_source.grad.cpu(), cpu_source.grad, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_threshold_early_stopping_reaches_same_solution(log_space, device):
+    """threshold > 0 stops the torch paths early without changing the solution."""
+    source = _make_positive((2, 8), device=device, seed=1301)
+    target = _make_positive((2, 8), device=device, seed=1303)
+    source, target = source.double(), target.double()  # force the torch paths
+
+    full = SinkhornSolver(epsilon=1.5, max_iter=5000, log_space=log_space)
+    early = SinkhornSolver(epsilon=1.5, max_iter=5000, threshold=1e-12, log_space=log_space)
+
+    assert torch.allclose(early(source, target), full(source, target), rtol=1e-8, atol=1e-10)
+    assert torch.allclose(
+        early.plan(source, target), full.plan(source, target), rtol=1e-6, atol=1e-10
+    )

From ed6d1c34645c4920435743abc54ad1d76078f5d9 Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Wed, 8 Jul 2026 23:55:15 +0800
Subject: [PATCH 19/20] Optimize Sinkhorn CUDA path: batch tiling, online
 logsumexp, CUDA graphs

Kernel changes (both memory-bandwidth-bound, so all wins target matrix
traffic and launch latency):

- Batch tiling: one block now streams a matrix row once and applies it to
  up to 8 batch items held in registers, cutting d^2 matrix reads by up to
  8x for batched inputs. A TILE=1 template instantiation keeps the n=1
  fast path free of predication overhead.
- Online logsumexp: the log-domain update now computes the LSE in a single
  pass over the row (running max with rescaled running sum, NaN-free -inf
  merging), halving matrix traffic vs the two-pass max-then-sum version.
- Warp-shuffle reductions replace the shared-memory tree (one barrier
  instead of eight per reduction).

Python-side:

- Long fused runs (max_iter >= 100) capture a 25-iteration chunk into a
  torch.cuda.CUDAGraph and replay it, hiding the 2-launches-per-iteration
  overhead that dominates small problems; falls back to plain launches
  when capture is unavailable (extra fixed-point iterations are harmless).
- Skip the (d, d) transpose copy when the cost/kernel matrix is symmetric,
  which build_cost_matrix outputs always are.

Tests: fused-vs-CPU now covers single-item, partial-tile, and multi-tile
batches (n = 1, 3, 20) plus a short-run case below the CUDA-graph
threshold to exercise the plain-launch fused path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 test/test_optimal_transport.py  |  21 ++-
 torchmorph/csrc/ot_kernel.cu    | 260 ++++++++++++++++++++++----------
 torchmorph/optimal_transport.py |  78 +++++++---
 3 files changed, 261 insertions(+), 98 deletions(-)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 6ceb62f..0289b49 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -153,13 +153,14 @@ def test_plan_matches_pot_sinkhorn_grid_cost(log_space, device):
 
 
 @pytest.mark.parametrize("log_space", LOG_SPACE)
-def test_fused_cuda_kernels_match_cpu(log_space):
+@pytest.mark.parametrize("n", [1, 3, 20])  # single-item, partial-tile, and multi-tile batches
+def test_fused_cuda_kernels_match_cpu(log_space, n):
     """CUDA float32 inputs dispatch to the fused kernels; results must match CPU."""
     _require_cuda()
     solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
 
-    source = _make_positive((3, 6), device="cpu", seed=503)
-    target = _make_positive((3, 6), device="cpu", seed=607)
+    source = _make_positive((n, 6), device="cpu", seed=503)
+    target = _make_positive((n, 6), device="cpu", seed=607)
 
     cpu_distance = solver(source, target)
     cpu_plan = solver.plan(source, target)
@@ -174,6 +175,20 @@ def test_fused_cuda_kernels_match_cpu(log_space):
     assert torch.allclose(cuda_g.cpu(), cpu_g, rtol=1e-4, atol=1e-4)
 
 
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_fused_short_run_without_cuda_graph_matches_cpu(log_space):
+    """max_iter below the CUDA-graph threshold takes the plain-launch fused path."""
+    _require_cuda()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=60, log_space=log_space)
+
+    source = _make_positive((2, 6), device="cpu", seed=1409)
+    target = _make_positive((2, 6), device="cpu", seed=1423)
+
+    cpu_distance = solver(source, target)
+    cuda_distance = solver(source.cuda(), target.cuda())
+    assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5)
+
+
 @pytest.mark.parametrize("log_space", LOG_SPACE)
 def test_dtype_device_combinations_agree(log_space):
     """float32/float64 on CPU/CUDA must all reach the same solution.
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
index 5481268..193d1ac 100644
--- a/torchmorph/csrc/ot_kernel.cu
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -5,85 +5,167 @@
 
 namespace {
 
-constexpr int kThreads = 256;  // power of two, assumed by the reductions below
+constexpr int kThreads = 256;  // multiple of the warp size, assumed below
+constexpr int kBatchTile = 8;  // batch items processed per block
+constexpr int kWarps = kThreads / 32;
 
-// One Sinkhorn scaling update. Vectors are batched (n, d); the (d, d) matrix
-// is shared across the batch. One block per (row, batch) pair computes:
+__device__ inline float warp_reduce_sum(float v){
+    for (int offset = 16; offset > 0; offset >>= 1){
+        v += __shfl_down_sync(0xffffffff, v, offset);
+    }
+    return v;
+}
+
+// Block-level sum reduction; the result is valid in thread 0 only.
+__device__ inline float block_reduce_sum(float v){
+    __shared__ float partial[kWarps];
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    v = warp_reduce_sum(v);
+    if (lane == 0) partial[warp] = v;
+    __syncthreads();
+    if (warp == 0){
+        v = lane < kWarps ? partial[lane] : 0.0f;
+        v = warp_reduce_sum(v);
+    }
+    __syncthreads();  // let the next call reuse `partial`
+    return v;
+}
+
+// Merge two online-logsumexp states; (m, s) represents m + log(s).
+// s is 0 exactly when m is -inf (the empty state), which the guards below
+// keep NaN-free.
+__device__ inline void lse_merge(float& m, float& s, float m_other, float s_other){
+    const float hi = fmaxf(m, m_other);
+    if (hi == -INFINITY){
+        m = -INFINITY;
+        s = 0.0f;
+        return;
+    }
+    s = (m == -INFINITY ? 0.0f : s * expf(m - hi)) +
+        (m_other == -INFINITY ? 0.0f : s_other * expf(m_other - hi));
+    m = hi;
+}
+
+__device__ inline void warp_reduce_lse(float& m, float& s){
+    for (int offset = 16; offset > 0; offset >>= 1){
+        const float m_other = __shfl_down_sync(0xffffffff, m, offset);
+        const float s_other = __shfl_down_sync(0xffffffff, s, offset);
+        lse_merge(m, s, m_other, s_other);
+    }
+}
+
+// Block-level logsumexp reduction; the result is valid in thread 0 only.
+__device__ inline void block_reduce_lse(float& m, float& s){
+    __shared__ float partial_m[kWarps];
+    __shared__ float partial_s[kWarps];
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    warp_reduce_lse(m, s);
+    if (lane == 0){
+        partial_m[warp] = m;
+        partial_s[warp] = s;
+    }
+    __syncthreads();
+    if (warp == 0){
+        m = lane < kWarps ? partial_m[lane] : -INFINITY;
+        s = lane < kWarps ? partial_s[lane] : 0.0f;
+        warp_reduce_lse(m, s);
+    }
+    __syncthreads();  // let the next call reuse the partial arrays
+}
+
+// One Sinkhorn scaling update. One block per (row, batch tile): the matrix
+// row is streamed once and applied to TILE batch items, so the d^2 matrix is
+// read ceil(n / TILE) times per update instead of n times.
 //   out[b, i] = num[b, i] / (sum_j mat[i, j] * scale[b, j] + 1e-12)
 // The u-update passes (a, K, v); the v-update passes (b, K^T, u).
+template <int TILE>
 __global__ void scaling_update(
     const float* __restrict__ num,
     const float* __restrict__ mat,
     const float* __restrict__ scale,
     float* __restrict__ out,
-    int d
+    int d,
+    int n_batch
 ){
-    extern __shared__ float sdata[];
     const int row = blockIdx.x;
-    const int tid = threadIdx.x;
-    const size_t offset = (size_t)blockIdx.y * d;
+    const int batch0 = blockIdx.y * TILE;
+    const int tile = min(TILE, n_batch - batch0);
+    const float* mat_row = mat + (size_t)row * d;
 
-    float local_sum = 0.0f;
-    for (int j = tid; j < d; j += blockDim.x){
-        local_sum += mat[(size_t)row * d + j] * scale[offset + j];
-    }
-    sdata[tid] = local_sum;
-    __syncthreads();
-    for (int stride = blockDim.x / 2; stride > 0; stride /= 2){
-        if (tid < stride) sdata[tid] += sdata[tid + stride];
-        __syncthreads();
+    float acc[TILE];
+#pragma unroll
+    for (int t = 0; t < TILE; t++) acc[t] = 0.0f;
+
+    for (int j = threadIdx.x; j < d; j += blockDim.x){
+        const float k = mat_row[j];
+#pragma unroll
+        for (int t = 0; t < TILE; t++){
+            if (t < tile) acc[t] += k * scale[(size_t)(batch0 + t) * d + j];
+        }
     }
 
-    if (tid == 0) out[offset + row] = num[offset + row] / (sdata[0] + 1e-12f);
+    for (int t = 0; t < tile; t++){
+        const float total = block_reduce_sum(acc[t]);
+        if (threadIdx.x == 0){
+            const size_t idx = (size_t)(batch0 + t) * d + row;
+            out[idx] = num[idx] / (total + 1e-12f);
+        }
+    }
 }
 
-// One log-domain Sinkhorn update, one block per (row, batch) pair:
+// One log-domain Sinkhorn update with the same batch tiling:
 //   log_out[b, i] = log_num[b, i] - logsumexp_j(-cost[i, j] / eps + log_scale[b, j])
-// evaluated with a two-pass block reduction (max, then sum of shifted exps).
+// The logsumexp is evaluated online in a single pass over the row (running
+// max with a rescaled running sum), halving the matrix traffic of the
+// classic two-pass max-then-sum evaluation.
 // The u-update passes (log_a, M, log_v); the v-update passes (log_b, M^T, log_u).
+template <int TILE>
 __global__ void log_scaling_update(
     const float* __restrict__ log_num,
     const float* __restrict__ cost,
     const float* __restrict__ log_scale,
     float* __restrict__ log_out,
     int d,
+    int n_batch,
     float inv_eps
 ){
-    extern __shared__ float sdata[];
     const int row = blockIdx.x;
-    const int tid = threadIdx.x;
-    const size_t offset = (size_t)blockIdx.y * d;
+    const int batch0 = blockIdx.y * TILE;
+    const int tile = min(TILE, n_batch - batch0);
+    const float* cost_row = cost + (size_t)row * d;
 
-    float local_max = -INFINITY;
-    for (int j = tid; j < d; j += blockDim.x){
-        local_max = fmaxf(local_max, log_scale[offset + j] - cost[(size_t)row * d + j] * inv_eps);
-    }
-    sdata[tid] = local_max;
-    __syncthreads();
-    for (int stride = blockDim.x / 2; stride > 0; stride /= 2){
-        if (tid < stride) sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]);
-        __syncthreads();
+    float m[TILE], s[TILE];
+#pragma unroll
+    for (int t = 0; t < TILE; t++){
+        m[t] = -INFINITY;
+        s[t] = 0.0f;
     }
-    const float row_max = sdata[0];
-    __syncthreads();
 
-    float local_sum = 0.0f;
-    for (int j = tid; j < d; j += blockDim.x){
-        local_sum += expf(log_scale[offset + j] - cost[(size_t)row * d + j] * inv_eps - row_max);
-    }
-    sdata[tid] = local_sum;
-    __syncthreads();
-    for (int stride = blockDim.x / 2; stride > 0; stride /= 2){
-        if (tid < stride) sdata[tid] += sdata[tid + stride];
-        __syncthreads();
+    for (int j = threadIdx.x; j < d; j += blockDim.x){
+        const float c = -cost_row[j] * inv_eps;
+#pragma unroll
+        for (int t = 0; t < TILE; t++){
+            if (t >= tile) continue;
+            const float z = c + log_scale[(size_t)(batch0 + t) * d + j];
+            if (z > m[t]){
+                s[t] = s[t] * expf(m[t] - z) + 1.0f;  // expf(-inf) == 0 covers the first element
+                m[t] = z;
+            } else if (m[t] != -INFINITY){  // skip z == m == -inf, whose contribution is 0
+                s[t] += expf(z - m[t]);
+            }
+        }
     }
 
-    // row_max == -inf means the opposite marginal is all zero; pin the
-    // potential instead of propagating -inf - (-inf) = NaN.
-    if (tid == 0){
-        log_out[offset + row] = row_max == -INFINITY
-            ? -INFINITY
-            : log_num[offset + row] - row_max - logf(sdata[0]);
+    for (int t = 0; t < tile; t++){
+        block_reduce_lse(m[t], s[t]);
+        if (threadIdx.x == 0){
+            const size_t idx = (size_t)(batch0 + t) * d + row;
+            // m == -inf means the opposite marginal is all zero; pin the
+            // potential instead of producing NaN.
+            log_out[idx] = m[t] == -INFINITY ? -INFINITY : log_num[idx] - m[t] - logf(s[t]);
+        }
     }
 }
 
@@ -94,13 +176,11 @@ void check_input(const torch::Tensor& t, int64_t numel, const char* name){
     TORCH_CHECK(t.numel() == numel, name, " must have ", numel, " elements, got ", t.numel());
 }
 
-// Validate batched (n, d) vectors against a shared (d, d) matrix and return
-// the launch grid: one block per (row, batch) pair.
-dim3 make_grid(const torch::Tensor& u, const torch::Tensor& mat){
+// Validate batched (n, d) vectors against a shared (d, d) matrix.
+void check_shapes(const torch::Tensor& u, const torch::Tensor& mat){
     TORCH_CHECK(mat.dim() == 2 && mat.size(0) == mat.size(1), "cost/kernel matrix must be square");
     TORCH_CHECK(u.dim() == 2 && u.size(1) == mat.size(0),
                 "scaling vectors must be (n, d) with d matching the matrix");
-    return dim3(static_cast<unsigned>(u.size(1)), static_cast<unsigned>(u.size(0)));
 }
 
 }  // namespace
@@ -114,26 +194,38 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
     torch::Tensor v,
     int64_t n_iter
 ){
-    const dim3 grid = make_grid(u, K);
-    const int64_t d = K.size(0);
-    check_input(a, u.numel(), "a");
-    check_input(b, u.numel(), "b");
+    check_shapes(u, K);
+    const int64_t n = u.size(0);
+    const int64_t d = u.size(1);
+    check_input(a, n * d, "a");
+    check_input(b, n * d, "b");
     check_input(K, d * d, "K");
     check_input(K_T, d * d, "K_T");
-    check_input(u, u.numel(), "u");
-    check_input(v, u.numel(), "v");
+    check_input(u, n * d, "u");
+    check_input(v, n * d, "v");
 
     const at::cuda::CUDAGuard guard(u.device());
     const auto stream = at::cuda::getCurrentCUDAStream();
-    const int shmem = kThreads * sizeof(float);
+    const bool single = n == 1;
+    const dim3 grid(static_cast<unsigned>(d),
+                    single ? 1u : static_cast<unsigned>((n + kBatchTile - 1) / kBatchTile));
 
     for (int64_t i = 0; i < n_iter; i++){
-        scaling_update<<<grid, kThreads, shmem, stream>>>(
-            a.data_ptr<float>(), K.data_ptr<float>(), v.data_ptr<float>(),
-            u.data_ptr<float>(), static_cast<int>(d));
-        scaling_update<<<grid, kThreads, shmem, stream>>>(
-            b.data_ptr<float>(), K_T.data_ptr<float>(), u.data_ptr<float>(),
-            v.data_ptr<float>(), static_cast<int>(d));
+        if (single){
+            scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                a.data_ptr<float>(), K.data_ptr<float>(), v.data_ptr<float>(),
+                u.data_ptr<float>(), (int)d, (int)n);
+            scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                b.data_ptr<float>(), K_T.data_ptr<float>(), u.data_ptr<float>(),
+                v.data_ptr<float>(), (int)d, (int)n);
+        } else {
+            scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                a.data_ptr<float>(), K.data_ptr<float>(), v.data_ptr<float>(),
+                u.data_ptr<float>(), (int)d, (int)n);
+            scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                b.data_ptr<float>(), K_T.data_ptr<float>(), u.data_ptr<float>(),
+                v.data_ptr<float>(), (int)d, (int)n);
+        }
     }
     C10_CUDA_KERNEL_LAUNCH_CHECK();
 
@@ -151,27 +243,39 @@ std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
     double epsilon
 ){
     TORCH_CHECK(epsilon > 0, "epsilon must be positive, got ", epsilon);
-    const dim3 grid = make_grid(log_u, M);
-    const int64_t d = M.size(0);
-    check_input(log_a, log_u.numel(), "log_a");
-    check_input(log_b, log_u.numel(), "log_b");
+    check_shapes(log_u, M);
+    const int64_t n = log_u.size(0);
+    const int64_t d = log_u.size(1);
+    check_input(log_a, n * d, "log_a");
+    check_input(log_b, n * d, "log_b");
     check_input(M, d * d, "M");
     check_input(M_T, d * d, "M_T");
-    check_input(log_u, log_u.numel(), "log_u");
-    check_input(log_v, log_u.numel(), "log_v");
+    check_input(log_u, n * d, "log_u");
+    check_input(log_v, n * d, "log_v");
 
     const at::cuda::CUDAGuard guard(log_u.device());
     const auto stream = at::cuda::getCurrentCUDAStream();
-    const int shmem = kThreads * sizeof(float);
     const float inv_eps = static_cast<float>(1.0 / epsilon);
+    const bool single = n == 1;
+    const dim3 grid(static_cast<unsigned>(d),
+                    single ? 1u : static_cast<unsigned>((n + kBatchTile - 1) / kBatchTile));
 
     for (int64_t i = 0; i < n_iter; i++){
-        log_scaling_update<<<grid, kThreads, shmem, stream>>>(
-            log_a.data_ptr<float>(), M.data_ptr<float>(), log_v.data_ptr<float>(),
-            log_u.data_ptr<float>(), static_cast<int>(d), inv_eps);
-        log_scaling_update<<<grid, kThreads, shmem, stream>>>(
-            log_b.data_ptr<float>(), M_T.data_ptr<float>(), log_u.data_ptr<float>(),
-            log_v.data_ptr<float>(), static_cast<int>(d), inv_eps);
+        if (single){
+            log_scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                log_a.data_ptr<float>(), M.data_ptr<float>(), log_v.data_ptr<float>(),
+                log_u.data_ptr<float>(), (int)d, (int)n, inv_eps);
+            log_scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                log_b.data_ptr<float>(), M_T.data_ptr<float>(), log_u.data_ptr<float>(),
+                log_v.data_ptr<float>(), (int)d, (int)n, inv_eps);
+        } else {
+            log_scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                log_a.data_ptr<float>(), M.data_ptr<float>(), log_v.data_ptr<float>(),
+                log_u.data_ptr<float>(), (int)d, (int)n, inv_eps);
+            log_scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                log_b.data_ptr<float>(), M_T.data_ptr<float>(), log_u.data_ptr<float>(),
+                log_v.data_ptr<float>(), (int)d, (int)n, inv_eps);
+        }
     }
     C10_CUDA_KERNEL_LAUNCH_CHECK();
 
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
index 3415283..3652845 100644
--- a/torchmorph/optimal_transport.py
+++ b/torchmorph/optimal_transport.py
@@ -22,21 +22,63 @@ def _use_fused_kernels(t: Tensor) -> bool:
     return t.is_cuda and t.dtype == torch.float32
 
 
+def _transpose_or_self(matrix: Tensor) -> Tensor:
+    """Return the contiguous transpose, reusing the input when it is symmetric.
+
+    Cost matrices from build_cost_matrix are always symmetric, so this
+    usually saves a (d, d) copy.
+    """
+    return matrix if torch.equal(matrix, matrix.mT) else matrix.mT.contiguous()
+
+
+_GRAPH_MIN_ITER = 100  # below this, plain launches beat the graph-capture setup cost
+_GRAPH_CHUNK = 25
+
+
+def _run_fused(launch, max_iter, device):
+    """Run `launch(k)`, which enqueues k in-place Sinkhorn iterations, max_iter times in total.
+
+    Long runs are launch-latency-bound (two tiny kernels per iteration), so a
+    chunk of iterations is captured into a CUDA graph once and replayed. The
+    iterations are fixed-point steps on ping-pong buffers, so a fallback that
+    runs extra iterations is always safe.
+    """
+    if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing():
+        launch(max_iter)
+        return
+    try:
+        with torch.cuda.device(device):
+            # Warm up on a side stream (required before capture), doing the
+            # first chunk of real iterations in the process.
+            side_stream = torch.cuda.Stream()
+            side_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(side_stream):
+                launch(_GRAPH_CHUNK)
+            torch.cuda.current_stream().wait_stream(side_stream)
+
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                launch(_GRAPH_CHUNK)
+            for _ in range(max_iter // _GRAPH_CHUNK - 1):
+                graph.replay()
+            remainder = max_iter % _GRAPH_CHUNK
+            if remainder:
+                launch(remainder)
+    except RuntimeError:
+        launch(max_iter)  # graphs unavailable; extra iterations are harmless
+
+
 def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold):
     """Plain scaling-form Sinkhorn on (n, d) marginals; returns (log_u, log_v)."""
     K = torch.exp(-cost_matrix / epsilon)
     if _use_fused_kernels(a):
         from torchmorph import _C
 
-        u, v = _C.sinkhorn_fastiter(
-            a.contiguous(),
-            b.contiguous(),
-            K,
-            K.mT.contiguous(),
-            torch.ones_like(a),
-            torch.ones_like(b),
-            max_iter,
-        )
+        a, b = a.contiguous(), b.contiguous()
+        K_T = _transpose_or_self(K)
+        u = torch.ones_like(a)
+        v = torch.ones_like(b)
+        _run_fused(lambda k: _C.sinkhorn_fastiter(a, b, K, K_T, u, v, k), max_iter, a.device)
         return u.log(), v.log()
 
     u = torch.ones_like(a)
@@ -57,16 +99,18 @@ def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold):
     if _use_fused_kernels(a):
         from torchmorph import _C
 
-        return _C.sinkhorn_logiter(
-            log_a.contiguous(),
-            log_b.contiguous(),
-            cost_matrix,
-            cost_matrix.mT.contiguous(),
-            torch.zeros_like(a),
-            torch.zeros_like(b),
+        log_a, log_b = log_a.contiguous(), log_b.contiguous()
+        cost_T = _transpose_or_self(cost_matrix)
+        log_u = torch.zeros_like(a)
+        log_v = torch.zeros_like(b)
+        _run_fused(
+            lambda k: _C.sinkhorn_logiter(
+                log_a, log_b, cost_matrix, cost_T, log_u, log_v, k, epsilon
+            ),
             max_iter,
-            epsilon,
+            a.device,
         )
+        return log_u, log_v
 
     log_K = -cost_matrix / epsilon
     log_u = torch.zeros_like(a)

From 3f374ad1584fd278708a6b34a14663bad0ce184c Mon Sep 17 00:00:00 2001
From: Kai ZHAO <694691@qq.com>
Date: Thu, 9 Jul 2026 00:13:36 +0800
Subject: [PATCH 20/20] Test gradients w.r.t. both source and target

The backward pass already returned gradients for both marginals (the
centered dual potentials f and g), but the tests only verified the source
side. The numeric-derivative test now perturbs source and target
independently, and the fused-kernel backward test compares both gradients
against CPU.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 test/test_optimal_transport.py | 40 +++++++++++++++++++++-------------
 1 file changed, 25 insertions(+), 15 deletions(-)

diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
index 0289b49..785e069 100644
--- a/test/test_optimal_transport.py
+++ b/test/test_optimal_transport.py
@@ -254,7 +254,8 @@ def test_log_space_survives_small_epsilon(device):
 @pytest.mark.parametrize("device", DEVICES)
 @pytest.mark.parametrize("log_space", LOG_SPACE)
 def test_backward_matches_numeric_directional_derivative(log_space, device):
-    """Autograd through forward() must match a numeric derivative of the entropic cost.
+    """Autograd gradients w.r.t. BOTH source and target must match numeric
+    derivatives of the entropic cost.
 
     Runs in float64, i.e. through the torch paths on both CPU and CUDA.
     """
@@ -263,23 +264,26 @@ def test_backward_matches_numeric_directional_derivative(log_space, device):
     cost_matrix = build_cost_matrix((d,), device=device).double()
 
     source = _make_positive((n, d), device=device, seed=1009).double().requires_grad_(True)
-    target = _make_positive((n, d), device=device, seed=2003).double()
+    target = _make_positive((n, d), device=device, seed=2003).double().requires_grad_(True)
     direction = _make_positive((n, d), device=device, seed=3001).double() - 0.6
 
     solver(source, target, cost_matrix).sum().backward()
-    analytic = (source.grad * direction).sum()
+    assert source.grad is not None and target.grad is not None
 
-    def entropic_cost(x):
-        plan = solver.plan(x, target, cost_matrix)
+    def entropic_cost(src, tgt):
+        plan = solver.plan(src, tgt, cost_matrix)
         entropy = (plan * plan.clamp(min=1e-300).log()).sum()
         return (plan * cost_matrix).sum() + solver.epsilon * entropy
 
     eps_fd = 1e-4
-    raw = source.detach()
-    numeric = (
-        entropic_cost(raw + eps_fd * direction) - entropic_cost(raw - eps_fd * direction)
-    ) / (2 * eps_fd)
-    assert torch.allclose(analytic, numeric, rtol=1e-5, atol=1e-7)
+    src, tgt = source.detach(), target.detach()
+    for grad, plus, minus in [
+        (source.grad, (src + eps_fd * direction, tgt), (src - eps_fd * direction, tgt)),
+        (target.grad, (src, tgt + eps_fd * direction), (src, tgt - eps_fd * direction)),
+    ]:
+        analytic = (grad * direction).sum()
+        numeric = (entropic_cost(*plus) - entropic_cost(*minus)) / (2 * eps_fd)
+        assert torch.allclose(analytic, numeric, rtol=1e-5, atol=1e-7)
 
 
 @pytest.mark.parametrize("log_space", LOG_SPACE)
@@ -292,12 +296,18 @@ def test_backward_through_fused_kernels_matches_cpu(log_space):
     target = _make_positive((3, 6), device="cpu", seed=1217)
 
     cpu_source = source.clone().requires_grad_(True)
-    solver(cpu_source, target).sum().backward()
+    cpu_target = target.clone().requires_grad_(True)
+    solver(cpu_source, cpu_target).sum().backward()
     cuda_source = source.cuda().requires_grad_(True)
-    solver(cuda_source, target.cuda()).sum().backward()
-
-    assert torch.isfinite(cuda_source.grad).all()
-    assert torch.allclose(cuda_source.grad.cpu(), cpu_source.grad, rtol=1e-4, atol=1e-5)
+    cuda_target = target.cuda().requires_grad_(True)
+    solver(cuda_source, cuda_target).sum().backward()
+
+    for cuda_grad, cpu_grad in [
+        (cuda_source.grad, cpu_source.grad),
+        (cuda_target.grad, cpu_target.grad),
+    ]:
+        assert torch.isfinite(cuda_grad).all()
+        assert torch.allclose(cuda_grad.cpu(), cpu_grad, rtol=1e-4, atol=1e-5)
 
 
 @pytest.mark.parametrize("device", DEVICES)