TorchGW – Fast Sampled Gromov-Wasserstein Optimal Transport

TorchGW is a scalable solver for Gromov-Wasserstein optimal transport, implemented in pure PyTorch with GPU-accelerated Triton fused Sinkhorn kernels.

It aligns two point clouds by matching their internal distance structures — even when the point clouds live in different dimensions — making it ideal for manifold alignment, single-cell multi-omics integration, and cross-domain graph matching.

Key features:

• Up to 175x faster than POT on typical workloads (spiral 4000×5000: 1s vs 183s)

• Triton fused Sinkhorn — single-pass online logsumexp, no intermediate N×K matrices

• Mixed precision — float32 Sinkhorn + float64 output, zero quality loss

• Smart early stopping — cost plateau detection, not just transport plan norm

• Differentiable — exact gradients via implicit differentiation at the Sinkhorn fixed point

• No POT dependency at runtime — pure PyTorch + scipy + scikit-learn

pip install torchgw

git clone https://github.com/chansigit/torchgw.git
cd torchgw && pip install -e ".[dev]"

from torchgw import sampled_gw

T = sampled_gw(X_source, X_target, distance_mode="landmark", mixed_precision=True)
# T[i,j] = optimal coupling weight between source point i and target point j

# Clone and install from source
git clone https://github.com/chansigit/torchgw.git
cd torchgw
pip install -e .

@software{torchgw,
  author = {Sijie Chen},
  title = {TorchGW: Fast Sampled Gromov-Wasserstein Optimal Transport},
  url = {https://github.com/chansigit/torchgw},
  version = {0.4.1},
  year = {2026},
}